# autopep8: off
# fmt: off
import datetime
import os
import shutil
import subprocess
import threading
import time
from typing import Any

from utils.signal_utils import signal_manager
from vwa_utils.captioner_utils import is_captioner_running

# Usage: python -u -m scripts.runs.p_run_many.py 
# nohup python -u -m offline_experiments.offline_verify > verify.log 2>&1 & disown  # Run detached in the background
#===============================================
# Configs
#===============================================
# Params for this script

# FIXME: MEMO: new_som is commented out

# Dict of agent config and the tasks to run on
# agent_config = "2p-verifier_t-30_yes-expert_verifier_dcv5_r4_inf_q_fixK.yaml"
# agent_config = "1p-verifier_t-30_yes-expert_verifier_dcv5_r4_inf_q_fixK.yaml"
# agent_config = "noverifier_t-30.yaml"
# agent_config = "debug_reflexion.yaml"
# dir_to_save = "experiments/gemini-2.5-flash-001/debug_reflexion-4"

# agent_config = "agent-thinking_verifier-thinking.yaml"
task_list1p = "evaluation_harness/task_subsets/"
task_list2p = "evaluation_harness/task_subsets/"

agent_config = "reflexion_1p-verifier-oracle-execution.yaml"
agent_config2p = "reflexion_2p-verifier-oracle-execution.yaml"
dir_to_save = f"experiments/gemini-2.5-flash-001/reflexion-1p-verifier-oracle-execution"
dir_to_save2p = f"experiments/gemini-2.5-flash-001/reflexion-2p-verifier-oracle-execution"
run_configs = [
    {"agent_config":agent_config2p, "task_list": f"{task_list2p}/shopping_not_eval_rand_reduced.txt", "results_dir": f"{dir_to_save2p}/shopping", "block":True, "num_processes":6},
    {"agent_config": "reflexion_oracle.yaml", "task_list": f"{task_list1p}/shopping_not_eval_rand_reduced.txt", "results_dir": f"{dir_to_save}/shopping", "block":True, "num_processes":6},
    {"agent_config": agent_config, "task_list": f"{task_list1p}/shopping_not_eval_rand_reduced.txt", "results_dir": f"{dir_to_save}/shopping", "block":True, "num_processes":6},
    
    # {"agent_config":agent_config2p, "task_list": f"{task_list2p}/classifieds_not_eval.txt", "results_dir": f"{dir_to_save2p}/classifieds", "block":True, "num_processes":6},
    # {"agent_config": agent_config, "task_list": f"{task_list1p}/classifieds_not_eval.txt", "results_dir": f"{dir_to_save}/classifieds", "block":True, "num_processes":6},

    # {"agent_config":agent_config2p, "task_list": f"{task_list2p}/reddit_not_eval.txt", "results_dir": f"{dir_to_save2p}/reddit", "block":True, "num_processes":6},
    # {"agent_config": agent_config, "task_list": f"{task_list1p}/reddit_not_eval.txt", "results_dir": f"{dir_to_save}/reddit", "block":True, "num_processes":6},


    # {"agent_config": "noverifier_t-30.yaml", "task_list": "evaluation_harness/tasks_vwa/reddit.txt", "results_dir": "experiments/gemini-2.5-flash-001/noverifier_t-30/reddit"},
    # {"agent_config": "noverifier_t-30.yaml", "task_list": "evaluation_harness/tasks_vwa/shopping.txt", "results_dir": "experiments/gemini-2.5-flash-001/noverifier_t-30/shopping"},
    # {"agent_config": "noverifier_t-30.yaml", "task_list": "evaluation_harness/tasks_vwa/classifieds.txt", "results_dir": "experiments/gemini-2.5-flash-001/noverifier_t-30/classifieds"},
]
non_blocking = True  # If false, next launch of parallel_run.py will wait for the previous one to finish.
replace_keys = False

# Params for parallel_run.py instances
start_env_idx = 0
skip_completed = False
tasks_per_process = 2                   # Max num of tasks in each `run.py` execution (i.e., size of a 'batch' of tasks)
max_running_time = -1                   # Max runtime (mins) of each run.py execution: -1 => no limit; 0 => dynamically set based on `avg_run_time_per_task` * `# tasks in the batch`
avg_run_time_per_task = 15*60           # Avg runtime (secs) of each task. Used to compute the max running time if `max-running-time` is set to 0.
shuffle_tasks = True                    # Shuffle the task list before parallel execution
seed = 42                               # Random seed
max_attempts_per_task = 1               # If error on task execution, it will be retried up to `max_attempts_per_task` times (e.g.: 2 => try at most 2 times).
attempt_num = -1                        # Attempt number. Used to identify the attempt for a given task id. Not implemented yet.


# Environment params
domains_to_reset = "all_vwa"                                # Domains to start/reset. If empty, infer from `test_config_dir`
captioner_device = "server-cuda:0"                          # Device to hold the captioner

for run_config in run_configs:
    num_envs = run_config["num_processes"] + 1
    run_config["num_envs_per_process"] = num_envs/run_config["num_processes"]

#===============================================
# Helpers
#===============================================
def build_command(run_config: dict[str, Any]) -> list[str]:
    """
    Build the command list to be executed for the given agent configuration
    and task list.
    """
    agent_config = run_config["agent_config"]
    task_list = run_config["task_list"]
    results_dir = run_config.get("results_dir", "")
    cmd = [
        "python",
        "-u",                # Unbuffered mode for parallel_run.py
        "-m", "scripts.runs.parallel_run",
        "-a", agent_config,
        "-t", task_list,
        "-n", str(run_config["num_processes"]),
        "-b", str(tasks_per_process),
        "-mrt", str(max_running_time),
        "-art", str(avg_run_time_per_task),
        "-d", domains_to_reset,
        "-ma", str(max_attempts_per_task),
        "-cd", captioner_device,
        "-seed", str(seed),
        "-nenvs", str(run_config["num_envs_per_process"]),
        "-r", results_dir,
        "-at", str(attempt_num),
        "-esi", str(run_config.get("start_env_idx", start_env_idx)),
    ]
    if skip_completed:
        cmd.append("-sc")
    if shuffle_tasks:
        cmd.append("-st")
    return cmd

def get_header(run_config: dict[str, Any]) -> str:
    """
    Generate a header string for the per-process log file.
    """
    agent_config = run_config["agent_config"]
    task_list = run_config["task_list"]
    header = (
        f"\n{'-'*80}\n"
        f"Running agent config: {agent_config}\n"
        f"Task list: {task_list}\n"
        f"{'-'*80}\n\n"
    )
    return header

def stream_output(proc: subprocess.Popen[Any], log_file: str) -> None:
    """Stream subprocess output, writing to both global logger and process-specific log."""
    with open(log_file, "a") as lf:
        for line in proc.stdout:  # type: ignore
            print(line.strip())
            lf.write(line)
            lf.flush()

# Global variable to track child processes for cleanup
child_processes = []

# Cleanup function and signal handler
def cleanup_processes() -> None:
    print("Cleaning up processes and resources...", flush=True)
    
    # Send termination signal to all child processes
    print("Sending termination signal to child processes...", flush=True)
    for proc in child_processes:
        try:
            if proc.poll() is None:  # Process is still running
                proc.terminate()
        except Exception as e:
            print(f"Error terminating process {proc.pid}: {e}", flush=True)
    
    # Wait for processes to finish naturally (with timeout)
    print("Waiting for processes to finish gracefully...", flush=True)
    timeout = 60  # 60 seconds timeout
    start_time = time.time()
    
    while time.time() - start_time < timeout:
        if all(proc.poll() is not None for proc in child_processes):
            print("All processes finished gracefully", flush=True)
            break
        time.sleep(1)
    else:
        print("Timeout reached, force killing remaining processes...", flush=True)
        for proc in child_processes:
            try:
                if proc.poll() is None:
                    proc.kill()
            except Exception as e:
                os.system(
                    r"ps aux | grep 'parallel_run' | grep -v grep | awk '{print $2}' | xargs -r kill"
                )
    
    # Kill tmux server to clean up all tmux sessions
    # print("Killing tmux server...", flush=True)
    # os.system("tmux kill-server")
            
    
    # Clean up any temporary files
    if os.path.exists(".temp_files"):
        print("Removing temporary files...", flush=True)
        shutil.rmtree(".temp_files", ignore_errors=True)
    
    print("Cleanup complete.", flush=True)

signal_manager.add_cleanup_function(cleanup_processes)

#===============================================
# Main
#===============================================
try:    
    os.makedirs("logs", exist_ok=True)

    # Start captioner
    if not is_captioner_running():
        print("Starting captioner...", flush=True)    
        result = subprocess.run([
            "python", "-m", "vwa_utils.captioner_utils", 
            "--model_name", "Salesforce/blip2-flan-t5-xl", 
            "--model_device", captioner_device, 
            "--port", "9555", 
            "--endpoint", "http://localhost:9555/caption/", 
            "--conda_env", "vwebarena", 
            "--tmux_session_name", "vwa_captioner"], check=True)
        print(result, flush=True)

    # Copy api_keys.json to api_keys_copy.json; overwrite if exists
    if replace_keys:
        # Run python api_keys.py
        subprocess.run(["python", "api_keys.py"], check=False)

    # Launch parallel runs
    print(f"Launching parallel runs in {'non-blocking' if non_blocking else 'blocking'} mode.", flush=True)
    print(f"Configs: {run_configs}", flush=True)

    if non_blocking:
        # Non-blocking mode: launch each process in a separate thread.
        # TODO: possibly deprecate this given new parallel_run.py orchestrator.
        threads = []
        env_id_start = start_env_idx
        for run_config in run_configs:
            unique_log_file = f"logs/prun_many_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.log"
            
            # Write header to the per-process log file.
            with open(unique_log_file, "a") as lf:
                lf.write(get_header(run_config))
                lf.flush()

            run_config["start_env_idx"] = env_id_start
            # Ensure env ids for next run are non-overlapping with the current run.
            env_id_start += int(run_config["num_envs_per_process"] * run_config["num_processes"]) + 1
            cmd = build_command(run_config)
            proc = subprocess.Popen(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                universal_newlines=True,
                cwd=os.getcwd(),
            )
            # Track the child process for cleanup
            child_processes.append(proc)
            if not run_config.get("block", False):
                t = threading.Thread(target=stream_output, args=(proc, unique_log_file))
                t.start()
                threads.append(t)
                time.sleep(5)
            else:
                t = threading.Thread(target=stream_output, args=(proc, unique_log_file))
                t.start()
                proc.wait()
                t.join()

        # Wait for all non-blocking process threads to finish.
        for t in threads:
            t.join()

        print("All processes finished.", flush=True)
    else:
        # Blocking mode: launch each process sequentially with its own unique log file.
        #TODO: deprecate this given new blocking per config above.
        process_counter = 0
        for run_config in run_configs:
            unique_log_file = f"logs/prun_many_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.log"
            with open(unique_log_file, "a") as lf:
                lf.write(get_header(run_config))
                lf.flush()
                cmd = build_command(run_config)
                proc = subprocess.Popen(
                    cmd,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.STDOUT,
                    universal_newlines=True,
                    cwd=os.getcwd(),
                )
                # Track the child process for cleanup
                child_processes.append(proc)
                
                for line in proc.stdout:  # type: ignore
                    print(line.strip())
                    lf.write(line)
                    lf.flush()

except Exception as e:
    print(f"Error: {e}", flush=True)
finally:
    cleanup_processes()
