import unittest
import asyncio
from pygeai_orchestration import (
    ReflectionPattern,
    ToolUsePattern,
    ReActPattern,
    PlanningPattern,
    MultiAgentPattern,
    AgentRole,
    BaseTool,
    ToolConfig,
    ToolResult,
    ToolCategory,
    PatternType
)


class MockAgent:
    """Mock agent for testing."""

    async def generate(self, prompt: str) -> str:
        if "Review your previous response" in prompt:
            return "Improved response"
        elif "TOOL_CALL" in prompt or "FINAL_ANSWER" in prompt:
            return "FINAL_ANSWER: Test answer"
        elif "Thought:" in prompt or "Action:" in prompt:
            return "Thought: Test\nAction: ANSWER Test result\nObservation: Done"
        elif "step-by-step plan" in prompt:
            return "1. First step\n2. Second step"
        elif "Execute the following step" in prompt:
            return "Step completed"
        elif "synthesized answer" in prompt or "Synthesize" in prompt:
            return "Final synthesized result"
        elif "coordination plan" in prompt:
            return "Coordination plan created"
        return "Test response"


class MockTool(BaseTool):
    """Mock tool for testing."""

    def __init__(self):
        config = ToolConfig(
            name="mock_tool",
            description="Mock tool",
            category=ToolCategory.CUSTOM
        )
        super().__init__(config)

    async def execute(self, **kwargs) -> ToolResult:
        return ToolResult(success=True, result="Tool result")

    def validate_parameters(self, parameters):
        return True


class TestReflectionPattern(unittest.TestCase):

    def test_creation(self):
        agent = MockAgent()
        pattern = ReflectionPattern(agent=agent)

        self.assertEqual(pattern.name, "reflection")
        self.assertEqual(pattern.pattern_type, PatternType.REFLECTION)
        self.assertEqual(pattern.current_iteration, 0)

    def test_execute(self):
        agent = MockAgent()
        pattern = ReflectionPattern(agent=agent)

        result = asyncio.run(pattern.execute("Test task"))

        self.assertTrue(result.success)
        self.assertIsNotNone(result.result)
        self.assertGreater(result.iterations, 0)


class TestToolUsePattern(unittest.TestCase):

    def test_creation(self):
        agent = MockAgent()
        tool = MockTool()
        pattern = ToolUsePattern(agent=agent, tools=[tool])

        self.assertEqual(pattern.name, "tool_use")
        self.assertEqual(pattern.pattern_type, PatternType.TOOL_USE)
        self.assertIn("mock_tool", pattern.tools)

    def test_execute(self):
        agent = MockAgent()
        tool = MockTool()
        pattern = ToolUsePattern(agent=agent, tools=[tool])

        result = asyncio.run(pattern.execute("Test task"))

        self.assertTrue(result.success)
        self.assertIsNotNone(result.result)


class TestReActPattern(unittest.TestCase):

    def test_creation(self):
        agent = MockAgent()
        pattern = ReActPattern(agent=agent)

        self.assertEqual(pattern.name, "react")
        self.assertEqual(pattern.pattern_type, PatternType.REACT)

    def test_execute(self):
        agent = MockAgent()
        pattern = ReActPattern(agent=agent, tools=[])

        result = asyncio.run(pattern.execute("Test task"))

        self.assertTrue(result.success)
        self.assertIn('thoughts', result.metadata)
        self.assertIn('actions', result.metadata)


class TestPlanningPattern(unittest.TestCase):

    def test_creation(self):
        agent = MockAgent()
        pattern = PlanningPattern(agent=agent)

        self.assertEqual(pattern.name, "planning")
        self.assertEqual(pattern.pattern_type, PatternType.PLANNING)

    def test_execute(self):
        agent = MockAgent()
        pattern = PlanningPattern(agent=agent)

        result = asyncio.run(pattern.execute("Test task"))

        self.assertTrue(result.success)
        self.assertIn('plan', result.metadata)
        self.assertGreater(len(pattern.get_plan()), 0)


class TestMultiAgentPattern(unittest.TestCase):

    def test_creation(self):
        agent = MockAgent()
        role = AgentRole(name="test", agent=agent, role_description="Test role")
        pattern = MultiAgentPattern(agents=[role], coordinator_agent=agent)

        self.assertEqual(pattern.name, "multi_agent")
        self.assertEqual(pattern.pattern_type, PatternType.MULTI_AGENT)
        self.assertIn("test", pattern.agent_roles)

    def test_execute(self):
        agent = MockAgent()
        role1 = AgentRole(name="agent1", agent=agent, role_description="Role 1")
        role2 = AgentRole(name="agent2", agent=agent, role_description="Role 2")
        pattern = MultiAgentPattern(agents=[role1, role2], coordinator_agent=agent)

        result = asyncio.run(pattern.execute("Test task"))

        self.assertTrue(result.success)
        self.assertIn('contributions', result.metadata)
        self.assertEqual(len(result.metadata['contributions']), 2)


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