import pytest
import numpy as np
from datetime import datetime, UTC
from pathlib import Path
import tempfile
import shutil
from multi_swarm.core.rag import Document, SearchResult, RAGSystem

@pytest.fixture
def temp_storage():
    """Temporary directory for RAG storage"""
    temp_dir = tempfile.mkdtemp()
    yield temp_dir
    shutil.rmtree(temp_dir)

class TestDocument:
    def test_document_creation(self):
        """Test basic document creation"""
        doc = Document(content="Test content")
        assert doc.content == "Test content"
        assert doc.metadata == {}
        assert doc.embedding is None
        assert isinstance(doc.created_at, datetime)
        assert doc.id is not None

    def test_document_with_metadata(self):
        """Test document creation with metadata"""
        metadata = {"source": "test", "author": "user"}
        doc = Document(content="Test content", metadata=metadata)
        assert doc.metadata == metadata

    def test_document_serialization(self):
        """Test document to_dict and from_dict methods"""
        original = Document(
            content="Test content",
            metadata={"source": "test"},
            embedding=[0.1, 0.2, 0.3]
        )
        dict_form = original.to_dict()
        reconstructed = Document.from_dict(dict_form)
        
        assert reconstructed.content == original.content
        assert reconstructed.metadata == original.metadata
        assert reconstructed.embedding == original.embedding
        assert reconstructed.id == original.id

class TestSearchResult:
    @pytest.fixture
    def sample_document(self):
        return Document(content="Test content")

    def test_search_result_creation(self, sample_document):
        """Test basic search result creation"""
        result = SearchResult(
            document=sample_document,
            score=0.95,
            metadata={"rank": 1}
        )
        assert result.document == sample_document
        assert result.score == 0.95
        assert result.metadata == {"rank": 1}
        assert result.reflection is None

    def test_search_result_with_reflection(self, sample_document):
        """Test search result with reflection"""
        reflection = "High confidence match with good temporal relevance"
        result = SearchResult(
            document=sample_document,
            score=0.95,
            reflection=reflection
        )
        assert result.reflection == reflection

class TestRAGSystem:
    def test_rag_initialization(self, temp_storage):
        """Test RAG system initialization"""
        rag = RAGSystem(storage_path=temp_storage)
        assert rag.model_name == RAGSystem.DEFAULT_MODEL
        assert rag.storage_path == Path(temp_storage)
        assert rag.dimension > 0
        assert len(rag.documents) == 0

    def test_add_document(self, temp_storage):
        """Test adding a document to RAG system"""
        rag = RAGSystem(storage_path=temp_storage)
        doc = rag.add_document(
            content="Test content",
            metadata={"source": "test"}
        )
        assert doc.id in rag.documents
        assert doc.embedding is not None
        assert len(doc.embedding) == rag.dimension

    def test_batch_document_addition(self, temp_storage):
        """Test adding multiple documents in batch"""
        rag = RAGSystem(storage_path=temp_storage)
        documents = [
            {"content": f"Test content {i}", "metadata": {"index": i}}
            for i in range(3)
        ]
        results = rag.add_documents(documents)
        assert len(results) == 3
        assert all(doc.embedding is not None for doc in results)

    def test_search_functionality(self, temp_storage):
        """Test search functionality"""
        rag = RAGSystem(storage_path=temp_storage)
        # Add some test documents
        docs = [
            "The quick brown fox jumps over the lazy dog",
            "A quick brown dog jumps over the lazy cat",
            "The lazy fox sleeps under the tree"
        ]
        for doc in docs:
            rag.add_document(content=doc)

        # Perform search
        results = rag.search("quick brown animal jumping", k=2)
        assert len(results) <= 2
        assert all(isinstance(r, SearchResult) for r in results)
        assert all(r.score >= 0 for r in results)

    def test_search_with_reflection(self, temp_storage):
        """Test search with reflection capabilities"""
        rag = RAGSystem(storage_path=temp_storage)
        # Add documents with different timestamps
        rag.add_document(
            content="Recent content",
            metadata={"created_at": datetime.now(UTC)}
        )
        rag.add_document(
            content="Old content",
            metadata={"created_at": datetime(2020, 1, 1, tzinfo=UTC)}
        )

        results = rag.search_with_reflection("content", k=2)
        assert len(results) <= 2
        assert all(r.reflection is not None for r in results)
        # Check if reflection mentions temporal aspects
        assert any("temporal" in r.reflection.lower() for r in results)

    def test_state_persistence(self, temp_storage):
        """Test saving and loading RAG state"""
        # Create and populate first instance
        rag1 = RAGSystem(storage_path=temp_storage)
        doc1 = rag1.add_document(content="Test content")
        
        # Create second instance pointing to same storage
        rag2 = RAGSystem(storage_path=temp_storage)
        assert doc1.id in rag2.documents
        assert np.array_equal(
            rag1.documents[doc1.id].embedding,
            rag2.documents[doc1.id].embedding
        )

    def test_document_deletion(self, temp_storage):
        """Test document deletion"""
        rag = RAGSystem(storage_path=temp_storage)
        doc = rag.add_document(content="Test content")
        assert rag.delete_document(doc.id)
        assert doc.id not in rag.documents

    def test_index_rebuilding(self, temp_storage):
        """Test index rebuilding"""
        rag = RAGSystem(storage_path=temp_storage)
        # Add documents
        docs = ["First document", "Second document", "Third document"]
        added_docs = [rag.add_document(content=doc) for doc in docs]
        
        # Rebuild index
        rag.rebuild_index()
        
        # Verify search still works
        results = rag.search("document")
        assert len(results) > 0 