import pytest
import tempfile
import shutil
from pathlib import Path
from datetime import datetime, timezone
from multi_swarm.core import Agency, Agent, RAGAgent, Document

@pytest.fixture
def temp_workspace():
    """Create a temporary workspace for testing"""
    temp_dir = tempfile.mkdtemp()
    yield temp_dir
    shutil.rmtree(temp_dir)

@pytest.fixture
def test_agency(temp_workspace):
    """Create a test agency with RAG-enabled agents"""
    # Create storage paths
    storage_path = Path(temp_workspace) / "agency"
    storage_path.mkdir(parents=True)
    
    # Create agents with RAG capabilities
    researcher = RAGAgent(
        name="Researcher",
        description="Research and knowledge gathering specialist",
        storage_path=str(storage_path / "researcher")
    )
    
    analyst = RAGAgent(
        name="Analyst",
        description="Data analysis and synthesis specialist",
        storage_path=str(storage_path / "analyst")
    )
    
    coordinator = RAGAgent(
        name="Coordinator",
        description="Knowledge coordination specialist",
        storage_path=str(storage_path / "coordinator")
    )
    
    # Create agency
    agency = Agency(
        name="ResearchAgency",
        description="Research and analysis agency",
        agents=[researcher, analyst, coordinator],
        flows=[
            (researcher, analyst),
            (analyst, coordinator),
            (coordinator, researcher)
        ],
        storage_path=str(storage_path),
        use_rag=True
    )
    
    return agency

class TestMultiAgentRAG:
    def test_knowledge_distribution(self, test_agency):
        """Test distribution of knowledge across agents"""
        # Add knowledge to shared storage
        docs = test_agency.add_to_shared_knowledge(
            content="Important research findings about AI",
            metadata={"source": "research_paper", "date": "2024-03-01"},
            target_agents=["Researcher", "Analyst"]
        )
        
        # Verify knowledge distribution
        assert len(docs) >= 2  # At least in shared storage and target agents
        # Check if each target agent has the knowledge
        for agent_name in ["Researcher", "Analyst"]:
            agent = test_agency.get_agent(agent_name)
            assert any(
                "Important research findings" in doc.content
                for doc in agent.rag_system.documents.values()
            )

    def test_coordinated_search(self, test_agency):
        """Test coordinated search across multiple agents"""
        # Add different knowledge to different agents
        researcher = test_agency.get_agent("Researcher")
        researcher.rag_system.add_document(
            content="Technical details about neural networks",
            metadata={"type": "technical"}
        )
        
        analyst = test_agency.get_agent("Analyst")
        analyst.rag_system.add_document(
            content="Statistical analysis of AI performance",
            metadata={"type": "analysis"}
        )
        
        # Perform coordinated search
        results = test_agency.search_knowledge_with_coordination(
            query="AI technology research",
            k=2,
            coordinating_agents=["Researcher", "Analyst"]
        )
        
        # Verify results
        assert "primary_results" in results
        assert "collaborative_results" in results
        assert len(results["collaborative_results"]) >= 2
        assert "synthesis" in results
        assert results["synthesis"]["confidence_score"] > 0

    def test_knowledge_synthesis(self, test_agency):
        """Test knowledge synthesis across agents"""
        # Add complementary knowledge to different agents
        test_agency.add_to_shared_knowledge(
            content="Part 1: Introduction to AI",
            metadata={"part": 1},
            target_agents=["Researcher"]
        )
        
        test_agency.add_to_shared_knowledge(
            content="Part 2: Advanced AI Concepts",
            metadata={"part": 2},
            target_agents=["Analyst"]
        )
        
        # Search with coordination
        results = test_agency.search_knowledge_with_coordination(
            query="comprehensive AI guide",
            coordinating_agents=["Researcher", "Analyst"]
        )
        
        # Verify synthesis
        synthesis = results["synthesis"]
        assert synthesis["total_sources"] >= 2
        assert len(synthesis["knowledge_gaps"]) >= 0
        assert len(synthesis["recommendations"]) >= 0

    def test_temporal_relevance(self, test_agency):
        """Test temporal relevance in coordinated search"""
        # Add documents with different timestamps
        old_date = datetime(2020, 1, 1, tzinfo=timezone.utc)
        new_date = datetime.now(timezone.utc)
        
        test_agency.add_to_shared_knowledge(
            content="Old AI research",
            metadata={"created_at": old_date},
            target_agents=["Researcher"]
        )
        
        test_agency.add_to_shared_knowledge(
            content="Latest AI developments",
            metadata={"created_at": new_date},
            target_agents=["Analyst"]
        )
        
        # Search with coordination
        results = test_agency.search_knowledge_with_coordination(
            query="AI research timeline",
            coordinating_agents=["Researcher", "Analyst"]
        )
        
        # Verify temporal awareness
        assert any(
            "temporal" in gap.lower()
            for gap in results["synthesis"]["knowledge_gaps"]
        )

    def test_confidence_scoring(self, test_agency):
        """Test confidence scoring in multi-agent search"""
        # Add high-confidence document
        test_agency.add_to_shared_knowledge(
            content="Verified AI research findings",
            metadata={"confidence": "high"},
            target_agents=["Researcher"]
        )
        
        # Add low-confidence document
        test_agency.add_to_shared_knowledge(
            content="Preliminary AI experiments",
            metadata={"confidence": "low"},
            target_agents=["Analyst"]
        )
        
        # Search with coordination
        results = test_agency.search_knowledge_with_coordination(
            query="AI research quality"
        )
        
        # Verify confidence scoring
        assert "confidence_score" in results["synthesis"]
        assert results["synthesis"]["confidence_score"] > 0

    def test_knowledge_gaps(self, test_agency):
        """Test knowledge gap identification"""
        # Add partial knowledge
        test_agency.add_to_shared_knowledge(
            content="Introduction to AI",
            metadata={"section": "intro"},
            target_agents=["Researcher"]
        )
        
        # Search for comprehensive information
        results = test_agency.search_knowledge_with_coordination(
            query="complete AI guide"
        )
        
        # Verify gap identification
        gaps = results["synthesis"]["knowledge_gaps"]
        assert len(gaps) > 0
        recommendations = results["synthesis"]["recommendations"]
        assert len(recommendations) > 0 