import json
import os
import re
from enum import Enum
from pathlib import Path
from typing import Any

from agent.agent import Agent
from agent.prompt_constructor import LMParsingError, PromptConstructor
from llms.prompt_utils import get_message, get_messages
from llms.types import Message
from utils.logger_utils import logger
from utils.string_utils import clean_spaces, partial_format
from utils.trajectory_view import TrajectoryView
from utils.types import ImageInput


# Enum for Critic Modes
class ReflexionMode(Enum):
    SGV = "sgv"
    NO_SGV = "no_sgv"


class EvalMode(Enum):
    ORACLE = "oracle"
    MODEL = "model"


def load_reflections_from_disk(log_file: str, domain_task_id: str = "") -> dict[str, list[str]]:
    reflection_data: dict[str, list[dict[str, Any]]] = {}
    if not os.path.exists(log_file):
        return {}

    with open(log_file, "r") as f:
        reflection_data = json.load(f)

    reflections: dict[str, list[str]] = {}

    for _domain_task_id, data in reflection_data.items():
        if domain_task_id and _domain_task_id != domain_task_id:
            continue
        reflections[_domain_task_id] = [reflection["reflection"] for reflection in data]

    return reflections


class ReflexionPromptConstructor(PromptConstructor):
    def __init__(self, lm_config: dict[str, Any], agent_config: dict[str, Any]):
        super().__init__(lm_config, agent_config)
        self.mode: str = ReflexionMode(agent_config["mode"]).value
        self.parsing_error_msg_key: str = self.instruction["meta_data"]["parsing_error_msg_key"]
        self.parsing_error_msg_template: str = self.instruction["parsing_error_msg_template"]
        self.agent_config = agent_config
        self.reflection_memory: dict[str, list[str]] = {}
        self.image_info: str = agent_config.get("image_info", "som")

        if agent_config.get("previous_reflections_path"):
            self.reflection_memory = load_reflections_from_disk(agent_config["previous_reflections_path"])

    def build_execution_history(
        self, trajectory: TrajectoryView, meta_data: dict[str, Any], idxs_trajectory: list[int]
    ) -> list[Message]:
        # TODO: add other types of execution history prompt
        # Without utterances, parsed actions instead of utterances, etc.
        return self.build_interaction_history(trajectory, meta_data, idxs_trajectory)

    def construct(
        self,
        trajectory: TrajectoryView,
        objective: str,
        objective_imgs: list[ImageInput],
        meta_data: dict[str, Any],
        idxs_trajectory: list[int] = [],
        previous_reflections: list[str] = [],
    ) -> list[Message]:
        messages: list[Message | str | ImageInput] = []

        # -----------------------------------------------------------------------
        # System prompt, evaluation prompt
        # -----------------------------------------------------------------------
        sys_prompt_template: str = self.instruction[f"system_prompt_{self.mode}"]
        reflection_request: str = self.instruction[f"reflection_request_{self.mode}"]

        sys_prompt = partial_format(
            sys_prompt_template,
            image_info=self.img_obs_info_trace,
            text_obs_info=self.text_obs_info,
            trace_info=self.trace_info,
        )

        # -----------------------------------------------------------------------
        # Examples
        # -----------------------------------------------------------------------

        # -----------------------------------------------------------------------
        # Build intent input
        # -----------------------------------------------------------------------
        # text intro for objective
        text_input = self.instruction["objective_template"].format(objective=objective)

        # Add intent images and corresponding textual prefixes
        state_img_intros, add_state_idxs = [], []

        # If SGV, and is the first pass, initial page is necessary to define the intent
        if self.mode == ReflexionMode.SGV.value:
            raise NotImplementedError("TODO: Implement SGV")

        messages.extend(
            self.build_intent_message(
                trajectory,
                text_input,
                objective_imgs,
                meta_data,
                objective_image_captions=self.get_image_captions(objective_imgs, meta_data),
                state_img_intros=state_img_intros,
                add_states_idxs=add_state_idxs,
                add_state_text=False,
                add_state_img=True,
                role="user",
                name=self.lm_config.get("name_user", ""),
            )
        )

        # -----------------------------------------------------------------------
        # Execution history
        # -----------------------------------------------------------------------
        # For answers providing only URLs, give additional context for being possible to evaluate
        msg_img_links, img_metadata = self.get_data_for_urls(
            trajectory.actions[-1].get("parsed_action", ""),
            trajectory.states[-1].get("observation", {}).get("text", ""),
        )
        # If raw link answers, provide associated data from the text observation (if any)
        if img_metadata and not self.use_text_observation:
            add_text_metadata = {"intro": "STATE `t-{t}` TEXT REPRESENTATION:\n", "text": img_metadata}
            trajectory.states[-1]["add_text_metadata"] = (
                add_text_metadata  # This adds a text observation to the specific state
            )
        execution_history_msgs = self.build_interaction_history(
            trajectory,
            meta_data,
            idxs_trajectory,
        )
        messages.extend(execution_history_msgs)

        # If image link answers, provide the image or images
        if msg_img_links:
            messages.append(msg_img_links)  # add the image or images
            # get_message(inputs = execution_history_msgs[-1].contents + msg_img_links.contents, role = "assistant", name = "assistant")

        # -----------------------------------------------------------------------
        # Get Reflexion Request
        # -----------------------------------------------------------------------
        # Inject previous reflections, if any
        reflection_injection_prompt = self.instruction["reflections_injection_template"]
        if previous_reflections:
            previous_reflections_str = ""
            for i, reflection in enumerate(previous_reflections):
                previous_reflections_str += f"\n- {i + 1}: {reflection}"
        else:
            previous_reflections_str = " None"

        messages.append(
            get_message(
                inputs=reflection_injection_prompt.format(previous_reflections=previous_reflections_str),
                role="user",
                name=self.lm_config.get("name_user", ""),
            )
        )

        # Parsing error feedback
        parsing_error_msg = ""
        if meta_data.get(self.parsing_error_msg_key):
            parsing_error_msg = meta_data[self.parsing_error_msg_key]
            meta_data[self.parsing_error_msg_key] = ""

        messages.append(
            get_message(
                inputs=reflection_request + parsing_error_msg,
                role="user",
                name=self.lm_config.get("name_user", ""),
            )
        )

        return get_messages(
            inputs=messages,  # type: ignore
            sys_prompt=sys_prompt,
            role="user",
            name=self.lm_config.get("name_user", ""),
        )

    def parse_reflection(self, response: str) -> dict[str, Any]:
        splitters: list[str] = self.instruction["meta_data"]["splitters"]
        parsed_data: dict[str, Any] = {}
        splitters_group: str = "|".join(map(re.escape, splitters))

        # Iterate over splitters and parse corresponding sections
        for splitter in splitters:
            pattern = rf"{re.escape(splitter)}(.*?)(?=\n(?:{splitters_group})|$)"
            matches = re.findall(pattern, response, re.IGNORECASE | re.DOTALL)
            if matches:
                content = matches[-1].strip()
                parsed_data[re.sub(r":$", "", splitter)] = content  # Remove trailing ":"

        if not all(req_splitter in parsed_data for req_splitter in self.instruction["meta_data"]["required_splitters"]):
            raise LMParsingError(f"Cannot find all required splitters in {response}")

        return parsed_data


class Reflexion(Agent):
    def __init__(self, lm_config: dict[str, Any], agent_config: dict[str, Any]) -> None:
        super().__init__()
        self.lm_config = lm_config
        self.prompt_constructor: ReflexionPromptConstructor = ReflexionPromptConstructor(
            lm_config=lm_config, agent_config=agent_config
        )
        self.eval_mode: EvalMode = EvalMode(agent_config["eval_mode"])
        self.max_reflexion_attempts = agent_config["max_reflexion_attempts"]
        self.agent_config = agent_config
        self.num_previous_state_actions = agent_config["num_previous_state_actions"]
        self.mode = agent_config["mode"]
        self.out_utterance = agent_config["out_utterance"]
        self.max_model_call = agent_config["max_model_call"]
        self.conversation_dir = agent_config.get("conversation_dir", None)
        self.usage_dir = agent_config.get("usage_dir", None)
        self.num_attempts_per_task = {}
        self.reflection_memory: dict[str, list[str]] = {}
        self.previous_reflections: dict[str, list[str]] = {}
        self.num_previous_reflections = agent_config["num_previous_reflections"]

        if agent_config.get("previous_reflections_path"):
            self.previous_reflections = load_reflections_from_disk(agent_config["previous_reflections_path"])

    def get_reflection_from_parsed_response(self, parsed_response: dict[str, Any]) -> str:
        reflection_str = parsed_response[self.prompt_constructor.instruction["meta_data"]["reflection_key"]]  # type: ignore
        # strip left ":" and space
        reflection_str = reflection_str.lstrip(":").strip()
        reflection_str = clean_spaces(reflection_str)

        # Remove "#" markers
        reflection_str = re.sub(r"#+", "", reflection_str)
        reflection_str = re.sub(r"(?m)^\s+", "", reflection_str)
        return reflection_str

    def should_retry(self, domain_task_id: str) -> bool:
        return self.get_num_attempts_per_task(domain_task_id) < self.max_reflexion_attempts

    def get_eval_mode(self) -> str:
        return self.eval_mode.value

    def select_states(self, trajectory: TrajectoryView, num_states: int = -1) -> list[int]:
        # Return indices of the last `num_previous_state_actions` states
        num_states = min(self.num_previous_state_actions, len(trajectory.states))
        return list(range(len(trajectory.states) - num_states, len(trajectory.states)))

    def select_reflections(self, domain_task_id: str, num_previous_reflections: int = -1) -> list[str]:
        prev_reflections_list = self.reflection_memory.get(domain_task_id, [])
        if not prev_reflections_list:
            return []
        if num_previous_reflections == -1:
            return prev_reflections_list
        return prev_reflections_list[:num_previous_reflections]

    def next_action(
        self,
        trajectory: TrajectoryView,
        intent: str,
        intent_images: list[ImageInput],
        meta_data: dict[str, Any],
    ) -> tuple[Any, str]:
        idxs_trajectory = self.select_states(trajectory)
        num_tries = 0
        logger.info("Reflexion Agent: generating reflection...")
        prev_reflections = self.select_reflections(
            meta_data["domain_task_id"], num_previous_reflections=self.num_previous_reflections
        )
        parsed_response, raw_response, _ = self.act_parse_retry(
            trajectory=trajectory,
            intent=intent,
            intent_images=intent_images,
            meta_data=meta_data,
            idxs_trajectory=idxs_trajectory,
            parser_fn=self.prompt_constructor.parse_reflection,
            max_tries=self.max_model_call - num_tries,
            error_msg_template=self.prompt_constructor.instruction["parsing_error_msg_template"],
            error_msg_key=self.prompt_constructor.instruction["meta_data"]["parsing_error_msg_key"],
            previous_reflections=prev_reflections,
        )

        if self.out_utterance:
            logger.info(f"\n[Reflexion Agent]: {raw_response}")

        reflection_str = self.get_reflection_from_parsed_response(parsed_response)
        self.update_reflection_memory(meta_data["domain_task_id"], [reflection_str])
        self.save_reflection_to_disk(meta_data["domain_task_id"], reflection_str, raw_response)

        return parsed_response, raw_response

    def get_num_attempts_per_task(self, domain_task_id: str) -> int:
        return len(self.reflection_memory.get(domain_task_id, []))

    def save_reflection_to_disk(self, domain_task_id: str, reflection: str, raw_response: str) -> None:
        # Open existing json
        if os.path.exists(self.reflection_memory_path):
            with open(self.reflection_memory_path, "r") as f:
                reflection_data = json.load(f)
        else:
            reflection_data = {}

        if domain_task_id not in reflection_data:
            reflection_data[domain_task_id] = []
        reflection_data[domain_task_id].append({"reflection": reflection, "raw_response": raw_response})

        dest_file_temp = Path(self.reflection_memory_path).with_suffix(".tmp")
        with open(dest_file_temp, "w") as f:
            json.dump(reflection_data, f)
        dest_file_temp.rename(self.reflection_memory_path)

    def update_reflection_memory(self, domain_task_id: str, reflections: list[str]) -> None:
        if domain_task_id not in self.reflection_memory:
            self.reflection_memory[domain_task_id] = []
        self.reflection_memory[domain_task_id].extend(reflections)

    def populate_reflection_memory(self, domain_task_id: str, args: dict[str, Any]) -> None:
        self.reflection_memory = {}
        self.reflection_memory_path = f"{args['result_dir']}/reflections.json"

        # If existing reflections from current run, load them
        self.reflection_memory = load_reflections_from_disk(self.reflection_memory_path, domain_task_id)

        # Combine current reflection memory with reflections from previous runs, if any
        if self.previous_reflections:
            for domain_task_id, prev_reflections in self.previous_reflections.items():
                self.update_reflection_memory(domain_task_id, prev_reflections)

    def reset(self, config_dict: dict[str, Any], args: dict[str, Any] = {}) -> None:
        super().reset(config_dict, args)
        task_id = config_dict["task_id"]
        attempt_num = args["attempts_per_task"][task_id]
        task_id_attempt = f"{task_id}_{attempt_num}"
        self.conversation_dir = f"{args['lm_logs_dir']}/{task_id_attempt}/conversations"
        self.usage_dir = f"{args['lm_logs_dir']}/{task_id_attempt}/lm_usage"
        self.call_id = f"reflexion"
