# fmt: off

# ===============================================
# LINK: System Prompt parts
# ===============================================
import random
import re

from .common_prompts import common_prompts
from .prompts import prompts_per_env

new_criterias_id: list[str] = ["num", "degree", "letter", "option"]

def reassign_criteria(criteria_id: str, criteria_prompt: str, shuffle_type: str) -> str:
    if criteria_id in new_criterias_id:
        # Lettered options (A., B., ...)
        if "letter" in criteria_id:
            lines = criteria_prompt.strip().split("\n")
            option_start = None
            for i, line in enumerate(lines):
                if re.match(r"^[A-Z]\. ", line.strip()):
                    option_start = i
                    break
            if option_start is None:
                raise ValueError(f"Error shuffling criteria: No match found for criteria_id: {criteria_id}")
            intro_section = "\n".join(lines[:option_start])
            option_lines = lines[option_start:]
            # Remove the labels
            option_texts = [re.sub(r"^[A-Z]\. ", "", li.strip(), 1) for li in option_lines]
            if shuffle_type == "random":
                random.shuffle(option_texts)
            elif shuffle_type == "rev":
                option_texts = option_texts[::-1]
            # Reassign labels
            labels = [chr(ord('A') + i) for i in range(len(option_texts))]
            reassigned_options = [f"{label}. {text}" for label, text in zip(labels, option_texts)]
            new_criteria = intro_section + ("\n" if intro_section else "") + "\n".join(reassigned_options)
            return new_criteria
        # Roman numeral options (I., II., ...)
        elif "option" in criteria_id:
            lines = criteria_prompt.strip().split("\n")
            option_start = None
            for i, line in enumerate(lines):
                if re.match(r"^(I{1,3}|IV|V|VI{0,3}|IX|X)\. ", line.strip()):
                    option_start = i
                    break
            if option_start is None:
                raise ValueError(f"Error shuffling criteria: No match found for criteria_id: {criteria_id}")
            intro_section = "\n".join(lines[:option_start])
            option_lines = lines[option_start:]
            # Remove the labels
            option_texts = [re.sub(r"^(I{1,3}|IV|V|VI{0,3}|IX|X)\. ", "", li.strip(), 1) for li in option_lines]
            if shuffle_type == "random":
                random.shuffle(option_texts)
            elif shuffle_type == "rev":
                option_texts = option_texts[::-1]
            # Reassign labels (Roman numerals)
            roman_labels = [
                "I", "II", "III", "IV", "V", "VI", "VII", "VIII", "IX", "X"
            ]
            reassigned_options = [f"{roman_labels[i]}. {text}" for i, text in enumerate(option_texts)]
            new_criteria = intro_section + ("\n" if intro_section else "") + "\n".join(reassigned_options)
            return new_criteria
        # Numeric/degree: shuffle/reverse score lines, reassign score labels
        elif ("num" in criteria_id) or ("degree" in criteria_id):
            # Find the score lines (start with 'Score ...:')
            lines = criteria_prompt.strip().split("\n")
            score_start = None
            for i, line in enumerate(lines):
                if re.match(r"^Score ", line.strip()):
                    score_start = i
                    break
            if score_start is None:
                raise ValueError(f"Error shuffling criteria: No score section found for criteria_id: {criteria_id}")
            intro_section = "\n".join(lines[:score_start])
            score_lines = lines[score_start:]
            # Extract score labels and texts
            score_labels = []
            score_texts = []
            for lin in score_lines:
                m = re.match(r"^(Score [^:]+:)(.*)$", lin.strip())
                if m:
                    score_labels.append(m.group(1))
                    score_texts.append(m.group(2).strip())
            if not score_labels:
                raise ValueError(f"Error shuffling criteria: No valid score lines for criteria_id: {criteria_id}")
            # Shuffle or reverse the texts
            if shuffle_type == "random":
                random.shuffle(score_texts)
            elif shuffle_type == "rev":
                score_texts = score_texts[::-1]
            # Reassign score labels in order
            reassigned_scores = [f"{label}{' ' if text else ''}{text}" for label, text in zip(score_labels, score_texts)]
            new_criteria = intro_section + ("\n" if intro_section else "") + "\n".join(reassigned_scores)
            return new_criteria
        else:
            return criteria_prompt
    return criteria_prompt

def shuffle_criteria(criteria_id: str, criteria_prompt: str, shuffle_type: str) -> str:
    try:
        new_criteria = criteria_prompt
        # Numeric or degree criteria
        if ("num" in criteria_id) or ("degree" in criteria_id):
            pattern = re.compile(r"(Score[\s\S]*)")
            match = pattern.search(criteria_prompt)
            if not match:
                raise ValueError(f"Error shuffling criteria: No match found for criteria_id: {criteria_id}")
            score_section = match.group(1)
            intro_section = new_criteria.strip()[: match.start()]

            # Shuffle the score section
            lines = score_section.strip().split("\n")
            if shuffle_type == "random":
                random.shuffle(lines)
            elif shuffle_type == "rev":
                lines = lines[::-1]

            shuffled_section = "\n".join(lines)
            new_criteria = intro_section + "\n" + shuffled_section

        # Lettered options (A., B., ...)
        elif "letter" in criteria_id:
            # Find the first line that starts with a capital letter and a period (A., B., ...)
            lines = new_criteria.strip().split("\n")
            option_start = None
            for i, line in enumerate(lines):
                if re.match(r"^[A-Z]\. ", line.strip()):
                    option_start = i
                    break
            if option_start is None:
                raise ValueError(f"Error shuffling criteria: No match found for criteria_id: {criteria_id}")
            intro_section = "\n".join(lines[:option_start])
            option_lines = lines[option_start:]
            if shuffle_type == "random":
                random.shuffle(option_lines)
            elif shuffle_type == "rev":
                option_lines = option_lines[::-1]
            shuffled_section = "\n".join(option_lines)
            new_criteria = intro_section + ("\n" if intro_section else "") + shuffled_section

        # Roman numeral options (I., II., ...)
        elif "option" in criteria_id:
            # Find the first line that starts with a Roman numeral and a period (I., II., ...)
            lines = new_criteria.strip().split("\n")
            option_start = None
            for i, line in enumerate(lines):
                if re.match(r"^(I{1,3}|IV|V|VI{0,3}|IX|X)\. ", line.strip()):
                    option_start = i
                    break
            if option_start is None:
                raise ValueError(f"Error shuffling criteria: No match found for criteria_id: {criteria_id}")
            intro_section = "\n".join(lines[:option_start])
            option_lines = lines[option_start:]
            if shuffle_type == "random":
                random.shuffle(option_lines)
            elif shuffle_type == "rev":
                option_lines = option_lines[::-1]
            shuffled_section = "\n".join(option_lines)
            new_criteria = intro_section + ("\n" if intro_section else "") + shuffled_section

        # Default: shuffle all lines
        else:
            lines = new_criteria.strip().split("\n")
            if shuffle_type == "random":
                random.shuffle(lines)
            elif shuffle_type == "rev":
                lines = lines[::-1]
            shuffled_section = "\n".join(lines)
            new_criteria = shuffled_section
        return new_criteria
    except Exception as e:
        print(f"Error shuffling criteria: {e}")
        return criteria_prompt
    
    


def safe_format(string_template: str, fill_with: str = "", **kwargs) -> str:
    """
    Formats a given template using the provided keyword arguments.
    Missing keys in the template are replaced with an empty string.

    Args:
        template (str): The string template with placeholders.
        **kwargs: Key-value pairs for formatting.

    Returns:
        str: The formatted string with missing keys as empty strings.
    """

    class DefaultDict(dict):
        def __missing__(self, key):
            return fill_with

    return string_template.format_map(DefaultDict(**kwargs))


def get_query_prompt_eval(
    env: str,
    eval_criteria: str = "", 
    cot_part: str = "", 
    add_retrieval_step: bool = False,
) -> str:
    prompts = prompts_per_env[env]
    cot_parts = prompts["cot_parts"]
    eval_criterias = prompts["eval_criterias"]
    k_retrieval_1p = prompts["k_retrieval_queries"]["k_1p"]
    response_format = prompts["response_format"]
    eval_prompt_templates = prompts["eval_prompt_templates"]
    
    # Get CoT parts
    if cot_part:
        _cot_parts = cot_parts[cot_part].strip()
    else:
        _cot_parts = ""

    # Get eval criteria
    if eval_criteria:
        if "random" in eval_criteria:
            _eval_criteria = eval_criterias[eval_criteria.replace("_random", "").strip("_")].strip()
            _eval_criteria = shuffle_criteria(eval_criteria, _eval_criteria, "random")
        elif "rev" in eval_criteria:
            _eval_criteria = eval_criterias[eval_criteria.replace("_rev", "").strip("_")].strip()
            _eval_criteria = shuffle_criteria(eval_criteria, _eval_criteria, "rev")
        elif "reassign_random" in eval_criteria:
            _eval_criteria = eval_criterias[eval_criteria.replace("_reassign_random", "").strip("_")].strip()
            _eval_criteria = reassign_criteria(eval_criteria, _eval_criteria, "random")
        elif "reassign_rev" in eval_criteria:
            _eval_criteria = eval_criterias[eval_criteria.replace("_reassign_rev", "").strip("_")].strip()
            _eval_criteria = reassign_criteria(eval_criteria, _eval_criteria, "rev")
            
        else:
            _eval_criteria = eval_criterias[eval_criteria].strip()
        if any(c in eval_criteria for c in new_criterias_id):
            eval_prompt_template = eval_prompt_templates["ins"]
        else:
            eval_prompt_template = eval_prompt_templates["cat"]
    else:
        _eval_criteria = ""
        eval_prompt_template = ""


    # If one-pass, but retrieving K, add K retrieval step before CoT
    k_step = k_retrieval_1p.strip() if add_retrieval_step else ""

    _response_format = response_format.format(k_retrieval_step=k_step, cot_parts=_cot_parts).strip()

    # Get eval prompt
    eval_prompt = safe_format(
        eval_prompt_template,
        eval_criteria=_eval_criteria,
        response_format=_response_format,
    )
    return eval_prompt.strip()

def _get_trace_info_key(trace_info: str) -> str:
    key=""
    if "utt" in trace_info and "actions" in trace_info:        
        key = "utt"

    elif "actions" in trace_info and "last_u" in trace_info:
        key = "actions"

    elif "utt" in trace_info:
        key = "utt"

    elif "actions" in trace_info:
        key = "actions"

    elif not trace_info or trace_info.lower() == "none":
        key = "none"

    else:
        raise ValueError(f"Invalid trace_info: {trace_info}")
    return key

def _get_image_info_key(img_ann_types: list[str]) -> str:
    if not img_ann_types:
        return "image_info_none"
    if "som" and "dot" in img_ann_types:
        key = "image_info_som_dot"
    elif "som" in img_ann_types:
        key = "image_info_som"
    elif "dot" in img_ann_types:
        key = "image_info_dot"
    else:
        raise NotImplementedError(f"Annotation type {img_ann_types} not implemented")
    return key

def get_sys_prompt_verifier(
    env: str,
    trace_info: str = "", 
    add_summ_info: bool = False, 
    add_k_info: bool = False,
    two_pass: bool = False,
    prompt_config: dict = {},
    config: dict = {},
) -> str:
    prompts = prompts_per_env[env]
    trace_infos = prompts["trace_infos"]
    summ_info = prompts["summ_info"]
    k_info = prompts["k_info"]

    # FIXME: clean this up
    if "rubrics" in prompt_config.get("rule", ""):
        k_info = prompts["k_info_rubrics"]
    elif "strict" in prompt_config.get("rule", ""):
        k_info = prompts["k_info_strict"]

    if prompt_config.get("rule", ""):
        # TODO: Add rules dynamically
        rule_key = prompt_config["rule"]
        rule = prompts[rule_key]

    elif prompt_config.get("rules", ""):
        rule_key = prompt_config["rules"][0]
        rule = prompts[rule_key]
    else:
        rule = ""

    key_sys_prompt = config.get("meta_data", {}).get("sys_prompt", "base") # TODO: clean this up

    if "aeval_refine" in key_sys_prompt:
        return prompts["sys_prompts_critic"][key_sys_prompt]

    if "base" in key_sys_prompt and "agrb" in config["env"]:
        key_sys_prompt = "base_agrb"


    sys_prompt_critic_base = prompts["sys_prompts_critic"][key_sys_prompt]  

    key_trace_info = _get_trace_info_key(trace_info)
    _trace_info = trace_infos[key_trace_info]
    _summ_info = summ_info if add_summ_info else ""

    if add_k_info:
        if two_pass:
            _k_info = k_info
            _rule = rule
        else:
            _k_info = ""
            _rule = rule
    else:
        _k_info = ""
        _rule = rule

    img_info_key = _get_image_info_key(config.get("additional_config", {}).get("img_ann_types", []))
    img_info = prompts["image_infos"][img_info_key]
    
    sys_prompt = safe_format(
        sys_prompt_critic_base,
        trace_info=_trace_info,
        summ_info=_summ_info,
        k_info=_k_info,
        rule=_rule,
        image_info=img_info,
    )
    return sys_prompt.strip()

def get_prompts_first_pass(
    env: str,
    config,
) -> tuple[str, str]:
    prompt_config = config["prompt_args"]
    k_prompt_config = prompt_config["k_config"]
    query: str = k_prompt_config["query"]

    k_retrieval_queries = prompts_per_env[env]["k_retrieval_queries"]

    if "expert" in query:
        k_retrieval_query = k_retrieval_queries[query]
        _sys_prompt = prompts_per_env[env]["sys_prompts_expert"][query]
        # FIXME: clean this up
        trace_info_fp = config.get("meta_data", {}).get("trace_info_fp", "")
        
        if trace_info_fp:  
            trace_info_fp = prompts_per_env[env]["trace_infos"][trace_info_fp].strip()
            _sys_prompt = safe_format(
                _sys_prompt,
                trace_info=trace_info_fp,
            )
    else:
        k_retrieval_query = k_retrieval_queries[query]
        _sys_prompt = get_sys_prompt_verifier(
            env=env,
            trace_info=config["prompt_args"]["trace_info"],
            add_summ_info=config["prompt_args"]["add_summ_info"],
            add_k_info=config["prompt_args"]["add_expectation_info"],
            two_pass=True,
            prompt_config=prompt_config,
            config=config,
        )
        
    return k_retrieval_query, _sys_prompt

def get_verifier_prompts(
    config: dict,
) -> dict[str, str]:
    env = config["env"]
    k_config = config["prompt_args"]["k_config"]
    prompt_args = config["prompt_args"]
    add_expectation_info = prompt_args.get("add_expectation_info", False)
    eval_criteria = prompt_args.get("eval_criteria", "")
    cot_part = prompt_args.get("cot_part", "")
    trace_info = prompt_args.get("trace_info", "")
    add_summ_info = prompt_args.get("add_summ_info", False)
    if k_config:
        two_pass = True
    else:
        two_pass = False

    add_retrieval_step: bool = True if (add_expectation_info and not two_pass) else False

    key_sys_prompt = config.get("meta_data", {}).get("sys_prompt", "base")
    eval_prompt = get_query_prompt_eval(env, eval_criteria, cot_part, add_retrieval_step)

    if "aeval_refine" in key_sys_prompt:
        eval_prompt = ""
        
    sys_prompt = get_sys_prompt_verifier(env, trace_info, add_summ_info, add_expectation_info, two_pass, prompt_args, config)
    return {"eval_prompt": eval_prompt, "sys_prompt": sys_prompt}


def add_noise_to_k(
    k: str, 
    noise_key: str = "latin_noise", 
    noise_injections: dict = common_prompts["noise_injections"],
) -> str:
    """
    Add noise to a given text by randomly inserting a noise line between content lines.
    Args:
        k: The input text to add noise to
        noise_key: Key to select which noise to use from noise_injections
        noise_injections: Dictionary containing different types of noise
    Returns:
        The text with randomly inserted noise
    """
    lines = k.split("\n")
    noise: str = f"{noise_injections[noise_key]}"
    # Choose a random line index to insert noise (can be at the start, between, or end)
    insert_idx = random.randint(0, len(lines))
    # Choose a random noise line from the noise list
    # Insert the noise line
    if insert_idx < len(lines) and not lines[insert_idx] == "":
        noise = noise + "\n"

    if insert_idx > 0 and not lines[insert_idx-1] == "":
        noise = "\n" + noise

    new_lines = lines[:insert_idx] + [noise] + lines[insert_idx:]
    return "\n".join(new_lines)


def get_k_injection(env: str, k_query: str) -> str:
    prompts = prompts_per_env[env]
    k_injections = prompts["k_injections"]
    k_injection = k_injections[k_query]
    return k_injection
