from posix import stat
from langchain_community.chat_models import human
from langchain_core.runnables import history
import requests
import time
import dspy

from colmind.agents import memory
from langchain.memory import ConversationBufferMemory
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
from langchain.chains import ConversationChain
from typing import Any
from dataclasses import dataclass
from .agents.llm import get_llm

from colmind.prompts import load_prompt

CONVERSATION_TEMPLATE_WEAK = """
You are a Minecraft agent named Weak and you are having a conversation with a more knowleadgeable agent named Strong.
Your are trying to complete the following task: {task}.

Below you have a mental model of the other agent, Strong, containing information about its state.
Mental Model of Strong: {other_world_model}

Based on the mental model of the Strong agent above, you have taken perspective of Strong in order to understand its motifs and reasoning.
Current perspective on Strong's state: {perspective}

Lastly, you have a mental model of yourself, Weak, containing information about your state.
World Model of Weak: {world_model}

Based on the current conversation, your perspective on Strong, and your own mental model, you have to respond to the conversation such that you get the most help from Strong in solving your task. Provide Strong with information about yourself and your problems.
Keep your responses concise and to the point. Do not provide unnecessary information. Share with Strong possible problems and why you think you failed the task.

Here are a few example conversations to help you understand how to respond:

Current conversation:
Weak: Hey, can you help me with crafting a bucket?
Strong: Sure, you need to gather iron ingots and smelt. Did you consider trying something already?
Weak: I have tried to gather iron ingots, but I am not sure how to smelt them. I had an execution error.
Strong: What was you execution error and what do you have in your inventory?
Weak: I have iron ore in my inventory and I tried to smelt it in the furnace, but it did not work. Here is my code ...

Current conversation:
Weak: Hey, can you help me with crafting a bucket?
Strong: Of course, what did you try until now? Do you have any problems?
Weak: I have tried to gather iron ingots, but I am not sure how to smelt them. I had an execution error.
Strong: What was you execution error and what do you have in your inventory?
Weak: I have iron ore in my inventory and I tried to smelt it in the furnace, but it did not work. Could you share what code you would be using?

Current conversation:
Weak: Hey, can you help me with building a house?
Strong: Yes, what materials do you have in your inventory?
Weak: Based on my percepts, I have wood and stone. How can I proceed?
Strong: Which biome are you in?
Weak: I am in the forest biome.


Current conversation:\n{history}
Strong:{input}
Weak:
"""

CONVERSATION_TEMPLATE_STRONG = """
You are a Minecraft agent named Strong and you are having a conversation with a less knowleadgeable agent named Weak that needs your help.
Your are trying to help Weak solve his task.

Below you have a mental model of the other agent, Weak, containing information about his state.
Mental Model of Weak: {other_world_model}

Based on the mental model of the Weak agent above, you have taken perspective of Strong in order to understand its motifs and reasoning.
Current perspective on Weak's state: {perspective}

Lastly, you have a mental model of yourself, Strong, containing information about your state.
World Model of StrStrong {world_model}

Based on the current conversation, your perspective on Weak, and your own mental model, you have to respond to the conversation such that you help Weak complete his task as good as possible.
Keep your responses concise and to the point. Do not provide unnecessary information. Try to give actionable advice, with specific instructions and simple explanations. Before giving any concrete advice, ask questions about the mental model of the agent.

Make sure that when providing advice, you include both code and natural language explanations.

Here are a few example conversations to help you understand how to respond:

Current conversation:
Weak: Hey, can you help me with crafting a bucket?
Strong: Sure, you need to gather iron ingots and smelt. Did you consider trying something already?
Weak: I have tried to gather iron ingots, but I am not sure how to smelt them. I had an execution error.
Strong: What was you execution error and what do you have in your inventory?
Weak: I have iron ore in my inventory and I tried to smelt it in the furnace, but it did not work.
Strong: You should use the function smelt_iron_ore() to smelt the iron ore in the furnace. Here is the code I am using: smelt_iron_ore(iron_ore).

Current conversation:
Weak: Hey, can you help me with crafting a bucket?
Strong: Of course, what did you try until now? Do you have any problems?
Weak: I have tried to gather iron ingots, but I am not sure how to smelt them. I had an execution error.
Strong: What was you execution error and what do you have in your inventory?
Weak: I have iron ore in my inventory and I tried to smelt it in the furnace, but it did not work.
Strong: You should use the function smelt_iron_ore() to smelt the iron ore in the furnace. Here is the code I am using: smelt_iron_ore(iron_ore).

Current conversation:
Weak: Hey, can you help me with building a house?
Strong: Yes, what materials do you have in your inventory?
Weak: Based on my percepts, I have wood and stone. How can I proceed? I want to use a diamond shovel to build.
Strong: Which biome are you in? You do not need a diamond shovel to build.
Weak: I am in the forest biome.
Strong: Here is the function i was using to build ... 

Current conversation:\n{history}
Weak:{input}
Strong:
"""

@dataclass
class AgentConfig:
    name: str
    llm: str
    temperature: float
    internal: Any
    partners: dict
    template: str

class UpdateWorldModel(dspy.Signature):
    """Update the world model of an agent based on conversation, perspective, and previous world model."""

    your_name = dspy.InputField(desc="Your name")
    other_name = dspy.InputField(desc="Other agent's name")
    current_conversation = dspy.InputField(desc="Previous conversation history")
    perspective = dspy.InputField(desc="Perspective taken on the other agent")
    previous_code_from_last_round = dspy.InputField(desc="Previous belief about other's agent code from the last round")
    previous_execution_error = dspy.InputField(desc="Previous belief about other's agent execution errors")
    previous_biome = dspy.InputField(desc="Previous belief about other's agent biome")
    previous_time = dspy.InputField(desc="Previous belief about other's agent  time")
    previous_nearby_blocks = dspy.InputField(desc="Previous belief about other's agent  nearby blocks")
    previous_health = dspy.InputField(desc="Previous belief about other's agent health")
    previous_hunger = dspy.InputField(desc="Previous belief about other's agent hunger")
    previous_position = dspy.InputField(desc="Previous belief about other's agent position")
    previous_equipment = dspy.InputField(desc="Previous belief about other's agent equipment")
    previous_inventory = dspy.InputField(desc="Previous belief about other's agent inventory")
    previous_chests = dspy.InputField(desc="Previous belief about other's agent chests")
    previous_critique = dspy.InputField(desc="Previous belief about other's agent critique")
    previous_task = dspy.InputField(desc="Previous belief about other's agent task")
    previous_task_beliefs = dspy.InputField(desc="Previous belief about other's agent task beliefs")
    previous_interaction_beliefs = dspy.InputField(desc="Previous belief about other's agent interaction beliefs")
    previous_partner_beliefs = dspy.InputField(desc="Previous belief about other's agent partner beliefs")
    previous_perception_beliefs = dspy.InputField(desc="Previous belief about other's agent perception beliefs")

    code_from_last_round = dspy.OutputField(desc="Updated belief about other's agent code from the last round")
    execution_error = dspy.OutputField(desc="Updated belief about other's agent execution errors")
    biome = dspy.OutputField(desc="Updated belief about other's agent biome")
    time = dspy.OutputField(desc="Updated belief about other's agent time")
    nearby_blocks = dspy.OutputField(desc="Updated belief about other's agent nearby blocks")
    health = dspy.OutputField(desc="Updated belief about other's agent health")
    hunger = dspy.OutputField(desc="Updated belief about other's agent hunger")
    position = dspy.OutputField(desc="Updated belief about other's agent position")
    equipment = dspy.OutputField(desc="Updated belief about other's agent equipment")
    inventory = dspy.OutputField(desc="Updated belief about other's agent inventory")
    chests = dspy.OutputField(desc="Updated belief about other's agent chests")
    critique = dspy.OutputField(desc="Updated belief about other's agent critique")
    task = dspy.OutputField(desc="Updated belief about other's agent task")
    task_beliefs = dspy.OutputField(desc="Updated belief about other's agent task beliefs")
    interaction_beliefs = dspy.OutputField(desc="Updated belief about other's agent interaction beliefs")
    partner_beliefs = dspy.OutputField(desc="Updated belief about other's agent partner beliefs")
    perception_beliefs = dspy.OutputField(desc="Updated belief about other's agent perception beliefs")


TEACHER_MODEL_DSPY = "openai/gpt-4"
STUDENT_MODEL_DSPY = "together_ai/meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo"

class UpdateWorldModelModule(dspy.Module):
    def __init__(self, character: str):
        super().__init__()
        if character == "weak":
            llm = STUDENT_MODEL_DSPY
        else:
            llm = TEACHER_MODEL_DSPY
        self.llm = dspy.LM(model=llm)
        dspy.configure(lm=self.llm)


    def forward(self, **kwargs):
        return dspy.Predict(UpdateWorldModel)(**kwargs)

@dataclass
class WorldModel:
    code_from_last_round: str = "Unknown"
    execution_error: str = "Unknown"
    biome: str = "Unknown"
    time: str = "Unknown"
    nearby_blocks: str = "Unknown"
    health: str = "Unknown"
    hunger: str = "Unknown"
    position: str = "Unknown"
    equipment: str = "Unknown"
    inventory: str = "Unknown"
    chests: str = "Unknown"
    critique: str = "Unknown"
    task: str = "Unknown"
    task_beliefs: str = "Unknown"
    interaction_beliefs: str = "Unknown"
    partner_beliefs: str = "Unknown"
    perception_beliefs: str = "Unknown"


class CommunicationChannel:
    def __init__(self, config: dict):
        self.config = config
        self.num_acks = 0
        self.weak_agent_world_model = WorldModel() # from the perspective of the strong agent
        self.strong_agent_world_model = WorldModel() # from the perspective of the weak agent

    def run(self):
        active = False
        round = None
        task = None
        while True:
            while not active: # wait for all the agents to be ready
                time.sleep(1)
                res = requests.get("http://localhost:3002/conversation")
                active = res.json()["active"]
                task = res.json()["task"]
                if round is None and active: # new round, there is nothing there yet
                    round = CommunicationRound(self.config, task, self.weak_agent_world_model, self.strong_agent_world_model)
                elif round is not None and task != round.task and active: # if I do a new task, I should reset the round
                    round = CommunicationRound(self.config, task, self.weak_agent_world_model, self.strong_agent_world_model)

            assert round is not None
            round.run()
            res = requests.post("http://localhost:3002/finish-conversation", json={"finished": True})
            if res.status_code == 200:
                print(f"Fininshed conversation for round {round.rounds}")
            active = False


class CommunicationRound:
    def __init__(self, config: dict, task: str, weak_agent_world_model: WorldModel, strong_agent_world_model: WorldModel):
        self.config = config
        self.task = task
        self.weak_agent_world_model = weak_agent_world_model
        self.strong_agent_world_model = strong_agent_world_model
        self.weak_memory = ConversationBufferMemory(
            input_key="input",
            memory_key="history",
            return_messages=False,
            ai_prefix="Weak",
            human_prefix="Strong",
        )
        self.strong_memory = ConversationBufferMemory(
            input_key="input",
            memory_key="history",
            return_messages=False,
            ai_prefix="Strong",
            human_prefix="Weak",
        )

        self.weak_template = PromptTemplate(
            template = "You are a Minecraft agent named Weak and you are having a conversation with a more knowleadgeable agent named Strong.\n\nCurrent conversation: {history}\n\n{input}\n",
            input_variables = ["history", "input"]
        )

        self.strong_template = PromptTemplate(
            template = "You are a Minecraft agent named Strong and you are having a conversation with a more knowleadgeable agent named Weak.\n\nCurrent conversation: {history}\n\n{input}\n",
            input_variables = ["history", "input"]
        )

        self.weak_perspective_template = PromptTemplate(
            template=(
                "You are a Minecraft agent named Weak and you are having a conversation with a more knowleadgeable agent named Strong that wants to offer you help."
                "Based on the current conversation and your knowledge about the other agent, Strong, take the other agent's perspective to assess and describe your current understanding, knowledge state, and likely needs from Strong's perspective."
                "Here is the current conversation between you and Strong:\n\n{conversation}\n\n"
                "Here is the mental model of Strong:\n\n{world_model}\n\n"
                "Perspective Analysis:"
                ),
                input_variables=["conversation", "world_model"]
        )

        self.strong_perspective_template = PromptTemplate(
            template=(
                "You are a Minecraft agent named Strong and you are having a conversation with a less knowleadgeable agent named Weak that needs your help."
                "Based on the current conversation and your knowledge about the other agent, Weak, take the other agent's perspective to assess and describe your current understanding, knowledge state, and likely needs from Weak's perspective."
                "Here is the current conversation between you and Weak:\n\n{conversation}\n\n"
                "Here is the mental model of Weak:\n\n{world_model}\n\n"
                "Perspective Analysis:"
                ),
                input_variables=["conversation", "world_model"]
        )



        self.agents = {}
        weak_agent_config, internal_world_model_weak = self.create_conversational_agent(agent_name="weak")
        self.internal_world_model_weak = internal_world_model_weak
        self.weak_perspective_llm = get_llm(model_name=weak_agent_config.llm, temperature=weak_agent_config.temperature, max_tokens=256)
        self.weak_world_model_updater = UpdateWorldModelModule(character="weak")
        self.agents["weak"] = ConversationChain(
            llm=get_llm(model_name=weak_agent_config.llm, temperature=weak_agent_config.temperature),
            memory=self.weak_memory,
            prompt=self.weak_template,
        )
        strong_agent_config, internal_world_model_strong = self.create_conversational_agent(agent_name="strong")
        self.internal_world_model_strong = internal_world_model_strong
        self.strong_perspective_llm = get_llm(model_name=strong_agent_config.llm, temperature=strong_agent_config.temperature, max_tokens=256)
        self.strong_world_model_updater = UpdateWorldModelModule(character="strong")
        self.agents["strong"] = ConversationChain(
            llm=get_llm(model_name=strong_agent_config.llm, temperature=strong_agent_config.temperature),
            memory=self.strong_memory,
            prompt=self.strong_template
        )
        self.rounds = 0 # the max number of rounds should be equivalent with the max number of tries to solve a task

    def run(self):
        self.rounds += 1
        sender, receiver = "strong", "weak"

        # the first message that is sent to request foe help is hardcoded
        if self.rounds == 1:
            input_conversation = f"Hey, can you help me with solving the task: {self.task}?"
            #self.weak_memory.save_context({"input": input_conversation}, {"output": ""}) #add to memory
        else:
            input_conversation = f"I still have problems solving the task {self.task}. Could you give me additional help?"
            #self.weak_memory.save_context({"input": input_conversation}, {"output": ""})  #add to memory

        current_conversation = self.format_conversation()
        perspective = self.get_perspective(current_conversation, sender)
        for idx in range(self.config["num_conversation_turns"]):
            print(receiver, input_conversation, f"{idx}")
            current_conversation = self.format_conversation()
            if idx > 0:
                if receiver == "weak":
                    updated_world_model = self.update_world_model(
                        current_conversation=current_conversation,
                        perspective=perspective,
                        agent_name=receiver,
                        world_model=self.weak_agent_world_model # Pass the weak agent's world model
                    )
                    self.weak_agent_world_model = updated_world_model
                elif receiver == "strong":
                    updated_world_model = self.update_world_model(
                        current_conversation=current_conversation,
                        perspective=perspective,
                        agent_name=receiver,
                        world_model=self.strong_agent_world_model
                    )
                    self.strong_agent_world_model = updated_world_model

            perspective = self.get_perspective(current_conversation, sender)
            print(f"\n{sender.capitalize()} taking {receiver}'s perspective:", perspective, "\n")

            if sender == "weak":
                original_template = CONVERSATION_TEMPLATE_WEAK
            else:
                original_template = CONVERSATION_TEMPLATE_STRONG
            temp_prompt = PromptTemplate(
                            template=original_template,
                            input_variables=["history", "input", "perspective", "world_model", "task", "other_world_model"]
            )

            self.agents[sender].prompt = temp_prompt
            if sender == "weak":
                input_conversation = self.agents[sender].run(
                            input=input_conversation,
                            perspective=perspective,
                            task=self.task,
                            other_world_model=self.strong_agent_world_model,
                            world_model=self.internal_world_model_weak
                )
            else:
                input_conversation = self.agents[sender].run(
                            input=input_conversation,
                            perspective=perspective,
                            task=self.task,
                            world_model=self.weak_agent_world_model,
                            other_world_model=self.internal_world_model_strong
                )
            self.agents[sender].prompt.template = original_template
            # input_conversation = self.agents[sender].run(input=input_conversation)
            sender, receiver = receiver, sender

        conversation = self.format_conversation()
        print("\n\n\n", conversation, "\n\n\n")
        res = requests.post("http://localhost:3002/latest-conversation", json={"conversation": conversation})
        if res.status_code == 200:
            print(f"Sent conversation for round {self.rounds}")


    def get_perspective(self, conversation: str, agent: str) -> str:
        """Take perspective of the other agent based on conversation history."""
        if agent == "weak":
            return self.weak_perspective_llm.predict(
                self.weak_perspective_template.format(
                    task=self.task,
                    conversation=conversation,
                    world_model=self.strong_agent_world_model
                )
            )
        else:
            return self.strong_perspective_llm.predict(
                self.strong_perspective_template.format(
                    conversation=conversation,
                    world_model=self.weak_agent_world_model
                )
            )


    def update_world_model(self, current_conversation, perspective, agent_name, world_model: WorldModel) -> WorldModel:
        """Updates the world model using the DSPy module."""
        if agent_name == "weak":
            updated_model = self.weak_world_model_updater(
                your_name="Weak",
                other_name="Strong",
                current_conversation=current_conversation,
                perspective=perspective,
                previous_code_from_last_round=world_model.code_from_last_round,
                previous_execution_error=world_model.execution_error,
                previous_biome=world_model.biome,
                previous_time=world_model.time,
                previous_nearby_blocks=world_model.nearby_blocks,
                previous_health=world_model.health,
                previous_hunger=world_model.hunger,
                previous_position=world_model.position,
                previous_equipment=world_model.equipment,
                previous_inventory=world_model.inventory,
                previous_chests=world_model.chests,
                previous_critique=world_model.critique,
                previous_task=world_model.task,
                previous_task_beliefs=world_model.task_beliefs,
                previous_interaction_beliefs=world_model.interaction_beliefs,
                previous_partner_beliefs=world_model.partner_beliefs,
                previous_perception_beliefs=world_model.perception_beliefs,
            )
        elif agent_name == "strong":
            updated_model = self.strong_world_model_updater(
                your_name="Strong",
                other_name="Weak",
                current_conversation=current_conversation,
                perspective=perspective,
                previous_code_from_last_round=world_model.code_from_last_round,
                previous_execution_error=world_model.execution_error,
                previous_biome=world_model.biome,
                previous_time=world_model.time,
                previous_nearby_blocks=world_model.nearby_blocks,
                previous_health=world_model.health,
                previous_hunger=world_model.hunger,
                previous_position=world_model.position,
                previous_equipment=world_model.equipment,
                previous_inventory=world_model.inventory,
                previous_chests=world_model.chests,
                previous_critique=world_model.critique,
                previous_task=world_model.task,
                previous_task_beliefs=world_model.task_beliefs,
                previous_interaction_beliefs=world_model.interaction_beliefs,
                previous_partner_beliefs=world_model.partner_beliefs,
                previous_perception_beliefs=world_model.perception_beliefs,
            )
        else:
            raise ValueError(f"Invalid agent name {agent_name}")

        return WorldModel(
            code_from_last_round=updated_model.code_from_last_round,
            execution_error=updated_model.execution_error,
            biome=updated_model.biome,
            time=updated_model.time,
            nearby_blocks=updated_model.nearby_blocks,
            health=updated_model.health,
            hunger=updated_model.hunger,
            position=updated_model.position,
            equipment=updated_model.equipment,
            inventory=updated_model.inventory,
            chests=updated_model.chests,
            critique=updated_model.critique,
            task=updated_model.task,
            task_beliefs=updated_model.task_beliefs,
            interaction_beliefs=updated_model.interaction_beliefs,
            partner_beliefs=updated_model.partner_beliefs,
            perception_beliefs=updated_model.perception_beliefs,
        )
    def create_conversational_agent(self, agent_name: str):
        res = requests.get("http://localhost:3002/agent-info", json={"agent": agent_name})
        print(res.json())
        return AgentConfig(
            name=res.json()["name"],
            llm=res.json()["llm"],
            temperature=res.json()["temperature"],
            internal=res.json()["internal"],
            partners=res.json()["partners"],
            template=res.json()["template"]
        ), WorldModel(
           code_from_last_round=res.json()["code_from_last_round"],
            execution_error=res.json()["execution_error"],
            biome=res.json()["biome"],
            time=res.json()["time_of_day"],
            nearby_blocks=res.json()["voxels"],
            health=res.json()["health"],
            hunger=res.json()["food"],
            position=res.json()["position"],
            equipment=res.json()["equipment"],
            inventory=res.json()["inventory"],
            chests=res.json()["chest_observation"],
            critique=res.json()["critique"],
            task=res.json()["task"],
            task_beliefs=res.json()["task_beliefs"],
            interaction_beliefs=res.json()["interaction_beliefs"],
            partner_beliefs=res.json()["partner_beliefs"],
            perception_beliefs=res.json()["perception_beliefs"]
        )

    def format_conversation(self):
        formatted_conversation = [f"Weak: Hey, can you help me with {self.task}?" if self.rounds == 1 else f"Weak: I still have problems solving the task {self.task}. Could you give me additional help?"]
        for message in self.agents['weak'].memory.chat_memory.messages:
            if message.type == 'human':
                formatted_conversation.append(f"Strong: {message.content}")
            else:
                formatted_conversation.append(f"Weak: {message.content}")
        return "\n".join(formatted_conversation)
