import pytest
import tempfile
import shutil
import os
from pathlib import Path
from datetime import datetime
import json
from multi_swarm.core import Agency, RAGAgent

@pytest.fixture
def research_workspace():
    """Create a temporary research workspace with sample documents"""
    temp_dir = tempfile.mkdtemp()
    
    # Create sample documents
    docs_dir = Path(temp_dir) / "documents"
    docs_dir.mkdir()
    
    # Technical documentation
    with open(docs_dir / "technical_spec.md", "w") as f:
        f.write("""# AI System Technical Specification
        
## Architecture
- Neural network based
- Distributed processing
- Real-time inference
        
## Performance
- 95% accuracy on benchmark
- 10ms latency
- 1000 requests per second""")
    
    # Research paper
    with open(docs_dir / "research_paper.md", "w") as f:
        f.write("""# Advanced AI Research Findings
        
## Abstract
Novel approaches to distributed AI systems show promising results.
        
## Methods
Using multi-agent coordination with RAG capabilities.
        
## Results
30% improvement in knowledge synthesis.""")
    
    # Performance data
    with open(docs_dir / "performance_data.json", "w") as f:
        json.dump({
            "accuracy": 0.95,
            "latency": 10,
            "throughput": 1000,
            "timestamp": datetime.utcnow().isoformat()
        }, f)
    
    yield temp_dir
    shutil.rmtree(temp_dir)

@pytest.fixture
def research_agency(research_workspace):
    """Create a research agency with specialized agents"""
    # Create storage paths
    storage_path = Path(research_workspace) / "agency"
    storage_path.mkdir()
    
    # Create specialized agents
    tech_researcher = RAGAgent(
        name="TechResearcher",
        description="Technical research and documentation specialist",
        storage_path=str(storage_path / "tech_researcher")
    )
    
    data_analyst = RAGAgent(
        name="DataAnalyst",
        description="Data analysis and performance evaluation specialist",
        storage_path=str(storage_path / "data_analyst")
    )
    
    knowledge_coordinator = RAGAgent(
        name="KnowledgeCoordinator",
        description="Knowledge synthesis and gap analysis specialist",
        storage_path=str(storage_path / "coordinator")
    )
    
    # Create agency
    agency = Agency(
        name="AIResearchAgency",
        description="AI research and development agency",
        agents=[tech_researcher, data_analyst, knowledge_coordinator],
        flows=[
            (tech_researcher, data_analyst),
            (data_analyst, knowledge_coordinator),
            (knowledge_coordinator, tech_researcher)
        ],
        storage_path=str(storage_path),
        use_rag=True,
        use_file_storage=True
    )
    
    return agency

class TestRAGWorkflow:
    def test_complete_research_workflow(self, research_agency, research_workspace):
        """Test complete research workflow from document ingestion to synthesis"""
        # Step 1: Import documents
        docs_dir = Path(research_workspace) / "documents"
        
        # Technical docs to tech researcher
        with open(docs_dir / "technical_spec.md") as f:
            tech_content = f.read()
            research_agency.add_to_shared_knowledge(
                content=tech_content,
                metadata={
                    "type": "technical",
                    "source": "technical_spec.md",
                    "created_at": datetime.utcnow()
                },
                target_agents=["TechResearcher"]
            )
        
        # Research paper to all agents
        with open(docs_dir / "research_paper.md") as f:
            research_content = f.read()
            research_agency.add_to_shared_knowledge(
                content=research_content,
                metadata={
                    "type": "research",
                    "source": "research_paper.md",
                    "created_at": datetime.utcnow()
                }
            )
        
        # Performance data to data analyst
        with open(docs_dir / "performance_data.json") as f:
            perf_content = f.read()
            research_agency.add_to_shared_knowledge(
                content=perf_content,
                metadata={
                    "type": "performance",
                    "source": "performance_data.json",
                    "created_at": datetime.utcnow()
                },
                target_agents=["DataAnalyst"]
            )
        
        # Step 2: Verify document distribution
        tech_researcher = research_agency.get_agent("TechResearcher")
        assert any(
            "AI System Technical Specification" in doc.content
            for doc in tech_researcher.rag_system.documents.values()
        )
        
        data_analyst = research_agency.get_agent("DataAnalyst")
        assert any(
            "accuracy" in doc.content
            for doc in data_analyst.rag_system.documents.values()
        )
        
        # Step 3: Perform coordinated research
        results = research_agency.search_knowledge_with_coordination(
            query="AI system performance and architecture",
            coordinating_agents=["TechResearcher", "DataAnalyst"]
        )
        
        # Verify comprehensive results
        assert results["synthesis"]["total_sources"] >= 3
        assert results["synthesis"]["confidence_score"] > 0.7
        
        # Verify technical details
        tech_results = [
            r for r in results["primary_results"]
            if "technical" in r.get("metadata", {}).get("type", "").lower()
        ]
        assert len(tech_results) > 0
        assert any("architecture" in r["content"].lower() for r in tech_results)
        
        # Verify performance data
        perf_results = [
            r for r in results["collaborative_results"].get("DataAnalyst", [])
            if "performance" in r.get("metadata", {}).get("type", "").lower()
        ]
        assert len(perf_results) > 0
        assert any("accuracy" in r["content"].lower() for r in perf_results)
        
        # Step 4: Knowledge gap analysis
        gaps = results["synthesis"]["knowledge_gaps"]
        recommendations = results["synthesis"]["recommendations"]
        
        # Verify gap identification
        assert len(gaps) > 0
        assert len(recommendations) > 0
        
        # Step 5: Iterative improvement
        # Add missing information based on gaps
        if any("implementation" in gap.lower() for gap in gaps):
            research_agency.add_to_shared_knowledge(
                content="Implementation details: The system uses PyTorch for ML operations",
                metadata={
                    "type": "implementation",
                    "source": "implementation_guide.md",
                    "created_at": datetime.utcnow()
                }
            )
        
        # Verify improved results
        updated_results = research_agency.search_knowledge_with_coordination(
            query="AI system complete overview",
            coordinating_agents=["TechResearcher", "DataAnalyst", "KnowledgeCoordinator"]
        )
        
        assert updated_results["synthesis"]["confidence_score"] > results["synthesis"]["confidence_score"]
        assert len(updated_results["synthesis"]["knowledge_gaps"]) < len(results["synthesis"]["knowledge_gaps"])

    def test_temporal_knowledge_evolution(self, research_agency):
        """Test how the system handles evolving knowledge over time"""
        # Add knowledge from different time periods
        old_knowledge = {
            "content": "Initial AI system design with basic capabilities",
            "metadata": {
                "version": "1.0",
                "created_at": datetime(2023, 1, 1)
            }
        }
        
        updated_knowledge = {
            "content": "Enhanced AI system with improved performance",
            "metadata": {
                "version": "2.0",
                "created_at": datetime(2024, 1, 1)
            }
        }
        
        latest_knowledge = {
            "content": "Latest AI system with advanced features",
            "metadata": {
                "version": "3.0",
                "created_at": datetime.utcnow()
            }
        }
        
        # Add knowledge chronologically
        for knowledge in [old_knowledge, updated_knowledge, latest_knowledge]:
            research_agency.add_to_shared_knowledge(
                content=knowledge["content"],
                metadata=knowledge["metadata"]
            )
        
        # Search with temporal awareness
        results = research_agency.search_knowledge_with_coordination(
            query="AI system evolution"
        )
        
        # Verify temporal analysis
        assert "temporal_analysis" in str(results["synthesis"])
        assert any(
            "latest version" in r["reflection"].lower()
            for r in results["primary_results"]
            if "reflection" in r
        )
        
        # Verify version tracking
        versions = [
            r["metadata"].get("version")
            for r in results["primary_results"]
            if "version" in r.get("metadata", {})
        ]
        assert "3.0" in versions  # Latest version should be included
        
    def test_cross_agent_knowledge_validation(self, research_agency):
        """Test how agents validate and cross-reference knowledge"""
        # Add potentially conflicting information to different agents
        research_agency.add_to_shared_knowledge(
            content="System achieves 95% accuracy",
            metadata={"type": "performance", "confidence": "high"},
            target_agents=["TechResearcher"]
        )
        
        research_agency.add_to_shared_knowledge(
            content="System achieves 85% accuracy",
            metadata={"type": "performance", "confidence": "medium"},
            target_agents=["DataAnalyst"]
        )
        
        # Search for performance metrics
        results = research_agency.search_knowledge_with_coordination(
            query="system accuracy metrics",
            coordinating_agents=["TechResearcher", "DataAnalyst"]
        )
        
        # Verify conflict identification
        assert any(
            "conflicting" in gap.lower()
            for gap in results["synthesis"]["knowledge_gaps"]
        )
        
        # Verify confidence-based resolution
        high_confidence_results = [
            r for r in results["primary_results"]
            if r.get("metadata", {}).get("confidence") == "high"
        ]
        assert len(high_confidence_results) > 0
        
        # Verify recommendations for resolution
        assert any(
            "validate" in rec.lower()
            for rec in results["synthesis"]["recommendations"]
        ) 