"""
Core RAG (Retrieval-Augmented Generation) system for Multi-Swarm Framework.
"""

import uuid
from datetime import datetime, UTC
from pathlib import Path
from typing import Dict, Any, List, Optional, Union
import numpy as np
from sentence_transformers import SentenceTransformer
import faiss
import json
from pydantic import BaseModel, Field

class DateTimeEncoder(json.JSONEncoder):
    """Custom JSON encoder for datetime objects."""
    def default(self, obj):
        if isinstance(obj, datetime):
            return obj.isoformat()
        return super().default(obj)

class Document(BaseModel):
    """
    Represents a document in the RAG system.
    
    Attributes:
        id: Unique identifier for the document
        content: The text content of the document
        metadata: Additional metadata about the document
        embedding: Pre-computed embedding vector (optional)
        created_at: Timestamp when the document was created
    """
    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
    content: str
    metadata: Optional[Dict[str, Any]] = Field(default_factory=dict)
    embedding: Optional[List[float]] = None
    created_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert document to dictionary format."""
        return {
            "id": self.id,
            "content": self.content,
            "metadata": self.metadata,
            "embedding": self.embedding,
            "created_at": self.created_at.isoformat()
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "Document":
        """Create document from dictionary format."""
        if isinstance(data["created_at"], str):
            data["created_at"] = datetime.fromisoformat(data["created_at"])
        return cls(**data)

class SearchResult(BaseModel):
    """
    Represents a search result from the RAG system.
    
    Attributes:
        document: The matched document
        score: Similarity score of the match
        metadata: Additional metadata about the match
        reflection: Optional reflection on the relevance and quality of the match
    """
    document: Document
    score: float
    metadata: Dict = Field(default_factory=dict)
    reflection: Optional[str] = None

class RAGSystem:
    """
    Advanced Retrieval-Augmented Generation system with multiple embedding models
    and sophisticated retrieval strategies.
    """
    
    DEFAULT_MODEL = "all-MiniLM-L6-v2"
    
    def __init__(
        self,
        model_name: str = DEFAULT_MODEL,
        storage_path: Optional[Union[str, Path]] = None,
        device: str = None,
        use_gpu: bool = None
    ):
        """Initialize the RAG system."""
        self.model_name = model_name
        self.storage_path = Path(storage_path) if storage_path else None
        
        # Initialize embedding model
        if device is None:
            device = "cuda" if use_gpu else "cpu"
        self.model = SentenceTransformer(model_name, device=device)
        self.dimension = self.model.get_sentence_embedding_dimension()
        
        # Initialize FAISS index
        self.index = faiss.IndexFlatL2(self.dimension)
        
        # Document storage
        self.documents: Dict[str, Document] = {}
        
        # Load existing state if available
        if self.storage_path and self.storage_path.exists():
            self._load_state()
    
    def _compute_embedding(self, text: str) -> np.ndarray:
        """Compute embedding for text."""
        return self.model.encode(text, convert_to_numpy=True)
    
    def _save_state(self):
        """Save system state to disk."""
        if not self.storage_path:
            return
            
        self.storage_path.mkdir(parents=True, exist_ok=True)
        
        # Save documents
        docs_path = self.storage_path / "documents.json"
        with open(docs_path, "w") as f:
            json.dump(
                {id: doc.to_dict() for id, doc in self.documents.items()},
                f,
                cls=DateTimeEncoder
            )
        
        # Save FAISS index
        index_path = self.storage_path / "index.faiss"
        faiss.write_index(self.index, str(index_path))
    
    def _load_state(self):
        """Load system state from disk."""
        # Load documents
        docs_path = self.storage_path / "documents.json"
        if docs_path.exists():
            with open(docs_path, "r") as f:
                docs_data = json.load(f)
                self.documents = {
                    id: Document.from_dict(data)
                    for id, data in docs_data.items()
                }
        
        # Load FAISS index
        index_path = self.storage_path / "index.faiss"
        if index_path.exists():
            self.index = faiss.read_index(str(index_path))
    
    def add_document(
        self,
        content: str,
        metadata: Dict = None,
        compute_embedding: bool = True
    ) -> Document:
        """Add a document to the system."""
        # Create document
        doc = Document(
            content=content,
            metadata=metadata or {}
        )
        
        # Compute embedding if requested
        if compute_embedding:
            embedding = self._compute_embedding(content)
            doc.embedding = embedding.tolist()
            self.index.add(embedding.reshape(1, -1))
        
        # Store document
        self.documents[doc.id] = doc
        
        # Save state
        if self.storage_path:
            self._save_state()
        
        return doc
    
    def add_documents(
        self,
        documents: List[Union[str, Dict]],
        batch_size: int = 32
    ) -> List[Document]:
        """Add multiple documents in batches."""
        results = []
        for i in range(0, len(documents), batch_size):
            batch = documents[i:i + batch_size]
            batch_docs = self._process_batch(batch)
            results.extend(batch_docs)
        
        # Save state
        if self.storage_path:
            self._save_state()
        
        return results
    
    def _process_batch(self, batch: List[Dict]) -> List[Document]:
        """Process a batch of documents."""
        # Prepare documents
        docs = []
        texts = []
        for item in batch:
            if isinstance(item, str):
                doc = Document(content=item)
            else:
                doc = Document(
                    content=item["content"],
                    metadata=item.get("metadata", {})
                )
            docs.append(doc)
            texts.append(doc.content)
        
        # Compute embeddings
        embeddings = self.model.encode(texts, convert_to_numpy=True)
        
        # Store documents and embeddings
        for doc, embedding in zip(docs, embeddings):
            doc.embedding = embedding.tolist()
            self.documents[doc.id] = doc
        
        # Add to index
        self.index.add(embeddings)
        
        return docs
    
    def search(
        self,
        query: str,
        k: int = 5,
        threshold: float = None,
        filter_fn: callable = None
    ) -> List[SearchResult]:
        """Search for relevant documents."""
        # Compute query embedding
        query_embedding = self._compute_embedding(query)
        
        # Search index
        scores, indices = self.index.search(
            query_embedding.reshape(1, -1),
            k
        )
        
        # Get document IDs
        doc_ids = list(self.documents.keys())
        
        # Create results
        results = []
        for score, idx in zip(scores[0], indices[0]):
            if idx < 0 or idx >= len(doc_ids):
                continue
                
            doc = self.documents[doc_ids[idx]]
            
            # Apply threshold
            if threshold is not None and score > threshold:
                continue
                
            # Apply filter
            if filter_fn and not filter_fn(doc):
                continue
            
            result = SearchResult(
                document=doc,
                score=float(score),
                metadata={"rank": len(results) + 1}
            )
            results.append(result)
        
        return results
    
    def get_document(self, doc_id: str) -> Optional[Document]:
        """Get a document by ID."""
        return self.documents.get(doc_id)
    
    def delete_document(self, doc_id: str) -> bool:
        """Delete a document."""
        if doc_id not in self.documents:
            return False
            
        # Remove from storage
        doc = self.documents.pop(doc_id)
        
        # Rebuild index
        self.rebuild_index()
        
        # Save state
        if self.storage_path:
            self._save_state()
        
        return True
    
    def rebuild_index(self):
        """Rebuild the FAISS index."""
        # Clear index
        self.index = faiss.IndexFlatL2(self.dimension)
        
        # Add all embeddings
        embeddings = []
        for doc in self.documents.values():
            if doc.embedding:
                embeddings.append(doc.embedding)
        
        if embeddings:
            self.index.add(np.array(embeddings))
    
    def clear(self):
        """Clear all documents and reset the index."""
        self.documents.clear()
        self.index = faiss.IndexFlatL2(self.dimension)
        
        if self.storage_path:
            self._save_state()
    
    def search_with_reflection(self, query: str, k: int = 5) -> List[SearchResult]:
        """Search with added reflection on result quality."""
        # Perform basic search
        results = self.search(query, k=k)
        
        # Add reflection to each result
        for result in results:
            result.reflection = self._generate_reflection(query, result.document)
        
        return results
    
    def _generate_reflection(self, query: str, result: Document) -> str:
        """Generate reflection on search result quality."""
        # Calculate temporal relevance
        time_diff = datetime.now(UTC) - result.created_at
        days_old = time_diff.days
        
        # Calculate content relevance
        relevance_factors = []
        
        # Check temporal relevance
        if days_old < 30:
            relevance_factors.append("high temporal relevance (recent)")
        elif days_old < 180:
            relevance_factors.append("moderate temporal relevance")
        else:
            relevance_factors.append("low temporal relevance (old)")
        
        # Check metadata quality
        if result.metadata:
            relevance_factors.append("includes additional context")
        
        # Check content length
        if len(result.content) > 1000:
            relevance_factors.append("comprehensive content")
        elif len(result.content) < 100:
            relevance_factors.append("brief content")
        
        return f"Match quality assessment: {', '.join(relevance_factors)}"

class RAGTool(BaseModel):
    """
    A tool that provides RAG capabilities for agents.
    
    This tool can be used to search through a knowledge base and retrieve
    relevant information for generating responses.
    """
    query: str = Field(..., description="The search query")
    k: int = Field(default=5, description="Number of results to return")
    threshold: Optional[float] = Field(
        default=None,
        description="Optional similarity threshold"
    )
    
    def run(self) -> List[SearchResult]:
        """
        Search the knowledge base and return relevant results.
        
        Returns:
            List of search results
        """
        rag = RAGSystem()
        return rag.search(
            self.query,
            self.k,
            self.threshold
        ) 