from colmind.agents.agent import LLMAgent
import dspy
from colmind.prompts import load_prompt
from colmind.utils.json_utils import fix_and_parse_json
from pathlib import Path
from dspy.teleprompt import LabeledFewShot
from dspy.primitives import Example

CRITIC_INSTRUCTIONS = """
You are an assistant that assesses my progress of playing Minecraft and provides useful guidance.

You are required to evaluate if I have met the task requirements. Exceeding the task requirements is also considered a success while failing to meet them requires you to provide critique to help me improve.
"""

class Critic(dspy.Signature):
    """Generate action critiques."""
    biome: str = dspy.InputField(desc="The biome after the task execution.")
    time_of_day: str = dspy.InputField(desc="The current time.")
    nearby_blocks: list = dspy.InputField()
    health: float = dspy.InputField(desc="Current health.")
    hunger: float = dspy.InputField(desc="Current hunger level.")
    position: dict = dspy.InputField(desc="Current position.")
    equipment: str = dspy.InputField(desc="Final equipment.")
    inventory: dict = dspy.InputField(desc="Final inventory.")
    task: str = dspy.InputField(desc="The objective to accomplish.")
    context: str = dspy.InputField(desc="Task context.")
    chest_observation: str = dspy.InputField(desc="Observation of chests")

    reasoning: str = dspy.OutputField(desc="reasoning")
    success: bool = dspy.OutputField(desc="boolean")
    critique: str = dspy.OutputField(desc="critique")

class BeliefCritic(dspy.Signature):
    """Generate action critiques."""
    generation_instructions: str = dspy.InputField(desc="Instructions about your role and what you should generate.")
    biome: str = dspy.InputField(desc="The biome after the task execution.")
    time_of_day: str = dspy.InputField(desc="The current time.")
    nearby_blocks: list = dspy.InputField(desc="The surrounding blocks. These blocks are not collected yet. However, this is useful for some placing or planting tasks.")
    health: float = dspy.InputField(desc="Current health.")
    hunger: float = dspy.InputField(desc="My current hunger level. For eating task, if my hunger level is 20.0, then I successfully ate the food.")
    position: dict = dspy.InputField(desc="Current position.")
    equipment: str = dspy.InputField(desc="My final equipment. For crafting tasks, I sometimes equip the crafted item.")
    inventory: list = dspy.InputField(desc="My final inventory. For mining and smelting tasks, you only need to check inventory.")
    task: str = dspy.InputField(desc="The objective to accomplish.")
    context: str = dspy.InputField(desc="The context of the task.")
    chest_observation: str = dspy.InputField(desc="If the task requires me to place items in a chest, you can find chest information here.")

    reasoning: str = dspy.OutputField(desc="Reasoning for success or failure of the task based on the inventory, equipment, task, position, chest observations and context", prefix="You are required to evaluate if I have met the task requirements based on the input fields provided. Exceeding the task requirements is also considered a success while failing to meet them requires you to provide the reasoning why")
    success: bool = dspy.OutputField(desc="Task completion check.", prefix="Based on the reasoning above, decide if the task was completed successfully or not")
    critique: str = dspy.OutputField(desc="Instructive critique when you failed to complete the task.", prefix="In case of task failure, why is the task not complete and how can I improve.")


critic_demonstrations = [
    Example(
        generation_instructions=CRITIC_INSTRUCTIONS,
        biome="",
        time_of_day="",
        nearby_blocks=[],
        health=20.0,
        hunger=20.0,
        position={},
        equipment="",
        inventory={'oak_log':2, 'spruce_log':2},
        task="Mine 3 wood logs",
        context="",
        chest_observation="",
        reasoning="You need to mine 3 wood logs. You have 2 oak logs and 2 spruce logs, which add up to 4 wood logs.",
        success=True,
        critique=""
    ),
    Example(
        generation_instructions=CRITIC_INSTRUCTIONS,
       biome="",
        time_of_day="",
        nearby_blocks=[],
        health=20.0,
        hunger=20.0,
        position={},
        equipment="",
        inventory={'crafting_table': 1, 'spruce_planks': 6, 'stick': 4},
        task="Craft a wooden pickaxe",
        context="",
        chest_observation="",
        reasoning="You have enough materials to craft a wooden pickaxe, but you didn't craft it.",
        success=False,
        critique="Craft a wooden pickaxe with a crafting table using 3 spruce planks and 2 sticks."
    ),
    Example(
        generation_instructions=CRITIC_INSTRUCTIONS,
        biome="",
        time_of_day="",
        nearby_blocks=[],
        health=20.0,
        hunger=20.0,
        position={},
        equipment="",
        inventory={'raw_iron': 5, 'stone_pickaxe': 1},
        task="Mine 5 iron_ore",
        context="",
        chest_observation="",
        reasoning="Mining iron_ore in Minecraft will get raw_iron. You have 5 raw_iron in your inventory.",
        success=True,
        critique=""
    ),
    Example(
        generation_instructions=CRITIC_INSTRUCTIONS,
        biome="",
        time_of_day="",
        nearby_blocks=["stone", "dirt", "grass_block", "grass", "farmland", "wheat"],
        health=20.0,
        hunger=20.0,
        position={},
        equipment="",
        inventory={},
        task="Plant 1 wheat seed.",
        context="",
        chest_observation="",
        reasoning="For planting tasks, inventory information is useless. In nearby blocks, there is farmland and wheat, which means you succeed to plant the wheat seed.",
        success=True,
        critique=""
    ),
   Example(
       generation_instructions=CRITIC_INSTRUCTIONS,
        biome="",
        time_of_day="",
        nearby_blocks=[],
        health=20.0,
        hunger=20.0,
        position={},
        equipment="",
        inventory={'rotten_flesh': 1},
        task="Kill 1 zombie",
        context="",
        chest_observation="",
        reasoning="You have rotten flesh in your inventory, which means you successfully killed one zombie.",
        success=True,
        critique=""
    ),
    Example(
        generation_instructions=CRITIC_INSTRUCTIONS,
         biome="",
        time_of_day="",
        nearby_blocks=["chest"],
        health=20.0,
        hunger=20.0,
        position={},
        equipment="",
        inventory={'rail': 1, 'coal': 2, 'oak_planks': 13, 'copper_block': 1, 'diorite': 7, 'cooked_beef': 4, 'granite': 22, 'cobbled_deepslate': 23, 'feather': 4, 'leather': 2, 'cooked_chicken': 3, 'white_wool': 2, 'stick': 3, 'black_wool': 1, 'stone_sword': 2, 'stone_hoe': 1, 'stone_axe': 2, 'stone_shovel': 2, 'cooked_mutton': 4, 'cobblestone_wall': 18, 'crafting_table': 1, 'furnace': 1, 'iron_pickaxe': 1, 'stone_pickaxe': 1, 'raw_copper': 12},
        task="Deposit useless items into the chest at (81, 131, 16)",
        context="",
        chest_observation="(81, 131, 16): {'andesite': 2, 'dirt': 2, 'cobblestone': 75, 'wooden_pickaxe': 1, 'wooden_sword': 1}",
        reasoning="You have 28 items in your inventory after depositing, which is more than 20. You need to deposit more items from your inventory to the chest.",
        success=False,
        critique="Deposit more useless items such as copper_block, diorite, granite, cobbled_deepslate, feather, and leather to meet the requirement of having only 20 occupied slots in your inventory."
    ),
    Example(
        generation_instructions=CRITIC_INSTRUCTIONS,
         biome="",
        time_of_day="",
        nearby_blocks=[],
        health=20.0,
        hunger=20.0,
        position={},
        equipment="",
        inventory={},
        task="Eat 1 ...",
        context="",
        chest_observation="",
        reasoning="For all eating task, if the player's hunger is 20.0, then the player successfully ate the food.",
        success=True,
        critique=""
    )
]

class CriticAgent(LLMAgent):
    def __init__(
        self,
        name: str,
        llm: str,
        temperature: float,
        request_timeout: int,
        mode: str,
        logger
    ):
        super().__init__(name, llm, temperature, logger)
        assert mode in ["auto", "manual"]
        self.mode = mode
        self.request_timeout = request_timeout
        self.generate_critic = None

    def __call__(self, context: dict):
        generate_critic = dspy.Predict(BeliefCritic)
        optimizer = LabeledFewShot(k=len(critic_demonstrations))
        generate_critic = optimizer.compile(generate_critic, trainset=critic_demonstrations)
        critique = generate_critic(
            generation_instructions=CRITIC_INSTRUCTIONS,
            biome=context["biome"],
            time_of_day=context["timeOfDay"],
            nearby_blocks=context["voxels"] if context["voxels"] else [],
            health=context["health"],
            hunger=context["food"],
            position=context["position"],
            equipment=context["equipment"],
            inventory=context["inventory"],
            task=context["task"],
            context=context["context"],
            chest_observation=context["chest_observation"]
        )
        return critique

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

    def restore(self):
        # TODO: implement this
        pass

if __name__ == "__main__":
    import os
    os.environ["OPENAI_API_KEY"] = "ADD_KEY_HERE"
    TOGETHER_API_KEY = "ADD_KEY_HERE"
    os.environ["TOGETHER_API_KEY"] = TOGETHER_API_KEY
    critic_agent = CriticAgent(
        name="critic_agent",
        llm="together_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo",
        temperature=0.7,
        request_timeout=120,
        mode="auto",
        logger=None
    )

    # Test data
    context = {
        "biome": "forest",
        "timeOfDay": "day",
        "voxels": ['oak_log', 'dirt', 'grass_block', 'leaves'],
        "nearby_entities": {
            "sheep": [15, 64, 18],
            "cow": [12, 65, 20]
        },
        "health": 20.0,
        "food": 18.0,
        "position": {'x': 10.0, 'y': 64.0, 'z': 10.0},
        "equipment": "wooden_axe",
        "inventory": {'wooden_axe': 1, 'oak_log': 16, 'dirt': 5},
        "task": "Build a shelter",
        "context": "I have gathered some wood and need to find a suitable location for my shelter.",
        "chest_observation": "Chests:\n(12.0,64.0,12.0): {'oak_log': 10}\n\n"
    }

    critique = critic_agent(context)
    print(f"Success: {critique.success}, Critique: {critique.critique}, Reasoning: {critique.reasoning}")
    print(critic_agent.lm.inspect_history(n=1))
