import os
import re
from string import Formatter

from bs4 import BeautifulSoup

from offline_experiments.agrb_specific import get_trace_data_agrb
from offline_experiments.osworld_specific import get_intent_message_osw, get_trace_data_osw, get_trajectory_osw
from utils.string_utils import clean_spaces

from .vwa_specific import get_intent_message_vwa, get_trace_data_vwa, get_trajectory_vwa


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 partial_format(string_template: str, **kwargs) -> str:
    """
    Partially formats a string by replacing only those placeholders
    for which corresponding keyword arguments are provided.
    Placeholders with missing keys remain in the output.

    Example:
        template = "Hello, {name}! Today is {day}."
        result = partial_format(template, name="Alice")
        # result will be "Hello, Alice! Today is {day}."

    Args:
        string_template (str): The string template with placeholders.
        **kwargs: Keyword arguments with values for placeholders.

    Returns:
        str: A partially formatted string.
    """
    formatter = Formatter()
    result_str = ""

    for literal_text, field_name, format_spec, conversion in formatter.parse(string_template):
        # Add the literal text between placeholders
        result_str += literal_text

        # If there is no field, then nothing to format here.
        if field_name is None:
            continue

        if field_name in kwargs:
            value = kwargs[field_name]
            if conversion:
                if conversion == "r":
                    value = repr(value)
                elif conversion == "s":
                    value = str(value)
                elif conversion == "a":
                    value = ascii(value)
            try:
                format_spec = format_spec or ""
                formatted_value = format(value, format_spec)
            except Exception:
                # Fallback: if formatting fails, just convert to string.
                formatted_value = str(value)
            result_str += formatted_value
        else:
            # Key not provided. Reconstruct the original placeholder including
            # any conversion flags and format specifiers.
            placeholder = "{" + field_name
            if conversion:
                placeholder += "!" + conversion
            if format_spec:
                placeholder += ":" + format_spec
            placeholder += "}"
            result_str += placeholder
    return result_str


def is_aeval_refine(config):
    key_sys_prompt = config.get("meta_data", {}).get("sys_prompt", "base")
    if "aeval_refine" in key_sys_prompt:
        return True
    else:
        return False


def get_intent_message(
    config,
    trace_data,
    add_state_idxs: list[int] = [],
    state_img_intros: list[str] = [],
    aeval_refine: bool = False,
):
    if "vwa" in config["env"]:
        return get_intent_message_vwa(
            trace_data=trace_data,
            add_state_idxs=add_state_idxs,
            state_img_intros=state_img_intros,
            aeval_refine=is_aeval_refine(config),
        )

    elif config["env"] == "osw":
        return get_intent_message_osw(
            trace_data=trace_data,
            add_state_idxs=add_state_idxs,
            state_img_intros=state_img_intros,
        )
    else:
        raise ValueError(f"Environment {config['env']} not supported")


def get_trace_data(config, task_id):
    env = config["env"]
    trajectory_path = config["trace_path_template"].format(task_id=task_id)
    if env == "vwa":
        return get_trace_data_vwa(trajectory_path, task_id, config)

    elif env == "osw":
        return get_trace_data_osw(trajectory_path, task_id)

    elif "agrb" in env:
        return get_trace_data_agrb(trajectory_path, env, config)


def get_trajectory_msgs(config, trace_data):
    env = config["env"]

    if "vwa" in env:
        return get_trajectory_vwa(config, trace_data, is_aeval_refine(config))

    elif env == "osw":
        return get_trajectory_osw(config, trace_data)

    else:
        raise ValueError(f"Environment {env} not supported")


def parse_evaluation(response: str) -> list[str]:
    splitters = ["EVALUATION:", "FEEDBACK:"]
    utterances = []
    splitters_group = "|".join(map(re.escape, splitters))
    for splitter in splitters:
        pattern = rf"{splitter}(.*?)(?:\n|{splitters_group}|$)"
        # pattern = rf"(?:\*+|#+|\s*)({splitters_group}):\s*(.*?)(?=\n(?:\s*(?:\*+|#+|\s*)({splitters_group}):)|$)"

        match = re.search(pattern, response, re.IGNORECASE)
        if match:
            utterances.append(clean_spaces(match.group(1)))
        else:
            utterances.append("error")
            # raise ValueError(f"Cannot find {splitter} in {response}")
    return utterances


def get_response_from_html_file(file_path: str) -> str:
    if not os.path.exists(file_path):
        return ""

    try:
        with open(file_path, "r", encoding="utf-8") as file:
            html_content = file.read()

        soup = BeautifulSoup(html_content, "html.parser")

        # Get all section-like elements
        sections = soup.find_all("div", class_="section")

        if sections:
            last_section = sections[-1]
            all_text = last_section.get_text(separator="\n", strip=True)
            all_text = re.sub(r"ASSISTANT\n", "", all_text)
            return all_text
        else:
            return ""

    except Exception as e:
        print(f"Error: {e}")
        return ""


def get_response_from_txt_file(file_path: str) -> str:
    if not os.path.exists(file_path):
        return ""

    try:
        with open(file_path, "r", encoding="utf-8") as file:
            content = file.read()

        pattern_match = r"==================================\nGENERATION\n==================================(.*)"
        match = re.search(pattern_match, content, re.DOTALL)
        if match:
            response = match.group(1)
            response = response.strip()
            response = re.sub(r"CONTENT TYPE: .*\n", "", response)
            return response
        else:
            return ""
    except Exception as e:
        print(f"Error: {e}")
        return ""


def get_response_from_file(file_path: str) -> str:
    if file_path.endswith(".html"):
        return get_response_from_html_file(file_path)
    elif file_path.endswith(".txt"):
        return get_response_from_txt_file(file_path)
    else:
        raise ValueError(f"Unsupported file type: {file_path}")
