import asyncio
import json
import os
import re
import sys

from llms.llm_utils import visualize_prompt
from llms.prompt_utils import get_conversation_payload_size, get_messages
from utils.logger_utils import logger

if __name__ == "__main__" and not __package__:  # @debug
    # Insert the parent directory into sys.path so that the package can be found
    parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
    sys.path.insert(0, parent_dir)
    # Manually set the package name so that relative imports work
    __package__ = "offline_experiments"


from .build_prompt import add_noise_to_k, get_k_injection, get_verifier_prompts
from .config_run import build_all_eval_configs, gen_config_v, run_config
from .runners import run_batch_mode, run_sequential
from .utils_offline_exper import get_intent_message, get_response_from_file, get_trace_data, get_trajectory_msgs

cache_trajectory_msgs = {}
cache_k_responses = {}


def _get_k_from_cache(cache_key: str):
    try:
        if cache_key and cache_key in cache_k_responses:
            return cache_k_responses[cache_key]
        return None
    except Exception as _:
        return None


def get_k_resp_from_cache(
    file_path: str = "",
):
    # Try to get K from cache
    k_resp = _get_k_from_cache(file_path)
    if k_resp:
        return k_resp

    if file_path:
        k_resp = get_response_from_file(file_path)
        cache_k_responses[file_path] = k_resp
        return k_resp

    k_resp = get_response_from_file(file_path)
    if k_resp:
        cache_k_responses[file_path] = k_resp
    return k_resp


def extract_generations(k_resp: str) -> list[str]:
    """
    Extract all generations from a k_resp string that contains generation markers.

    Args:
        k_resp: String containing generation markers like "---------- GENERATION 0 ----------"

    Returns:
        List of strings, where each string is the content between generation markers
    """
    pattern_match = re.search(r"---------- GENERATION \d+ ----------\n", k_resp)
    if not pattern_match:
        return [k_resp.strip()]

    # Split by generation markers and extract content
    generations = re.split(r"---------- GENERATION \d+ ----------\n", k_resp)
    # Remove empty first element if it exists
    if generations and not generations[0].strip():
        generations = generations[1:]
    # Strip whitespace from each generation
    return [gen.strip() for gen in generations if gen.strip()]


def build_llm_call_args(task_id, config, run_config) -> tuple[list[dict] | None, str | None, str | None, str | None]:
    try:
        conversation_dir = f"{config['out_dir'].strip('-')}/conversation"
        usage_dir = f"{config['out_dir'].strip('-')}/usage"

        combined_ks = []
        if config["prompt_args"].get("k_config"):
            k_config = config["prompt_args"]["k_config"]

            k_resps_per_query = []
            for k_dir_name, query in zip(k_config["cached_k_dir"], k_config["queries"]):
                file_path = f"{k_dir_name}/conversation/{task_id}.txt"
                cached_k_resp = get_k_resp_from_cache(file_path=file_path)
                if not cached_k_resp:
                    logger.error(
                        f"Failed to build prompt for {config}:\n  task_id: {task_id}\n  Msg: Unable to get first pass response from {file_path}"
                    )
                    return None, None, None, None

                k_injection = get_k_injection(config["env"], query)
                # Get (possible) N generations for this first-step query (used in case of majority voting with multiple K retrievals)
                k_resps = extract_generations(cached_k_resp)
                final_ks = [k_injection.format(k=k_resp) for k_resp in k_resps]
                k_resps_per_query.append(final_ks)

            # Combine the generations for each query
            # Find the maximum number of generations across all queries
            max_generations = max(len(final_ks) for final_ks in k_resps_per_query) if k_resps_per_query else 0

            # Create combined_ks by taking the ith element from each query, or the last if it doesn't exist
            for i in range(max_generations):
                combined_k = []
                for final_ks in k_resps_per_query:
                    # If the ith element exists, use it; otherwise use the last element
                    if i < len(final_ks):
                        combined_k.append(final_ks[i])
                    else:
                        combined_k.append(final_ks[-1])  # Use the last element
                combined_ks.append(combined_k)

        trace_data = get_trace_data(config, task_id)
        if not trace_data:
            logger.error(
                f"Failed to build prompt for {config}:\n  task_id: {task_id}\n  Msg: Unable to get trace data from {config['trace_path_template']}"
            )
            return None, None, None, None

        msg_intent = get_intent_message(config, trace_data, add_state_idxs=[], state_img_intros=[])
        trajectory_msgs = get_trajectory_msgs(config, trace_data)

        verifier_prompts = get_verifier_prompts(config)
        sys_prompt, eval_prompt = verifier_prompts["sys_prompt"], verifier_prompts["eval_prompt"]

        full_prompts = []
        if config["prompt_args"].get("k_config"):
            for combined_k in combined_ks:
                full_prompt = [
                    {"role": "system", "content": sys_prompt},
                    msg_intent,
                    trajectory_msgs,
                    *combined_k,
                    eval_prompt,
                ]
                full_prompts.append(full_prompt)
        else:
            full_prompts.append([{"role": "system", "content": sys_prompt}, msg_intent, trajectory_msgs, eval_prompt])

        if run_config.get("skip_payload", 0):
            payload_size = get_conversation_payload_size(get_messages(full_prompts[0]))
            if payload_size > run_config["skip_payload"]:
                logger.info(
                    f"Skipping {task_id} because payload size {payload_size / 1024 / 1024} MB is greater than {run_config['skip_payload'] / 1024 / 1024} MB"
                )
                return None, None, None, None
        logger.info(f"VERIFY: Finished building llm call args for task {task_id}, config {config}")

        if len(combined_ks) > 1:
            call_ids = [f"{task_id}_k-{i}" for i in range(len(combined_ks))]
        else:
            call_ids = [task_id] * len(full_prompts)

        conversation_dirs = [conversation_dir] * len(full_prompts)
        if sys.gettrace():
            visualize_prompt(full_prompts[0], "debug.html")
        usage_dirs = [usage_dir] * len(full_prompts)

        if not run_config["overwrite"]:
            for call_id in call_ids:
                full_conversation_path = f"{conversation_dir}/{call_id}.txt"
                if os.path.exists(full_conversation_path):
                    logger.info(f"Skipping {call_id} because {full_conversation_path} exists.")
                    return None, None, None, None

        return full_prompts, conversation_dirs, usage_dirs, call_ids  # type: ignore

    except Exception as e:
        logger.error(f"Failed to build prompt for {config}:\n  task_id: {task_id}\n  Msg: {e}")
        return None, None, None, None


if __name__ == "__main__":
    batch_mode = run_config["batch_mode"]
    all_configs = build_all_eval_configs()

    if len(all_configs) == 0:
        logger.info("No configs to run")
        exit()

    if batch_mode:
        logger.info("Running verify in batch mode")
        run_batch_mode(
            all_configs,
            run_config=run_config,
            gen_config=gen_config_v,
            build_llm_call_args_fn=build_llm_call_args,
        )
    else:
        logger.info("Running verify in sequential mode")
        asyncio.run(
            run_sequential(
                all_configs,
                run_config=run_config,
                gen_config=gen_config_v,
                build_llm_call_args_fn=build_llm_call_args,
            )
        )
