import unittest
import asyncio
from pygeai_orchestration import (
    AgentConfig,
    BaseAgent,
    PatternConfig,
    BasePattern,
    PatternType,
    PatternResult,
    OrchestratorConfig,
    BaseOrchestrator,
    ToolConfig,
    BaseTool,
    ToolResult,
    ToolCategory
)


class ConcreteAgent(BaseAgent):
    """Concrete implementation for testing."""

    async def execute(self, task: str, context=None):
        return {"success": True, "result": "Executed"}

    async def generate(self, prompt: str, **kwargs):
        return "Generated response"


class ConcretePattern(BasePattern):
    """Concrete implementation for testing."""

    async def execute(self, task: str, context=None):
        return PatternResult(success=True, result="Done", iterations=1)

    async def step(self, state):
        return {"status": "ok"}


class ConcreteOrchestrator(BaseOrchestrator):
    """Concrete implementation for testing."""

    async def orchestrate(self, task: str, pattern: str, context=None):
        return PatternResult(success=True, result="Orchestrated", iterations=1)

    async def coordinate_agents(self, agents, task: str):
        return {"success": True, "results": {}}


class ConcreteTool(BaseTool):
    """Concrete implementation for testing."""

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

    def validate_parameters(self, parameters):
        return True


class TestBaseAgent(unittest.TestCase):

    def test_creation(self):
        config = AgentConfig(name="test-agent", model="gpt-4")
        agent = ConcreteAgent(config)

        self.assertEqual(agent.name, "test-agent")
        self.assertEqual(agent.config.model, "gpt-4")
        self.assertEqual(len(agent.get_history()), 0)

    def test_history(self):
        config = AgentConfig(name="test-agent", model="gpt-4")
        agent = ConcreteAgent(config)

        agent.add_to_history({"test": "data"})
        self.assertEqual(len(agent.get_history()), 1)

        agent.clear_history()
        self.assertEqual(len(agent.get_history()), 0)

    def test_execute(self):
        config = AgentConfig(name="test-agent", model="gpt-4")
        agent = ConcreteAgent(config)

        result = asyncio.run(agent.execute("Test task"))
        self.assertTrue(result["success"])


class TestBasePattern(unittest.TestCase):

    def test_creation(self):
        config = PatternConfig(
            name="test-pattern",
            pattern_type=PatternType.REFLECTION,
            max_iterations=5
        )
        pattern = ConcretePattern(config)

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

    def test_iteration_tracking(self):
        config = PatternConfig(name="test", pattern_type=PatternType.REFLECTION)
        pattern = ConcretePattern(config)

        self.assertEqual(pattern.current_iteration, 0)
        pattern.increment_iteration()
        self.assertEqual(pattern.current_iteration, 1)
        pattern.reset()
        self.assertEqual(pattern.current_iteration, 0)

    def test_execute(self):
        config = PatternConfig(name="test", pattern_type=PatternType.REFLECTION)
        pattern = ConcretePattern(config)

        result = asyncio.run(pattern.execute("Test"))
        self.assertTrue(result.success)
        self.assertEqual(result.iterations, 1)


class TestBaseOrchestrator(unittest.TestCase):

    def test_creation(self):
        config = OrchestratorConfig(name="test-orchestrator")
        orch = ConcreteOrchestrator(config)

        self.assertEqual(orch.name, "test-orchestrator")
        self.assertEqual(len(orch.list_agents()), 0)
        self.assertEqual(len(orch.list_patterns()), 0)

    def test_register_agent(self):
        config = OrchestratorConfig(name="test-orchestrator")
        orch = ConcreteOrchestrator(config)

        agent_config = AgentConfig(name="agent1", model="gpt-4")
        agent = ConcreteAgent(agent_config)
        orch.register_agent(agent)

        self.assertEqual(len(orch.list_agents()), 1)
        self.assertIsNotNone(orch.get_agent("agent1"))

    def test_register_pattern(self):
        config = OrchestratorConfig(name="test-orchestrator")
        orch = ConcreteOrchestrator(config)

        pattern_config = PatternConfig(name="pattern1", pattern_type=PatternType.REFLECTION)
        pattern = ConcretePattern(pattern_config)
        orch.register_pattern(pattern)

        self.assertEqual(len(orch.list_patterns()), 1)
        self.assertIsNotNone(orch.get_pattern("pattern1"))


class TestBaseTool(unittest.TestCase):

    def test_creation(self):
        config = ToolConfig(
            name="test-tool",
            description="A test tool",
            category=ToolCategory.CUSTOM
        )
        tool = ConcreteTool(config)

        self.assertEqual(tool.name, "test-tool")
        self.assertEqual(tool.category, ToolCategory.CUSTOM)

    def test_schema(self):
        config = ToolConfig(
            name="test-tool",
            description="A test tool",
            category=ToolCategory.SEARCH
        )
        tool = ConcreteTool(config)

        schema = tool.get_schema()
        self.assertEqual(schema["name"], "test-tool")
        self.assertEqual(schema["category"], "search")

    def test_execute(self):
        config = ToolConfig(name="test-tool", description="Test", category=ToolCategory.CUSTOM)
        tool = ConcreteTool(config)

        result = asyncio.run(tool.execute())
        self.assertTrue(result.success)


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