# TODO: cleanup of config dicts; consider flat dict for case configs

import argparse
import json
import os
import re
import subprocess
import sys
import warnings
from datetime import datetime
from itertools import product
from typing import Any

from offline_experiments.config.env_configs import all_env_configs
from offline_experiments.config.llm_configs import llm_configs
from utils.logger_utils import logger

# Ignore warnings whose message starts with "Pydantic serializer warnings:"
warnings.filterwarnings("ignore", message="Pydantic serializer warnings:")

from offline_experiments.utils_offline_exper import partial_format
from utils.signal_utils import signal_manager

NUM_RETRIES = 2
DUMP_HTML = False

ann_v = "noSOM-NA"
ann_k = "noSOM"
overwrite = False
max_api_keys = 30
max_batch_runners = 90
# envs_to_run = ["vwa"]
envs_to_run = ["agrb_vwa"]
# envs_to_run = ["vwa", "osw", "agrb_vwa"]
# model_fp = "gemini-2.0-flash-001_zero-random"
# model_v = "gemini-2.0-flash-001_zero-random"

# model_fp = "gemini-2.5-flash-base"
model_fp = "gemini-2.5-flash-zero-random"
model_v = "gemini-2.5-flash-zero-random"

# model_fp = "qwen-2.5-72b-base"
# model_v = "qwen-2.5-72b-base"


# model_fp = "gemini-2.5-flash-thinking_auto-zero-random"
# model_v = "gemini-2.5-flash-thinking_auto-majority"
# model_fp = "llama-4-maverick-zero-random"
# model_v = "llama-4-maverick-zero-random"

# model_fp = "qwen-2.5-72b-zero-random"
# model_v = "qwen-2.5-72b-zero-random"

# model_fp = "gpt-4.1-apr-zero-random"
# model_v = "gpt-4.1-apr-zero-random"

# Run api_keys.py to get the API keys

task_list_vwa = ""
task_list_osw = ""
task_list_agrb_vwa = ""
sort_by_config = True
# task_list_vwa = "offline_experiments/task_list_debug.txt"
# task_list_vwa = "offline_experiments/eval_set_vwa_shopping.txt"

# fmt: off
CASES = [


    {
        "1p_no_k": {
            "eval_criterias": [
                # "four_num_unc_accomplished",
                # "tri_unc",
                # "tri_rev",
                # "tri_random",
                # "tri_rev",
                # "four",
                # "bin",
                "bin",
                "tri",
                # "four_room"
                # "four_num_unc",
                # "four_num_unc_rev",
                # "four_num_unc_random",
                # "three_statement_unc_letter",
                # "four_statement_unc_letter",
                # "four_statement_unc_letter_t_f",
                # "three_option_unc",
                # "four_option_unc_question",
            ],
            "cot_parts": ["basic_cot", "desc_nocompare_v4", "desc_nocompare_v5"],
            "trace_info": ["utt_actions"],
            "add_summ_info": False,
            "add_expectation_info": False,
            "rules": ["rule_efficiency_vision_1p_v6"],
            "sys_prompt": "base",  # aeval_refine, base #

            "additional_config": {
                "img_ann_types": [],
            },

            # "trace_dirs": ["gpt4o-aug", "gemini-2.0-flash-001"],
            # "trace_dirs": ["ui-tars-50s_04-05"],
            "trace_dirs": ["agrb_claude-3.7", "agrb_qwen2.5-vl", "agrb_gpt4o"],
        }
    },
    
    # {
    #     "2p": {
    #         "eval_criterias": [
    #             "bin",
    #             "tri",
    #         ],
    #         "cot_parts": [
    #             # "no_cot",
    #             # "basic_cot",
    #             # "compare",
    #             # "desc_cot_v2",
    #             # # "desc_compare",
    #             # # "desc_compare_v2",
    #             "desc_compare_v4",
    #             "desc_compare_v5",
    #             "desc_compare_v4_no_cot",
    #             # "desc_trace",
    #             # "desc_cot_ss_trace",
    #             # "desc_compare",
    #             # "desc_ss",
    #             # "desc_cot",
    #         ],
    #         "trace_info": ["utt_actions"],
    #         "k_configs": [
    #             # {"conditional": [False], "queries": ["k_2p_expert_v6"], "fp_models": [model_fp], "trace_infos": [""]},
    #             {"conditional": [False], "queries": ["k_2p_expert_v4"], "fp_models": [model_fp], "trace_infos": [""]},
    #             # {"conditional": [False, True], "queries": ["k_2p_expert_v4", "k_2p_vis_expert"], "fp_models": [model_fp, model_fp], "trace_infos": ["actions", "none"]},
    #         ],
    #         "add_summ_info": False,
    #         "add_expectation_info": True,
    #         "rules": ["k_rule_efficiency_v4_6"],  
    #         "sys_prompt": "base",

    #         "additional_config": {
    #             "img_ann_types": [],
    #         },
    #         "trace_dirs": ["agrb_claude-3.7", "agrb_qwen2.5-vl", "agrb_gpt4o"],
    #     }
    # },
 ]

# fmt: on

timestamp = datetime.now().strftime("%Y-%m-%d")
exper_dir = "./offline_experiments"
out_dir_template = f"{exper_dir}/{{env}}/{{trace_source_path}}/{{model_name}}/{{full_config_name}}/{{domain}}"


run_config = {
    "sort_by_config": sort_by_config,
    "overwrite": overwrite,
    "batch_mode": True,
    "max_batch_size": -1,
    "max_api_keys": max_api_keys,
    "max_batch_size_runners": max_batch_runners,
    "task_lists": {"vwa": task_list_vwa, "osw": task_list_osw, "agrb_vwa": task_list_agrb_vwa},
    # "multiprocess_batch_mode": True,
    # "num_processes": 2,
    # "skip_payload": 15 * 1024 * 1024,
    "dump_html": DUMP_HTML,
}

env_configs = {env: all_env_configs[env] for env in envs_to_run}
gen_config_fp = {**llm_configs[model_fp]}
gen_config_v = {**llm_configs[model_v]}


def dump_exper_args(
    config: dict,
    gen_config: dict,
):
    output_dir = config["out_dir"]
    if os.path.exists(f"{output_dir}/exper_args.json"):
        return

    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    with open(f"{output_dir}/exper_args.json", "w") as f:
        args_to_dump = {
            "gen_config": gen_config,
            "date": timestamp,
        }

        if config is not None:
            args_to_dump.update(config)

        json.dump(args_to_dump, f, indent=4)


def build_config_name_first_pass(conditional, query, trace_info):
    k_config_str = ""
    if conditional:
        assert trace_info is not None, "trace_info is required for conditional k_config"
        k_config_str += f"-cond-{trace_info}"
    else:
        k_config_str += "-uncond"

    k_config_str += f"-{query}"

    # Sub one more "--"
    final_str = re.sub(r"-{2,}", "-", k_config_str)

    return f"k-{final_str.strip('-')}{ann_k}"


def build_config_name_verify(
    eval_criteria="", cot_part="", trace_info="", k_config=None, env="", rule=None, case_config=None
):
    config_str = ""
    if eval_criteria:
        config_str += f"{eval_criteria}"
    if cot_part:
        config_str += f"-{cot_part}"
    if trace_info:
        config_str += f"-{trace_info}"

    k_config_str = ""
    if k_config:
        for conditional, query, k_trace_info, fp_model in zip(
            k_config["conditional"], k_config["queries"], k_config["trace_infos"], k_config["fp_models"]
        ):
            k_config_str += f"-{build_config_name_first_pass(conditional, query, k_trace_info)}-{fp_model}"
    k_config_str = f"-{k_config_str.strip('-')}" if k_config_str else ""

    rules_str = f"-{rule}" if rule else ""

    final_str = f"{config_str}{k_config_str}{rules_str}"

    # Sub one more "--"
    final_str = re.sub(r"-{2,}", "-", final_str)

    if case_config:
        sys_p_str = case_config.get("sys_prompt", "")
        if sys_p_str and "base" not in sys_p_str:
            final_str += f"-sys_{sys_p_str}"

    return final_str.strip("-")


def build_all_first_pass_configs():
    model_name = re.sub("/", "-", model_fp.split("/")[-1])
    configs_per_env = {}
    for env in env_configs:
        all_configs = {}
        for case_config in CASES:
            for case_name, case_config in case_config.items():
                for trace_source_shortname in case_config.get("trace_dirs", []):
                    if trace_source_shortname not in env_configs[env]["traces_dirs"]:
                        logger.warning(f"Trace dir {trace_source_shortname} not found in {env} configs. Skipping...")
                        continue

                    trace_dir = env_configs[env]["traces_dirs"][trace_source_shortname]["full_path"]
                    trace_path_template = partial_format(env_configs[env]["trace_path_template"], trace_dir=trace_dir)

                    combinations = []
                    if not case_config.get("k_configs", []):
                        continue

                    for k_config in case_config["k_configs"]:
                        combinations.extend(
                            list(
                                product(
                                    env_configs[env]["domains"],
                                    [k_config],
                                )
                            )
                        )

                    # Create new cases with all combinations
                    for i, (domain, k_config) in enumerate(combinations):
                        for trace_info, conditional, query, fp_model in zip(
                            k_config["trace_infos"], k_config["conditional"], k_config["queries"], k_config["fp_models"]
                        ):
                            if not conditional:
                                trace_info = None
                            key = f"{env}_{trace_source_shortname}_{case_name}-"
                            full_config_name = build_config_name_first_pass(conditional, query, trace_info)
                            key += full_config_name

                            subkey = domain

                            if key not in all_configs:
                                all_configs[key] = {}

                            _k_config = {"conditional": conditional, "query": query, "trace_info": trace_info}
                            config: dict[str, Any] = {}
                            config = {
                                "prompt_args": {
                                    "eval_criteria": "",
                                    "cot_part": "",
                                    "trace_info": trace_info,
                                    "add_summ_info": case_config["add_summ_info"],
                                    "add_expectation_info": case_config["add_expectation_info"],
                                    "k_config": _k_config,
                                    "rules": case_config.get("rules", []),
                                },
                                "trace_path_template": partial_format(trace_path_template, domain=domain),
                                "env": env,
                                "domain": domain,
                                # "additional_config": env_configs[env].get("additional_config", {}),
                                "additional_config": case_config.get("additional_config", {}),
                                "meta_data": case_config,  # TODO: clean this up
                            }
                            out_dir = f"{out_dir_template.format(env=env, model_name=model_name, domain=domain, trace_source_path=trace_source_shortname, full_config_name=full_config_name)}"
                            config["out_dir"] = out_dir
                            all_configs[key][subkey] = config
            configs_per_env[env] = all_configs
    return configs_per_env


def build_all_eval_configs():
    model_name = re.sub("/", "-", model_v.split("/")[-1])
    configs_per_env = {env: {} for env in env_configs}
    for env in env_configs:
        all_configs = {}
        for case_config in CASES:
            for case_name, case_config in case_config.items():
                for trace_source_shortname in case_config.get("trace_dirs", []):
                    if trace_source_shortname not in env_configs[env]["traces_dirs"]:
                        logger.warning(f"Trace dir {trace_source_shortname} not found in {env} configs. Skipping...")
                        continue

                    trace_dir = env_configs[env]["traces_dirs"][trace_source_shortname]["full_path"]
                    trace_path_template = partial_format(env_configs[env]["trace_path_template"], trace_dir=trace_dir)

                    # Generate all combinations for each case
                    combinations = list(
                        product(
                            env_configs[env]["domains"],
                            case_config["eval_criterias"],
                            case_config["cot_parts"],
                            case_config["trace_info"],
                            case_config.get("k_configs", [None]),
                            case_config.get("rules", [None]),
                        )
                    )

                    # Create new cases with all combinations
                    for i, (domain, eval_criteria, cot_part, trace_info, k_config, rule) in enumerate(combinations):
                        key = f"{env}__{trace_source_shortname}__{case_name}-"
                        config_name = build_config_name_verify(
                            eval_criteria, cot_part, trace_info, k_config, env, rule, case_config
                        )
                        key += config_name

                        subkey = domain
                        if key not in all_configs:
                            all_configs[key] = {}

                        config: dict[str, Any] = {}
                        full_config_name = f"{case_name}-{config_name}{ann_v}"
                        config = {
                            "prompt_args": {
                                "eval_criteria": eval_criteria,
                                "cot_part": cot_part,
                                "trace_info": trace_info,
                                "add_summ_info": case_config["add_summ_info"],
                                "add_expectation_info": case_config["add_expectation_info"],
                                "k_config": k_config.copy() if k_config else None,
                                "rule": rule if rule else None,
                            },
                            "env": env,
                            "domain": domain,
                            "trace_path_template": partial_format(trace_path_template, domain=domain),
                            "out_dir": f"{out_dir_template.format(env=env, model_name=model_name, domain=domain, trace_source_path=trace_source_shortname, full_config_name=full_config_name)}",
                            "additional_config": case_config.get("additional_config", {}),
                            "meta_data": case_config,  # TODO: clean this up
                        }
                        if k_config:
                            fp_config_names = [
                                build_config_name_first_pass(conditional, query, trace_info)
                                for conditional, query, trace_info in zip(
                                    k_config["conditional"], k_config["queries"], k_config["trace_infos"]
                                )
                            ]
                            cached_k_dirs = [
                                f"{out_dir_template.format(domain=domain, env=env, model_name=fp_model, trace_source_path=trace_source_shortname, full_config_name=fp_config_name)}"
                                for fp_config_name, fp_model in zip(fp_config_names, k_config["fp_models"])
                            ]
                            config["prompt_args"]["k_config"]["cached_k_dir"] = cached_k_dirs
                        all_configs[key][subkey] = config

        configs_per_env[env] = all_configs
    return configs_per_env


def run_and_wait(command, logfile_path):
    if logfile_path:
        with open(logfile_path, "w") as logfile:
            process = subprocess.Popen(
                command,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                text=True,
            )
            signal_manager.register_termination_signals(lambda: process.terminate())  # type:ignore
            while True:
                output_line = process.stdout.readline()  # type:ignore
                if output_line == "" and process.poll() is not None:
                    break
                if output_line:
                    print(output_line, end="")  # Print to command line
                    logfile.write(output_line)
                    logfile.flush()
            process.wait()
    else:
        process = subprocess.Popen(
            command,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
        )
        signal_manager.register_termination_signals(lambda: process.terminate())  # type:ignore
        while True:
            output_line = process.stdout.readline()  # type:ignore
            if output_line == "" and process.poll() is not None:
                break
            if output_line:
                print(output_line, end="")  # Print to command line
        process.wait()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--fp", action="store_true")
    parser.add_argument("--v", action="store_true")
    args = parser.parse_args()

    # Log the proccess ID
    logger.info(f"Process ID: {os.getpid()}")

    if sys.gettrace():
        pass
        fp_configs = build_all_first_pass_configs()
        v_configs = build_all_eval_configs()
        args.fp = True
        args.v = True

    timestamp_time = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
    os.makedirs(f"{exper_dir}/logs", exist_ok=True)

    num_retries = 1 if overwrite else NUM_RETRIES

    run_and_wait(["python", "api_keys.py"], None)

    if args.fp:
        log_file_path = f"{exper_dir}/logs/{model_fp}-{timestamp_time}.log"
        for i in range(num_retries):
            logger.info(f"Running first pass: attempt {i + 1} of {num_retries}")
            run_and_wait(
                ["python", "-u", "-m", "offline_experiments.first_pass"],
                log_file_path,
            )

    if args.v:
        log_file_path = f"{exper_dir}/logs/{model_v}-{timestamp_time}.log"
        for i in range(num_retries):
            logger.info(f"Running verify: attempt {i + 1} of {num_retries}")
            run_and_wait(
                ["python", "-u", "-m", "offline_experiments.verify"],
                log_file_path,
            )
