"""
Test component serialization across the picoagents framework.

Ensures all core components can serialize/deserialize correctly.
"""

import pytest
from picoagents.llm import OpenAIChatCompletionClient
from picoagents.memory import ListMemory, FileMemory
from picoagents.tools import FunctionTool
from picoagents.agents import Agent
from picoagents.termination import MaxMessageTermination, TextMentionTermination, CompositeTermination
from picoagents.orchestration import RoundRobinOrchestrator, AIOrchestrator


def test_llm_client_serialization():
    """Test LLM client serialization roundtrip."""
    # Create client
    client = OpenAIChatCompletionClient(
        model="gpt-4o-mini",
        api_key="test-key",
        base_url="https://test.openai.com/v1/"
    )
    
    # Serialize and deserialize
    component_model = client.dump_component()
    loaded_client = OpenAIChatCompletionClient.load_component(component_model)
    
    # Verify
    assert loaded_client.model == client.model
    assert loaded_client.api_key == client.api_key
    assert str(loaded_client.client.base_url) == "https://test.openai.com/v1/"


def test_memory_serialization():
    """Test memory serialization roundtrip."""
    # Test ListMemory
    list_memory = ListMemory(max_memories=50)
    component_model = list_memory.dump_component()
    loaded_memory = ListMemory.load_component(component_model)
    
    assert loaded_memory.max_memories == 50
    assert loaded_memory.memories == []  # Should start empty
    
    # Test FileMemory
    file_memory = FileMemory("test.json", max_memories=100)
    component_model = file_memory.dump_component()
    loaded_file_memory = FileMemory.load_component(component_model)
    
    assert loaded_file_memory.file_path == "test.json"
    assert loaded_file_memory.max_memories == 100


def test_function_tool_serialization_blocked():
    """Test that FunctionTool serialization is properly blocked."""
    def test_func():
        return "test"
    
    tool = FunctionTool(test_func)
    
    with pytest.raises(NotImplementedError, match="cannot be serialized for security reasons"):
        tool.dump_component()


def test_termination_serialization():
    """Test termination condition serialization."""
    # Test MaxMessageTermination
    max_term = MaxMessageTermination(max_messages=10)
    component_model = max_term.dump_component()
    loaded_term = MaxMessageTermination.load_component(component_model)
    
    assert loaded_term.max_messages == 10
    assert loaded_term.message_count == 0  # Should reset
    
    # Test TextMentionTermination
    text_term = TextMentionTermination("DONE", case_sensitive=True)
    component_model = text_term.dump_component()
    loaded_text_term = TextMentionTermination.load_component(component_model)
    
    assert loaded_text_term.text == "DONE"
    assert loaded_text_term.case_sensitive is True
    
    # Test CompositeTermination
    composite = CompositeTermination([max_term, text_term], mode="any")
    component_model = composite.dump_component()
    loaded_composite = CompositeTermination.load_component(component_model)
    
    assert len(loaded_composite.conditions) == 2
    assert loaded_composite.mode == "any"


def test_agent_serialization():
    """Test agent serialization with nested components."""
    # Create components
    model_client = OpenAIChatCompletionClient(model="gpt-4o-mini", api_key="test-key")
    memory = ListMemory(max_memories=20)
    
    # Create agent
    agent = Agent(
        name="TestAgent",
        description="A test agent",
        instructions="Test instructions",
        model_client=model_client,
        memory=memory,
        max_iterations=8
    )
    
    # Serialize and deserialize
    component_model = agent.dump_component()
    loaded_agent = Agent.load_component(component_model)
    
    # Verify
    assert loaded_agent.name == "TestAgent"
    assert loaded_agent.description == "A test agent"
    assert loaded_agent.instructions == "Test instructions"
    assert loaded_agent.max_iterations == 8
    assert loaded_agent.model_client.model == "gpt-4o-mini"
    assert loaded_agent.memory is not None
    assert loaded_agent.memory.max_memories == 20
    assert isinstance(loaded_agent.memory, ListMemory)


def test_round_robin_orchestrator_serialization():
    """Test RoundRobinOrchestrator serialization."""
    # Create components
    model_client1 = OpenAIChatCompletionClient(model="gpt-4o-mini", api_key="test-1")
    model_client2 = OpenAIChatCompletionClient(model="gpt-4o-mini", api_key="test-2")
    
    agent1 = Agent("Agent1", "First agent", "You are agent 1", model_client1)
    agent2 = Agent("Agent2", "Second agent", "You are agent 2", model_client2)
    
    termination = MaxMessageTermination(max_messages=5)
    
    # Create orchestrator
    orchestrator = RoundRobinOrchestrator(
        agents=[agent1, agent2],
        termination=termination,
        max_iterations=10
    )
    
    # Serialize and deserialize
    component_model = orchestrator.dump_component()
    loaded_orchestrator = RoundRobinOrchestrator.load_component(component_model)
    
    # Verify
    assert len(loaded_orchestrator.agents) == 2
    assert loaded_orchestrator.agents[0].name == "Agent1"
    assert loaded_orchestrator.agents[1].name == "Agent2"
    assert loaded_orchestrator.max_iterations == 10
    assert isinstance(loaded_orchestrator.termination, MaxMessageTermination)
    assert loaded_orchestrator.termination.max_messages == 5
    # Runtime state should reset
    assert loaded_orchestrator.current_agent_index == 0


def test_ai_orchestrator_serialization():
    """Test AIOrchestrator serialization."""
    # Create components
    model_client = OpenAIChatCompletionClient(model="gpt-4o-mini", api_key="test-key")
    selector_client = OpenAIChatCompletionClient(model="gpt-4o-mini", api_key="selector-key")
    
    agent1 = Agent("Agent1", "First agent", "You are agent 1", model_client)
    agent2 = Agent("Agent2", "Second agent", "You are agent 2", model_client)
    
    termination = TextMentionTermination("COMPLETE")
    
    # Create AI orchestrator
    ai_orchestrator = AIOrchestrator(
        agents=[agent1, agent2],
        termination=termination,
        model_client=selector_client,
        max_iterations=15
    )
    
    # Serialize and deserialize
    component_model = ai_orchestrator.dump_component()
    loaded_ai = AIOrchestrator.load_component(component_model)
    
    # Verify
    assert len(loaded_ai.agents) == 2
    assert loaded_ai.agents[0].name == "Agent1"
    assert loaded_ai.agents[1].name == "Agent2"
    assert loaded_ai.max_iterations == 15
    assert loaded_ai.model_client.api_key == "selector-key"
    assert isinstance(loaded_ai.termination, TextMentionTermination)
    assert loaded_ai.termination.text == "COMPLETE"
    # Runtime state should reset
    assert loaded_ai.selection_history == []
    assert loaded_ai.agent_capabilities_cache is None


def test_nested_serialization_integrity():
    """Test that deeply nested serialization preserves all data integrity."""
    # Create a complex nested structure
    model_client = OpenAIChatCompletionClient(model="gpt-4o-mini", api_key="nested-test")
    memory = FileMemory("nested_test.json", max_memories=100)
    
    agent = Agent(
        name="NestedAgent",
        description="Complex nested agent",
        instructions="Complex instructions",
        model_client=model_client,
        memory=memory
    )
    
    # Create composite termination
    max_term = MaxMessageTermination(25)
    text_term = TextMentionTermination("FINISH", case_sensitive=False)
    composite_term = CompositeTermination([max_term, text_term], mode="all")
    
    # Create AI orchestrator with all the nested components
    orchestrator = AIOrchestrator(
        agents=[agent],
        termination=composite_term,
        model_client=OpenAIChatCompletionClient(model="gpt-4o-mini", api_key="orchestrator-key"),
        max_iterations=20
    )
    
    # Full serialization roundtrip
    component_model = orchestrator.dump_component()
    loaded_orchestrator = AIOrchestrator.load_component(component_model)
    
    # Deep verification
    loaded_agent = loaded_orchestrator.agents[0]
    assert loaded_agent.name == "NestedAgent"
    assert loaded_agent.model_client.api_key == "nested-test"
    assert loaded_agent.memory is not None
    assert isinstance(loaded_agent.memory, FileMemory)
    assert loaded_agent.memory.file_path == "nested_test.json"
    assert loaded_agent.memory.max_memories == 100
    
    loaded_termination = loaded_orchestrator.termination
    assert isinstance(loaded_termination, CompositeTermination)
    assert len(loaded_termination.conditions) == 2
    assert loaded_termination.mode == "all"
    
    # Verify nested termination conditions
    conditions = loaded_termination.conditions
    max_condition = next(c for c in conditions if isinstance(c, MaxMessageTermination))
    text_condition = next(c for c in conditions if isinstance(c, TextMentionTermination))
    
    assert max_condition.max_messages == 25
    assert text_condition.text == "FINISH"
    assert text_condition.case_sensitive is False


if __name__ == "__main__":
    # Run tests manually if needed
    pytest.main([__file__, "-v"])