"""RAG Pipeline for Context-Aware Agent Operations

Retrieval-Augmented Generation pipeline that enriches agent responses
with relevant context from docs/assets/test data corpus.
"""

import os
import json
from pathlib import Path
from typing import List, Dict, Any, Optional, Tuple
import numpy as np
from dataclasses import dataclass

try:
    from sentence_transformers import SentenceTransformer
    EMBEDDINGS_AVAILABLE = True
except ImportError:
    EMBEDDINGS_AVAILABLE = False
    print("Warning: sentence-transformers not available. Install with: pip install sentence-transformers")


@dataclass
class Document:
    """Document container with metadata"""
    content: str
    source: str
    metadata: Dict[str, Any]
    embedding: Optional[np.ndarray] = None


@dataclass
class RetrievalResult:
    """Result from RAG retrieval with context and citations"""
    documents: List[Document]
    context: str
    citations: List[str]
    scores: List[float]


class RAGPipeline:
    """Retrieval-Augmented Generation Pipeline
    
    Provides context-aware enrichment for agent operations by:
    - Indexing documents from docs/assets/test data
    - Retrieving relevant context for queries
    - Generating citations and evidence for responses
    """
    
    def __init__(
        self,
        corpus_dir: str = "./docs",
        model_name: str = "all-MiniLM-L6-v2",
        top_k: int = 5,
        cache_dir: str = "./rag_cache"
    ):
        """Initialize RAG pipeline
        
        Args:
            corpus_dir: Directory containing documents to index
            model_name: Sentence transformer model for embeddings
            top_k: Number of top documents to retrieve
            cache_dir: Directory for caching embeddings
        """
        self.corpus_dir = Path(corpus_dir)
        self.top_k = top_k
        self.cache_dir = Path(cache_dir)
        self.cache_dir.mkdir(exist_ok=True)
        
        # Initialize embedding model
        if EMBEDDINGS_AVAILABLE:
            try:
                self.model = SentenceTransformer(model_name)
                self.embedding_dim = self.model.get_sentence_embedding_dimension()
            except Exception as e:
                print(f"Error loading embedding model: {e}")
                self.model = None
        else:
            self.model = None
        
        # Document store
        self.documents: List[Document] = []
        self.embeddings: Optional[np.ndarray] = None
        
        # Load and index documents
        self._load_corpus()
    
    def _load_corpus(self):
        """Load documents from corpus directory"""
        if not self.corpus_dir.exists():
            print(f"Warning: Corpus directory {self.corpus_dir} does not exist")
            return
        
        # Load all text-based files
        extensions = [".txt", ".md", ".json", ".csv", ".py"]
        
        for ext in extensions:
            for file_path in self.corpus_dir.rglob(f"*{ext}"):
                try:
                    self._load_document(file_path)
                except Exception as e:
                    print(f"Error loading {file_path}: {e}")
        
        print(f"Loaded {len(self.documents)} documents from corpus")
        
        # Generate embeddings
        if self.model and self.documents:
            self._generate_embeddings()
    
    def _load_document(self, file_path: Path):
        """Load a single document"""
        with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
            content = f.read()
        
        # Skip empty files
        if not content.strip():
            return
        
        # Create document
        doc = Document(
            content=content,
            source=str(file_path.relative_to(self.corpus_dir)),
            metadata={
                "file_type": file_path.suffix,
                "file_name": file_path.name,
                "path": str(file_path)
            }
        )
        
        self.documents.append(doc)
    
    def _generate_embeddings(self):
        """Generate embeddings for all documents"""
        if not self.model:
            return
        
        # Check cache
        cache_file = self.cache_dir / "embeddings.npy"
        cache_meta = self.cache_dir / "embeddings_meta.json"
        
        if cache_file.exists() and cache_meta.exists():
            try:
                self.embeddings = np.load(cache_file)
                with open(cache_meta, 'r') as f:
                    meta = json.load(f)
                if meta.get("num_docs") == len(self.documents):
                    print("Loaded embeddings from cache")
                    return
            except Exception as e:
                print(f"Error loading cache: {e}")
        
        # Generate new embeddings
        print("Generating document embeddings...")
        texts = [doc.content for doc in self.documents]
        
        try:
            self.embeddings = self.model.encode(
                texts,
                show_progress_bar=True,
                convert_to_numpy=True
            )
            
            # Cache embeddings
            np.save(cache_file, self.embeddings)
            with open(cache_meta, 'w') as f:
                json.dump({"num_docs": len(self.documents)}, f)
            
            print("Embeddings generated and cached")
        except Exception as e:
            print(f"Error generating embeddings: {e}")
            self.embeddings = None
    
    def retrieve(
        self,
        query: str,
        top_k: Optional[int] = None,
        filter_metadata: Optional[Dict[str, Any]] = None
    ) -> RetrievalResult:
        """Retrieve relevant documents for a query
        
        Args:
            query: Search query
            top_k: Number of documents to retrieve (default: self.top_k)
            filter_metadata: Optional metadata filters
            
        Returns:
            RetrievalResult with documents, context, and citations
        """
        if top_k is None:
            top_k = self.top_k
        
        if not self.documents:
            return RetrievalResult(
                documents=[],
                context="No documents available in corpus.",
                citations=[],
                scores=[]
            )
        
        # Use semantic search if embeddings available
        if self.model and self.embeddings is not None:
            return self._semantic_retrieve(query, top_k, filter_metadata)
        else:
            # Fallback to keyword search
            return self._keyword_retrieve(query, top_k, filter_metadata)
    
    def _semantic_retrieve(
        self,
        query: str,
        top_k: int,
        filter_metadata: Optional[Dict[str, Any]]
    ) -> RetrievalResult:
        """Semantic retrieval using embeddings"""
        # Encode query
        query_embedding = self.model.encode([query], convert_to_numpy=True)[0]
        
        # Compute similarities
        similarities = np.dot(self.embeddings, query_embedding)
        
        # Apply metadata filters
        valid_indices = range(len(self.documents))
        if filter_metadata:
            valid_indices = [
                i for i in valid_indices
                if all(
                    self.documents[i].metadata.get(k) == v
                    for k, v in filter_metadata.items()
                )
            ]
        
        # Get top-k indices
        if valid_indices:
            filtered_sims = [(i, similarities[i]) for i in valid_indices]
            top_indices = sorted(filtered_sims, key=lambda x: x[1], reverse=True)[:top_k]
            indices, scores = zip(*top_indices) if top_indices else ([], [])
        else:
            indices, scores = [], []
        
        # Build result
        retrieved_docs = [self.documents[i] for i in indices]
        
        return self._build_result(retrieved_docs, list(scores))
    
    def _keyword_retrieve(
        self,
        query: str,
        top_k: int,
        filter_metadata: Optional[Dict[str, Any]]
    ) -> RetrievalResult:
        """Simple keyword-based retrieval fallback"""
        query_terms = set(query.lower().split())
        
        # Score documents by keyword overlap
        scores = []
        for doc in self.documents:
            doc_terms = set(doc.content.lower().split())
            overlap = len(query_terms & doc_terms)
            scores.append(overlap)
        
        # Apply metadata filters
        valid_indices = range(len(self.documents))
        if filter_metadata:
            valid_indices = [
                i for i in valid_indices
                if all(
                    self.documents[i].metadata.get(k) == v
                    for k, v in filter_metadata.items()
                )
            ]
        
        # Get top-k
        if valid_indices:
            filtered_scores = [(i, scores[i]) for i in valid_indices]
            top_indices = sorted(filtered_scores, key=lambda x: x[1], reverse=True)[:top_k]
            indices, top_scores = zip(*top_indices) if top_indices else ([], [])
        else:
            indices, top_scores = [], []
        
        retrieved_docs = [self.documents[i] for i in indices]
        
        return self._build_result(retrieved_docs, list(top_scores))
    
    def _build_result(self, documents: List[Document], scores: List[float]) -> RetrievalResult:
        """Build retrieval result with context and citations"""
        # Build context string
        context_parts = []
        citations = []
        
        for i, (doc, score) in enumerate(zip(documents, scores), 1):
            # Truncate long documents
            content_preview = doc.content[:500]
            if len(doc.content) > 500:
                content_preview += "..."
            
            context_parts.append(
                f"[{i}] {doc.source}\n{content_preview}\n"
            )
            citations.append(f"[{i}] {doc.source} (relevance: {score:.2f})")
        
        context = "\n".join(context_parts) if context_parts else "No relevant context found."
        
        return RetrievalResult(
            documents=documents,
            context=context,
            citations=citations,
            scores=scores
        )
    
    def enrich_query(self, query: str, max_context_length: int = 2000) -> Dict[str, Any]:
        """Enrich a query with retrieved context
        
        Args:
            query: User query
            max_context_length: Maximum context length to include
            
        Returns:
            Dict with enriched_query, context, and citations
        """
        result = self.retrieve(query)
        
        # Truncate context if needed
        context = result.context
        if len(context) > max_context_length:
            context = context[:max_context_length] + "..."
        
        enriched_query = f"""{query}

Relevant Context:
{context}

Please use the above context to provide an informed response and cite sources where applicable."""
        
        return {
            "enriched_query": enriched_query,
            "context": context,
            "citations": result.citations,
            "num_sources": len(result.documents)
        }


# Singleton instance
_rag_pipeline = None

def get_rag_pipeline(
    corpus_dir: str = "./docs",
    force_reload: bool = False
) -> RAGPipeline:
    """Get or create RAG pipeline singleton
    
    Args:
        corpus_dir: Directory containing documents
        force_reload: Force reload of the pipeline
        
    Returns:
        RAGPipeline instance
    """
    global _rag_pipeline
    
    if _rag_pipeline is None or force_reload:
        _rag_pipeline = RAGPipeline(corpus_dir=corpus_dir)
    
    return _rag_pipeline