#!/usr/bin/env python3
"""
Tests for Sub-Agent Context Management

Verifies:
1. Sub-agent context creation and retrieval
2. Phase progression tracking
3. Sibling agent discovery and coordination
4. Context persistence
5. Flow integration
"""

import pytest
import tempfile
from pathlib import Path
import sys
import os

# Add parent directory to path for imports
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'lib'))

from sub_agent_context import (
    SubAgentContext,
    SubAgentContextManager,
    FlowPhase,
)
from sub_agent_flow_integration import SubAgentFlowIntegrator


class TestSubAgentContextCreation:
    """Test sub-agent context creation"""

    def setup_method(self):
        """Setup test fixtures"""
        self.temp_dir = tempfile.TemporaryDirectory()
        self.manager = SubAgentContextManager(Path(self.temp_dir.name))

    def teardown_method(self):
        """Cleanup"""
        self.temp_dir.cleanup()

    def test_create_sub_agent_context(self):
        """Test creating a new sub-agent context"""
        context = self.manager.create_sub_agent_context(
            parent_task_id="task-123",
            parent_project="admin",
            parent_description="Test parent task",
            parent_context={"key": "value"},
            parent_tags=["important", "research"],
        )

        assert context.sub_agent_id is not None
        assert context.parent_task_id == "task-123"
        assert context.parent_project == "admin"
        assert context.parent_description == "Test parent task"
        assert len(context.phase_progression) == 9
        assert context.phase_progression[0].phase_name == "CONTEXT_PREP"

    def test_phase_progression_initialization(self):
        """Test that all 9 phases are initialized"""
        context = self.manager.create_sub_agent_context(
            parent_task_id="task-456",
            parent_project="test",
            parent_description="Phase test",
        )

        phase_names = [p.phase_name for p in context.phase_progression]
        expected_phases = [
            "CONTEXT_PREP",
            "RECEIVED",
            "PREDICTING",
            "ANALYZING",
            "CONSENSUS_CHECK",
            "AWAITING_APPROVAL",
            "STRATEGIZING",
            "EXECUTING",
            "LEARNING",
        ]

        assert phase_names == expected_phases

    def test_retrieve_sub_agent_context(self):
        """Test retrieving sub-agent context"""
        created = self.manager.create_sub_agent_context(
            parent_task_id="task-789",
            parent_project="admin",
            parent_description="Retrieve test",
        )

        retrieved = self.manager.get_sub_agent_context(created.sub_agent_id)

        assert retrieved is not None
        assert retrieved.sub_agent_id == created.sub_agent_id
        assert retrieved.parent_task_id == "task-789"


class TestSiblingDiscovery:
    """Test sibling agent discovery and awareness"""

    def setup_method(self):
        """Setup test fixtures"""
        self.temp_dir = tempfile.TemporaryDirectory()
        self.manager = SubAgentContextManager(Path(self.temp_dir.name))

    def teardown_method(self):
        """Cleanup"""
        self.temp_dir.cleanup()

    def test_single_sub_agent_no_siblings(self):
        """Test first sub-agent has no siblings"""
        context = self.manager.create_sub_agent_context(
            parent_task_id="parent-1",
            parent_project="admin",
            parent_description="First agent",
        )

        assert len(context.sibling_agents) == 0

    def test_multiple_sub_agents_discover_siblings(self):
        """Test multiple sub-agents discover each other as siblings"""
        # Create first sub-agent
        agent1 = self.manager.create_sub_agent_context(
            parent_task_id="parent-2",
            parent_project="admin",
            parent_description="Agent 1",
        )

        # Create second sub-agent for same parent
        agent2 = self.manager.create_sub_agent_context(
            parent_task_id="parent-2",
            parent_project="admin",
            parent_description="Agent 2",
        )

        # Create third sub-agent for same parent
        agent3 = self.manager.create_sub_agent_context(
            parent_task_id="parent-2",
            parent_project="admin",
            parent_description="Agent 3",
        )

        # Verify sibling relationships
        assert agent2.sub_agent_id in self.manager.get_sibling_agents(agent1.sub_agent_id)
        assert agent3.sub_agent_id in self.manager.get_sibling_agents(agent1.sub_agent_id)
        assert len(self.manager.get_sibling_agents(agent1.sub_agent_id)) == 2

        assert agent1.sub_agent_id in self.manager.get_sibling_agents(agent2.sub_agent_id)
        assert agent3.sub_agent_id in self.manager.get_sibling_agents(agent2.sub_agent_id)
        assert len(self.manager.get_sibling_agents(agent2.sub_agent_id)) == 2

    def test_agents_from_different_parents_not_siblings(self):
        """Test agents from different parents are not siblings"""
        agent1 = self.manager.create_sub_agent_context(
            parent_task_id="parent-a",
            parent_project="admin",
            parent_description="Agent 1",
        )

        agent2 = self.manager.create_sub_agent_context(
            parent_task_id="parent-b",
            parent_project="admin",
            parent_description="Agent 2",
        )

        assert agent2.sub_agent_id not in self.manager.get_sibling_agents(agent1.sub_agent_id)
        assert agent1.sub_agent_id not in self.manager.get_sibling_agents(agent2.sub_agent_id)


class TestPhaseProgression:
    """Test phase progression tracking"""

    def setup_method(self):
        """Setup test fixtures"""
        self.temp_dir = tempfile.TemporaryDirectory()
        self.manager = SubAgentContextManager(Path(self.temp_dir.name))
        self.context = self.manager.create_sub_agent_context(
            parent_task_id="task-phase",
            parent_project="admin",
            parent_description="Phase test",
        )

    def teardown_method(self):
        """Cleanup"""
        self.temp_dir.cleanup()

    def test_update_phase_status(self):
        """Test updating phase status"""
        success = self.manager.update_phase(
            self.context.sub_agent_id,
            "CONTEXT_PREP",
            "completed",
            output="Context prepared",
        )

        assert success is True

        updated = self.manager.get_sub_agent_context(self.context.sub_agent_id)
        phase = updated.phase_progression[0]
        assert phase.status == "completed"
        assert phase.output == "Context prepared"

    def test_get_current_phase(self):
        """Test getting current active phase"""
        # Initially should be first pending phase
        current = self.manager.get_current_phase(self.context.sub_agent_id)
        assert current == "CONTEXT_PREP"

        # Mark first phase as complete
        self.manager.update_phase(
            self.context.sub_agent_id,
            "CONTEXT_PREP",
            "completed",
        )

        # Now should be next pending phase
        current = self.manager.get_current_phase(self.context.sub_agent_id)
        assert current == "RECEIVED"

    def test_phase_duration_calculation(self):
        """Test duration calculation for completed phases"""
        # Mark phase as in progress
        self.manager.update_phase(
            self.context.sub_agent_id,
            "CONTEXT_PREP",
            "in_progress",
        )

        # Mark as completed
        self.manager.update_phase(
            self.context.sub_agent_id,
            "CONTEXT_PREP",
            "completed",
            output="Done",
        )

        updated = self.manager.get_sub_agent_context(self.context.sub_agent_id)
        phase = updated.phase_progression[0]
        assert phase.duration_seconds is not None
        assert phase.duration_seconds >= 0

    def test_phase_progression_sequence(self):
        """Test progressing through all phases"""
        sub_agent_id = self.context.sub_agent_id
        phases = [p.phase_name for p in self.context.phase_progression]

        for phase_name in phases:
            self.manager.update_phase(
                sub_agent_id,
                phase_name,
                "completed",
                output=f"Completed {phase_name}",
            )

        updated = self.manager.get_sub_agent_context(sub_agent_id)
        all_completed = all(p.status == "completed" for p in updated.phase_progression)
        assert all_completed is True


class TestCoordination:
    """Test sub-agent coordination and messaging"""

    def setup_method(self):
        """Setup test fixtures"""
        self.temp_dir = tempfile.TemporaryDirectory()
        self.manager = SubAgentContextManager(Path(self.temp_dir.name))

        # Create two sibling agents
        self.agent1 = self.manager.create_sub_agent_context(
            parent_task_id="parent-coord",
            parent_project="admin",
            parent_description="Agent 1",
        )
        self.agent2 = self.manager.create_sub_agent_context(
            parent_task_id="parent-coord",
            parent_project="admin",
            parent_description="Agent 2",
        )

    def teardown_method(self):
        """Cleanup"""
        self.temp_dir.cleanup()

    def test_send_message_to_sibling(self):
        """Test sending coordination message to sibling"""
        success = self.manager.send_message_to_sibling(
            self.agent1.sub_agent_id,
            self.agent2.sub_agent_id,
            "request",
            {"type": "need_data", "data_type": "context"},
        )

        assert success is True

    def test_message_appears_in_both_agents(self):
        """Test message is visible to both sender and receiver"""
        self.manager.send_message_to_sibling(
            self.agent1.sub_agent_id,
            self.agent2.sub_agent_id,
            "update",
            {"status": "ready"},
        )

        agent1_updated = self.manager.get_sub_agent_context(self.agent1.sub_agent_id)
        agent2_updated = self.manager.get_sub_agent_context(self.agent2.sub_agent_id)

        assert len(agent1_updated.coordination_messages) == 1
        assert len(agent2_updated.coordination_messages) == 1
        assert agent1_updated.coordination_messages[0]["type"] == "update"
        assert agent2_updated.coordination_messages[0]["type"] == "update"

    def test_cannot_message_non_sibling(self):
        """Test cannot send message to non-sibling agent"""
        # Create agent with different parent
        agent3 = self.manager.create_sub_agent_context(
            parent_task_id="parent-other",
            parent_project="admin",
            parent_description="Agent 3",
        )

        # Try to send message across parent boundary
        success = self.manager.send_message_to_sibling(
            self.agent1.sub_agent_id,
            agent3.sub_agent_id,
            "request",
            {"data": "test"},
        )

        assert success is False


class TestContextPersistence:
    """Test context persistence to disk"""

    def test_context_saved_and_loaded(self):
        """Test contexts are saved to disk and reloaded"""
        with tempfile.TemporaryDirectory() as temp_dir:
            manager1 = SubAgentContextManager(Path(temp_dir))

            # Create context in first manager
            context1 = manager1.create_sub_agent_context(
                parent_task_id="task-persist",
                parent_project="admin",
                parent_description="Persistence test",
            )
            sub_agent_id = context1.sub_agent_id

            # Create new manager pointing to same directory
            manager2 = SubAgentContextManager(Path(temp_dir))

            # Should be able to retrieve context from new manager
            context2 = manager2.get_sub_agent_context(sub_agent_id)

            assert context2 is not None
            assert context2.parent_task_id == "task-persist"
            assert context2.sub_agent_id == sub_agent_id


class TestFlowIntegration:
    """Test flow integration with sub-agent context"""

    def setup_method(self):
        """Setup test fixtures"""
        self.temp_dir = tempfile.TemporaryDirectory()
        self.context_manager = SubAgentContextManager(Path(self.temp_dir.name))
        self.integrator = SubAgentFlowIntegrator(self.context_manager)

    def teardown_method(self):
        """Cleanup"""
        self.temp_dir.cleanup()

    def test_execute_sub_agent_flow(self):
        """Test executing full sub-agent flow"""
        results = self.integrator.execute_sub_agent_flow(
            parent_task_id="task-flow",
            parent_project="admin",
            parent_description="Flow test",
            parent_context={"key": "value"},
        )

        assert results["sub_agent_id"] is not None
        assert "phases" in results
        # Should have results for all 9 phases
        assert len(results["phases"]) == 9

    def test_execute_single_phase(self):
        """Test executing a single phase"""
        context = self.context_manager.create_sub_agent_context(
            parent_task_id="task-single",
            parent_project="admin",
            parent_description="Single phase test",
        )

        result = self.integrator.execute_phase(context.sub_agent_id, "CONTEXT_PREP")

        assert result["status"] == "completed"
        assert "output" in result

    def test_get_sub_agent_progress(self):
        """Test getting progress report"""
        context = self.context_manager.create_sub_agent_context(
            parent_task_id="task-progress",
            parent_project="admin",
            parent_description="Progress test",
        )

        # Execute a phase
        self.integrator.execute_phase(context.sub_agent_id, "CONTEXT_PREP")
        self.integrator.execute_phase(context.sub_agent_id, "RECEIVED")

        progress = self.integrator.get_sub_agent_progress(context.sub_agent_id)

        assert progress["completed_phases"] == 2
        assert progress["in_progress_phases"] == 0
        assert progress["total_phases"] == 9

    def test_coordinate_sequential_sub_agents(self):
        """Test sequential coordination of sub-agents"""
        # Create multiple sub-agents for same parent
        for i in range(3):
            self.context_manager.create_sub_agent_context(
                parent_task_id="task-coord",
                parent_project="admin",
                parent_description=f"Agent {i+1}",
            )

        coordination = self.integrator.coordinate_sub_agents(
            parent_task_id="task-coord",
            coordination_strategy="sequential",
        )

        assert len(coordination["sub_agents"]) == 3
        assert coordination["strategy"] == "sequential"

    def test_collect_sub_agent_results(self):
        """Test collecting results from multiple sub-agents"""
        # Create and execute multiple sub-agents
        for i in range(2):
            context = self.context_manager.create_sub_agent_context(
                parent_task_id="task-collect",
                parent_project="admin",
                parent_description=f"Agent {i+1}",
            )
            self.integrator.execute_phase(context.sub_agent_id, "CONTEXT_PREP")

        results = self.integrator.collect_sub_agent_results("task-collect")

        assert results["sub_agents_total"] == 2
        assert len(results["sub_agents"]) == 2
        assert all("progress" in s for s in results["sub_agents"])


class TestContextSummary:
    """Test context summary generation"""

    def setup_method(self):
        """Setup test fixtures"""
        self.temp_dir = tempfile.TemporaryDirectory()
        self.manager = SubAgentContextManager(Path(self.temp_dir.name))

    def teardown_method(self):
        """Cleanup"""
        self.temp_dir.cleanup()

    def test_get_context_summary(self):
        """Test getting human-readable summary"""
        context = self.manager.create_sub_agent_context(
            parent_task_id="task-summary",
            parent_project="admin",
            parent_description="Summary test",
            parent_tags=["important", "urgent"],
        )

        # Create a sibling
        self.manager.create_sub_agent_context(
            parent_task_id="task-summary",
            parent_project="admin",
            parent_description="Sibling agent",
        )

        summary = self.manager.get_context_summary(context.sub_agent_id)

        assert summary is not None
        assert summary["sub_agent_id"] == context.sub_agent_id
        assert summary["parent_task_id"] == "task-summary"
        assert summary["sibling_count"] == 1
        assert summary["parent_tags"] == ["important", "urgent"]


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
