from colmind.agents.llm import customLLM, get_llm
from colmind.prompts import load_prompt
from colmind.utils.json_utils import fix_and_parse_json

from langchain.schema import HumanMessage, SystemMessage


class BeliefCriticAgent:
    def __init__(
        self,
        config,
        username,
        logger,
        temperature=0,
    ):
        self.config = config
        self.username = username
        self.logger = logger

        self.llm = get_llm(config["parameters"]["llm"])
        self.mode = config["parameters"]["mode"]
        assert self.mode in ["auto", "manual"]

    def render_system_message(self):
        system_message = SystemMessage(content=load_prompt("critic"))
        return system_message

    def render_system_message_conversation(self):
        system_message = SystemMessage(content=load_prompt("critic_conversation"))
        return system_message

    def is_item_in_inventory(self, item_name, events):
        inventory = events[-1][1]["inventory"]
        print("inventory", inventory)

        if item_name in inventory:
            return True
        else:
            return False



    def render_human_message(self, *, events, task, context, partner_beliefs, chest_observation):
        assert events[-1][0] == "observe", "Last event must be observe"
        biome = events[-1][1]["status"]["biome"]
        time_of_day = events[-1][1]["status"]["timeOfDay"]
        voxels = events[-1][1]["voxels"]
        health = events[-1][1]["status"]["health"]
        hunger = events[-1][1]["status"]["food"]
        position = events[-1][1]["status"]["position"]
        equipment = events[-1][1]["status"]["equipment"]
        inventory_used = events[-1][1]["status"]["inventoryUsed"]
        inventory = events[-1][1]["inventory"]

        for i, (event_type, event) in enumerate(events):
            if event_type == "onError":
                self.logger.warning(f"Critic Agent: Error occurs {event['onError']}")
                return None

        observation = ""

        observation += f"Biome: {biome}\n\n"

        observation += f"Time: {time_of_day}\n\n"

        if voxels:
            observation += f"Nearby blocks: {', '.join(voxels)}\n\n"
        else:
            observation += f"Nearby blocks: None\n\n"

        observation += f"Health: {health:.1f}/20\n\n"
        observation += f"Hunger: {hunger:.1f}/20\n\n"

        observation += f"Position: x={position['x']:.1f}, y={position['y']:.1f}, z={position['z']:.1f}\n\n"

        observation += f"Equipment: {equipment}\n\n"

        if inventory:
            observation += f"Inventory ({inventory_used}/36): {inventory}\n\n"
        else:
            observation += f"Inventory ({inventory_used}/36): Empty\n\n"

        observation += chest_observation

        observation += f"Task: {task}\n\n"

        if context:
            observation += f"Context: {context}\n\n"
        else:
            observation += f"Context: None\n\n"

        #observation += f"Interaction Beliefs: {partner_beliefs}\n\n"

        self.logger.info(f"****Critic Agent human message****\n{observation}")
        return HumanMessage(content=observation)

    def render_human_message_conversation(self, *, events, task, context, partner_beliefs, chest_observation, old_critique):
        assert events[-1][0] == "observe", "Last event must be observe"
        biome = events[-1][1]["status"]["biome"]
        time_of_day = events[-1][1]["status"]["timeOfDay"]
        voxels = events[-1][1]["voxels"]
        health = events[-1][1]["status"]["health"]
        hunger = events[-1][1]["status"]["food"]
        position = events[-1][1]["status"]["position"]
        equipment = events[-1][1]["status"]["equipment"]
        inventory_used = events[-1][1]["status"]["inventoryUsed"]
        inventory = events[-1][1]["inventory"]

        for i, (event_type, event) in enumerate(events):
            if event_type == "onError":
                self.logger.warning(f"Critic Agent: Error occurs {event['onError']}")
                return None

        observation = ""

        observation += f"Biome: {biome}\n\n"

        observation += f"Time: {time_of_day}\n\n"

        if voxels:
            observation += f"Nearby blocks: {', '.join(voxels)}\n\n"
        else:
            observation += f"Nearby blocks: None\n\n"

        observation += f"Health: {health:.1f}/20\n\n"
        observation += f"Hunger: {hunger:.1f}/20\n\n"

        observation += f"Position: x={position['x']:.1f}, y={position['y']:.1f}, z={position['z']:.1f}\n\n"

        observation += f"Equipment: {equipment}\n\n"

        if inventory:
            observation += f"Inventory ({inventory_used}/36): {inventory}\n\n"
        else:
            observation += f"Inventory ({inventory_used}/36): Empty\n\n"

        observation += chest_observation

        observation += f"Task: {task}\n\n"

        if context:
            observation += f"Context: {context}\n\n"
        else:
            observation += f"Context: None\n\n"

        observation += f"Interaction Beliefs: {partner_beliefs}\n\n"
        observation += f"Old Critique: {old_critique}\n\n"

        self.logger.info(f"****Critic Agent human message****\n{observation}")
        return HumanMessage(content=observation)

    def human_check_task_success(self):
        confirmed = False
        success = False
        critique = ""
        while not confirmed:
            success = input("Success? (y/n)")
            success = success.lower() == "y"
            critique = input("Enter your critique:")
            self.logger.info(f"Success: {success}\nCritique: {critique}")
            confirmed = input("Confirm? (y/n)") in ["y", ""]
        return success, critique

    def ai_check_task_success(self, messages, max_retries=5):
        if max_retries == 0:
            print(
                "\033[31mFailed to parse Critic Agent response. Consider updating your prompt.\033[0m"
            )
            return False, ""

        if messages[1] is None:
            return False, ""

        critic = self.llm(messages).content
        self.logger.info(f"****Critic Agent ai message****\n{critic}")
        try:
            response = fix_and_parse_json(critic)
            assert response["success"] in [True, False]
            if "critique" not in response:
                response["critique"] = ""
            return response["success"], response["critique"]
        except Exception as e:
            print(f"\033[31mError parsing critic response: {e} Trying again!\033[0m")
            return self.ai_check_task_success(
                messages=messages,
                max_retries=max_retries - 1,
            )

    def ai_check_task_success_conversation(self, messages, max_retries=5):
        if max_retries == 0:
            print(
                "\033[31mFailed to parse Critic Agent response. Consider updating your prompt.\033[0m"
            )
            return False, ""

        if messages[1] is None:
            return False, ""

        critic = self.llm(messages).content
        self.logger.info(f"****Critic Agent ai message****\n{critic}")
        try:
            response = fix_and_parse_json(critic)
            return response["new_critique"]
        except Exception as e:
            print(f"\033[31mError parsing critic response: {e} for {critic} for {self.bot_username} Trying again!\033[0m")
            return self.ai_check_task_success_conversation(
                messages=messages,
                max_retries=max_retries - 1,
            )

    def check_task_success(
        self, *, events, task, context, partner_beliefs, chest_observation, max_retries=5
    ):
        human_message = self.render_human_message(
            events=events,
            task=task,
            context=context,
            partner_beliefs=partner_beliefs,
            chest_observation=chest_observation,
        )

        messages = [
            self.render_system_message(),
            human_message,
        ]
        if human_message is not None:
            self.logger.info(f"****Critic Agent Human Check Success****\n{human_message.content}")
        else:
            self.logger.info(f"****Critic Agent Human Check Success****\nNone")

        if self.mode == "manual":
            return self.human_check_task_success()
        elif self.mode == "auto":
            return self.ai_check_task_success(
                messages=messages, max_retries=max_retries
            )
        else:
            raise ValueError(f"Invalid critic agent mode: {self.mode}")

    def check_task_success_conversation(
        self, *, events, task, context, partner_beliefs, chest_observation, old_critique, max_retries=5
    ):
        human_message = self.render_human_message_conversation(
            events=events,
            task=task,
            context=context,
            partner_beliefs=partner_beliefs,
            old_critique=old_critique,
            chest_observation=chest_observation,
        )

        messages = [
            self.render_system_message_conversation(),
            human_message,
        ]
        if human_message is not None:
            self.logger.info(f"****Critic Agent Conversation Change***\n{human_message.content}")
        else:
            self.logger.info(f"****Critic Agent Conversation Change****\nNone")

        if self.mode == "manual":
            return self.human_check_task_success()
        elif self.mode == "auto":
            return self.ai_check_task_success_conversation(
                messages=messages, max_retries=max_retries
            )
        else:
            raise ValueError(f"Invalid critic agent mode: {self.mode}")
