import json
import os
import datetime
import shutil
import subprocess
import signal
import sys
import concurrent.futures
from pathlib import Path
from typing import Dict, Any, List


dirs = [
    "history",
    "perf_jsons",
    "IR",
    "IR/base",
    "bin",
    "obj",
    "remarks",
    "features",
    "jsons",
]


def load_bench(bench_def_json: str) -> List[Dict]:
    benchs = []
    with open(bench_def_json, "r") as f:
        benchs = json.load(f)

    return benchs


def run_and_archive(bench_path, flags, run_flag, log_dir, target_machine):
    # create log file for every single workers
    bench_name = bench_path.split('/')[-1]
    process_log_file = f"./AutoPass_{bench_name}.log"

    sub_command = (
        f"yes | python -u AutoPassWorker.py "
        f"-p '{bench_path}' "
        f"-f '{flags}' "
        f"-r '{run_flag}' "
        f"-s '{log_dir}' "
        f"{'--target ' + target_machine if target_machine else ''}"
    )

    print(f"----- : {bench_name} (: {process_log_file}) -----")

    try:
        with open(process_log_file, 'w', encoding='utf-8') as log_f:
            result = subprocess.run(
                sub_command,
                shell=True,
                executable='/bin/bash',
                check=False,
                stdout=log_f,
                stderr=log_f
            )
        print(f"----- {bench_name} : {result.returncode} -----")
    except Exception as e:
        print(f"----- {bench_name} : {e}")
        return None
    try:
        archive_path = f"{log_dir}/{bench_name}.log"
        shutil.move(process_log_file, archive_path)
        print(f"----- {bench_name} {archive_path} -----")
    except Exception as e:
        print(f"----- {bench_name} : {e}")

    sys.stdout.flush()
    return result.returncode


def process_benchmark(bench, log_dir, target_machine):
    bench_path = bench["path"]
    ldflag = bench["ldflags"]
    runflag = bench["runflags"]
    bench_name = bench_path.split("/")[-1]
    print(f"--- : {bench_name} ---")

    if not Path(bench_path).is_dir():
        print(f"error: Benchmark : {bench_path}")
        return None

    for d in dirs:
        try:
            target = Path(f"{bench_path}/{d}")
            if not target.is_dir():
                target.mkdir(exist_ok=True)
        except OSError as e:
            print(f"{bench_name} {d} mkdir : {e}")
            return None

    return_code = run_and_archive(
        bench_path, ldflag, runflag, log_dir, target_machine=target_machine)

    if return_code == 0:
        print(f"--- {bench_name}  ---")
        return bench
    else:
        print(f"--- {bench_name} : {return_code} ---")
        return None


if __name__ == "__main__":
    bench_def_json = "./loadbenchs/cbench/benchmarks.json"

    # --- Target machine ---
    # "" for local test
    # to run on a remote machine, set to machines in remote configs ("machine1" for example)
    target_machine = ""
    # target_machine = "machine1"

    benchs = load_bench(bench_def_json)
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    log_dir = f"log/log_{timestamp}"
    os.makedirs(log_dir, exist_ok=True)

    tested_json = f"{log_dir}/benchmarks_tested.json"

    #  CPU  API
    # Adjusting according cpu cores and API limitations
    MAX_WORKERS = os.cpu_count() // 2 if os.cpu_count() else 4
    if MAX_WORKERS == 0:
        MAX_WORKERS = 2

    # MAX_WORKERS = 1
    MAX_WORKERS = 10

    print(f"---   {MAX_WORKERS}  ---")
    print(f"---  {MAX_WORKERS} workers running AutoPass. ---")

    tested_benchs = []

    executor = None

    try:
        executor = concurrent.futures.ProcessPoolExecutor(
            max_workers=MAX_WORKERS)

        #
        futures = {executor.submit(
            process_benchmark, bench, log_dir, target_machine): bench for bench in benchs}

        for future in concurrent.futures.as_completed(futures):
            bench_name = futures[future]["path"].split("/")[-1]
            try:
                result_bench = future.result()
                if result_bench:
                    tested_benchs.append(result_bench)

                    with open(tested_json, 'w', encoding="utf-8") as f:
                        json.dump(tested_benchs, f, indent=2,
                                  ensure_ascii=False)
                else:
                    print(f"--- {bench_name}  tested_json ---")
            except Exception as e:
                if not isinstance(e, concurrent.futures.CancelledError):
                    print(f"--- {bench_name} : {e} ---")

    #  Ctrl+C
    except KeyboardInterrupt:
        print("\n---  (Ctrl+C)... ---")

        if executor:
            executor.shutdown(wait=False, cancel_futures=True)

        print("---  ---")
        sys.exit(1)

    finally:
        if executor:
            executor.shutdown(wait=True)

    print(f"--- Benchmarks tested. Tested list saved to {tested_json} ---")
