import unittest
import pytest
from mcpuniverse.mcp.manager import MCPManager
from mcpuniverse.llm.manager import ModelManager
from mcpuniverse.agent.explore_and_exploit import ExploreAndExploit
from mcpuniverse.tracer import Tracer
from mcpuniverse.callbacks.base import Printer
from mcpuniverse.callbacks.handlers.vprint import get_vprint_callbacks


class TestExploreAndExploitAgent(unittest.IsolatedAsyncioTestCase):

    @unittest.skip("Skipping exploration phase test")
    async def test_prompt_exploration_phase(self):
        """Test prompt building during exploration phase"""
        agent = ExploreAndExploit(
            mcp_manager=MCPManager(),
            llm=ModelManager().build_model(name="openai"),
            config={"servers": [{"name": "weather"}], "exploration_phase_iterations": 3}
        )
        await agent.initialize()

        # Set up exploration phase
        agent._current_iteration = 1
        agent._exploration_complete = False
        agent._history.append("Thought: I want to explore the weather tools")
        agent._history.append("Action: Using tool `get_weather` in server `weather`")

        prompt = agent._build_prompt(question="What's the weather in San Francisco now?")

        # Check exploration phase content
        self.assertTrue("EXPLORATION PHASE" in prompt)
        self.assertTrue("exploration" in prompt)
        self.assertTrue("What's the weather in San Francisco now?" in prompt)
        self.assertTrue("I want to explore the weather tools" in prompt)

        await agent.cleanup()

    @unittest.skip("Skipping exploitation phase test")
    async def test_prompt_exploitation_phase(self):
        """Test prompt building during exploitation phase"""
        agent = ExploreAndExploit(
            mcp_manager=MCPManager(),
            llm=ModelManager().build_model(name="openai"),
            config={"servers": [{"name": "weather"}], "exploration_phase_iterations": 2}
        )
        await agent.initialize()

        # Set up exploitation phase
        agent._current_iteration = 3
        agent._exploration_complete = True
        agent._history.append("Knowledge_Update: Updated knowledge for weather.get_weather")

        # Add some tool knowledge
        agent._update_tool_knowledge(
            server="weather",
            tool_name="get_weather",
            purpose="Retrieves current weather information",
            caveats="Returns weather data including temperature and conditions",
        )

        prompt = agent._build_prompt(question="What's the weather in San Francisco now?")

        # Check exploitation phase content
        self.assertTrue("EXPLOITATION PHASE" in prompt)
        self.assertTrue("exploitation" in prompt)
        self.assertTrue("Tool: weather.get_weather" in prompt)
        self.assertTrue("Purpose: Retrieves current weather information" in prompt)
        self.assertTrue("Updated knowledge for weather.get_weather" in prompt)

        await agent.cleanup()

    @unittest.skip("Skipping tool knowledge management test")
    async def test_tool_knowledge_management(self):
        """Test tool knowledge storage and retrieval"""
        agent = ExploreAndExploit(
            mcp_manager=MCPManager(),
            llm=ModelManager().build_model(name="openai"),
            config={"servers": [{"name": "weather"}]}
        )
        await agent.initialize()

        # Test initial empty knowledge
        self.assertEqual(len(agent.get_tool_knowledge()), 0)

        # Add tool knowledge
        agent._update_tool_knowledge(
            server="weather",
            tool_name="get_weather",
            purpose="Gets weather data",
            caveats="Returns temperature and conditions"
        )

        agent._update_tool_knowledge(
            server="weather",
            tool_name="get_forecast",
            purpose="Gets weather forecast",
            caveats="Returns multi-day forecast"
        )

        # Test knowledge retrieval
        knowledge = agent.get_tool_knowledge()
        self.assertEqual(len(knowledge), 2)
        self.assertIn("weather.get_weather", knowledge)
        self.assertIn("weather.get_forecast", knowledge)
        self.assertEqual(knowledge["weather.get_weather"]["purpose"], "Gets weather data")
        self.assertEqual(knowledge["weather.get_forecast"]["caveats"], "Returns multi-day forecast")

        await agent.cleanup()

    @unittest.skip("Skipping knowledge formatting test")
    async def test_knowledge_formatting(self):
        """Test tool knowledge formatting for prompts"""
        agent = ExploreAndExploit(
            mcp_manager=MCPManager(),
            llm=ModelManager().build_model(name="openai"),
            config={"servers": [{"name": "weather"}]}
        )
        await agent.initialize()

        # Test empty knowledge formatting
        formatted = agent._format_tool_knowledge()
        self.assertEqual(formatted, "No tool knowledge accumulated yet.")

        # Add knowledge and test formatting
        agent._update_tool_knowledge(
            server="weather",
            tool_name="get_weather",
            purpose="Retrieves weather data",
            caveats="Checking current weather"
        )

        formatted = agent._format_tool_knowledge()
        self.assertTrue("Tool: weather.get_weather" in formatted)
        self.assertTrue("Purpose: Retrieves weather data" in formatted)

        await agent.cleanup()

    @unittest.skip("Skipping phase transition test")
    async def test_phase_transition(self):
        """Test automatic phase transition from exploration to exploitation"""
        agent = ExploreAndExploit(
            mcp_manager=MCPManager(),
            llm=ModelManager().build_model(name="openai"),
            config={"servers": [{"name": "weather"}], "exploration_phase_iterations": 2}
        )
        await agent.initialize()

        # Initially in exploration phase
        self.assertFalse(agent._exploration_complete)

        # Simulate iterations
        agent._current_iteration = 0
        self.assertFalse(agent._exploration_complete)

        agent._current_iteration = 1
        self.assertFalse(agent._exploration_complete)

        # Should transition to exploitation at iteration 2
        agent._current_iteration = 2
        if agent._current_iteration >= agent._config.exploration_phase_iterations:
            agent._exploration_complete = True
        self.assertTrue(agent._exploration_complete)

        await agent.cleanup()

    @unittest.skip("Skipping final answer restriction test")
    async def test_final_answer_restriction(self):
        """Test that final answers are blocked during exploration phase"""
        agent = ExploreAndExploit(
            mcp_manager=MCPManager(),
            llm=ModelManager().build_model(name="openai"),
            config={"servers": [{"name": "weather"}], "exploration_phase_iterations": 3}
        )
        await agent.initialize()

        from mcpuniverse.tracer import Tracer
        tracer = Tracer()

        # Test during exploration phase
        agent._current_iteration = 1
        agent._exploration_complete = False

        # Try to provide a final answer during exploration
        parsed_response = {"thought": "I think I know the answer", "answer": "It's sunny"}
        result = await agent._handle_final_answer(parsed_response, tracer, [])

        # Should return None (blocked)
        self.assertIsNone(result)

        # Check that error was added to history
        self.assertTrue(
            any("Final answers are not allowed during exploration phase" in entry for entry in agent._history))

        # Test during exploitation phase
        agent._exploration_complete = True

        # Try to provide a final answer during exploitation
        result = await agent._handle_final_answer(parsed_response, tracer, [])

        # Should return AgentResponse (allowed)
        self.assertIsNotNone(result)
        self.assertEqual(result.response, "It's sunny")

        await agent.cleanup()

    @unittest.skip("Skipping reset functionality test")
    async def test_reset_functionality(self):
        """Test agent reset clears all state"""
        agent = ExploreAndExploit(
            mcp_manager=MCPManager(),
            llm=ModelManager().build_model(name="openai"),
            config={"servers": [{"name": "weather"}]}
        )
        await agent.initialize()

        # Add some state
        agent._history.append("Some history")
        agent._update_tool_knowledge("weather", "get_weather", "purpose", "caveats")
        agent._exploration_complete = True
        agent._current_iteration = 5

        # Verify state exists
        self.assertTrue(len(agent._history) > 0)
        self.assertTrue(len(agent._tool_knowledge) > 0)
        self.assertTrue(agent._exploration_complete)
        self.assertEqual(agent._current_iteration, 5)

        # Reset and verify clean state
        agent.reset()
        self.assertEqual(len(agent._history), 0)
        self.assertEqual(len(agent._tool_knowledge), 0)
        self.assertFalse(agent._exploration_complete)
        self.assertEqual(agent._current_iteration, 0)

        await agent.cleanup()

    @unittest.skip("Skipping full execution test")
    async def test_execute(self):
        """Test full execution with exploration and exploitation phases"""
        question = "I live in San Francisco. Do I need to bring an umbrella if I need to go outside?"
        tracer = Tracer()
        agent = ExploreAndExploit(
            mcp_manager=MCPManager(),
            llm=ModelManager().build_model(name="openai"),
            config={
                "servers": [{"name": "weather"}],
                "max_iterations": 6,
                "exploration_phase_iterations": 2
            }
        )
        await agent.initialize()

        response = await agent.execute(
            message=question,
            tracer=tracer,
            callbacks=[Printer()] + get_vprint_callbacks()
        )

        print("=" * 50)
        print("RESPONSE:")
        print(response)
        print("=" * 50)
        print("HISTORY:")
        print(agent.get_history())
        print("=" * 50)
        print("TOOL KNOWLEDGE:")
        for tool, knowledge in agent.get_tool_knowledge().items():
            print(f"{tool}: {knowledge}")
        print("=" * 50)

        await agent.cleanup()

    @unittest.skip("Skipping config parameters test")
    async def test_config_parameters(self):
        """Test agent configuration parameters"""
        config = {
            "servers": [{"name": "weather"}],
            "max_iterations": 10,
            "exploration_phase_iterations": 4,
            "instruction": "Test agent for weather tasks"
        }

        agent = ExploreAndExploit(
            mcp_manager=MCPManager(),
            llm=ModelManager().build_model(name="openai"),
            config=config
        )
        await agent.initialize()

        # Verify configuration
        self.assertEqual(agent._config.max_iterations, 10)
        self.assertEqual(agent._config.exploration_phase_iterations, 4)
        self.assertEqual(agent._config.instruction, "Test agent for weather tasks")

        await agent.cleanup()


if __name__ == "__main__":
    unittest.main()