"""
Retrieval service for RAG pipeline.

Supports:
- Dense retrieval (embedding-based)
- Sparse retrieval (BM25)
- Hybrid retrieval (combination)
- Reranking
"""

import re
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple, Union

from .chunking import Chunk, ChunkType
from .indexing import EmbeddingModel
from ..infra.vectordb import VectorStore, SearchResult
from ..infra.logging import get_logger
from ..infra.tracing import trace_function

logger = get_logger(__name__)


@dataclass
class RetrievalResult:
    """Result from retrieval with relevance score."""
    content: str
    score: float
    chunk_type: ChunkType
    metadata: Dict[str, Any] = field(default_factory=dict)
    
    # Source information
    document_id: Optional[str] = None
    section_hierarchy: List[str] = field(default_factory=list)
    page_numbers: List[int] = field(default_factory=list)
    
    @classmethod
    def from_search_result(cls, result: SearchResult) -> "RetrievalResult":
        """Create from vector store search result."""
        metadata = result.metadata.copy()
        
        # Parse chunk type
        chunk_type_str = metadata.pop("chunk_type", "mixed")
        try:
            chunk_type = ChunkType(chunk_type_str)
        except ValueError:
            chunk_type = ChunkType.MIXED
        
        # Parse section hierarchy
        section_str = metadata.pop("section_hierarchy", "")
        section_hierarchy = section_str.split("|") if section_str else []
        
        # Parse page numbers
        pages_str = metadata.pop("page_numbers", "")
        page_numbers = [int(p) for p in pages_str.split(",") if p.strip().isdigit()]
        
        document_id = metadata.pop("document_id", None)
        
        return cls(
            content=result.content,
            score=result.score,
            chunk_type=chunk_type,
            metadata=metadata,
            document_id=document_id,
            section_hierarchy=section_hierarchy,
            page_numbers=page_numbers
        )


class Retriever(ABC):
    """Abstract base class for retrievers."""
    
    @abstractmethod
    def retrieve(
        self,
        query: str,
        top_k: int = 10,
        filter_metadata: Optional[Dict[str, Any]] = None
    ) -> List[RetrievalResult]:
        """
        Retrieve relevant documents for a query.
        
        Args:
            query: Search query
            top_k: Number of results to return
            filter_metadata: Optional metadata filters
            
        Returns:
            List of retrieval results
        """
        pass


class DenseRetriever(Retriever):
    """Dense retriever using embeddings."""
    
    def __init__(
        self,
        vector_store: VectorStore,
        embedding_model: EmbeddingModel
    ):
        """
        Initialize dense retriever.
        
        Args:
            vector_store: Vector store for search
            embedding_model: Model for query embedding
        """
        self.vector_store = vector_store
        self.embedding_model = embedding_model
    
    @trace_function(name="dense_retrieve")
    def retrieve(
        self,
        query: str,
        top_k: int = 10,
        filter_metadata: Optional[Dict[str, Any]] = None
    ) -> List[RetrievalResult]:
        """Retrieve using dense embeddings."""
        # Embed query
        query_embedding = self.embedding_model.embed(query)
        
        # Search vector store
        results = self.vector_store.search(
            query_embedding=query_embedding,
            top_k=top_k,
            filter_metadata=filter_metadata
        )
        
        # Convert to retrieval results
        return [RetrievalResult.from_search_result(r) for r in results]


class BM25Retriever(Retriever):
    """Sparse retriever using BM25."""
    
    def __init__(
        self,
        documents: Optional[List[Dict[str, Any]]] = None,
        k1: float = 1.5,
        b: float = 0.75
    ):
        """
        Initialize BM25 retriever.
        
        Args:
            documents: List of documents with 'content' and 'metadata'
            k1: BM25 k1 parameter
            b: BM25 b parameter
        """
        self.k1 = k1
        self.b = b
        self._documents = documents or []
        self._bm25 = None
        self._tokenized_corpus = None
        
        if documents:
            self._build_index(documents)
    
    def _build_index(self, documents: List[Dict[str, Any]]) -> None:
        """Build BM25 index from documents."""
        try:
            from rank_bm25 import BM25Okapi
        except ImportError:
            raise ImportError("rank_bm25 is required. Install with: pip install rank-bm25")
        
        self._documents = documents
        self._tokenized_corpus = [
            self._tokenize(doc.get("content", ""))
            for doc in documents
        ]
        
        self._bm25 = BM25Okapi(
            self._tokenized_corpus,
            k1=self.k1,
            b=self.b
        )
        
        logger.info(f"Built BM25 index with {len(documents)} documents")
    
    def _tokenize(self, text: str) -> List[str]:
        """Simple tokenization."""
        # Convert to lowercase and split on non-alphanumeric
        tokens = re.findall(r'\b\w+\b', text.lower())
        return tokens
    
    def add_documents(self, documents: List[Dict[str, Any]]) -> None:
        """Add documents and rebuild index."""
        self._build_index(self._documents + documents)
    
    @trace_function(name="bm25_retrieve")
    def retrieve(
        self,
        query: str,
        top_k: int = 10,
        filter_metadata: Optional[Dict[str, Any]] = None
    ) -> List[RetrievalResult]:
        """Retrieve using BM25."""
        if self._bm25 is None:
            logger.warning("BM25 index not built")
            return []
        
        # Tokenize query
        query_tokens = self._tokenize(query)
        
        # Get BM25 scores
        scores = self._bm25.get_scores(query_tokens)
        
        # Get top-k indices
        top_indices = sorted(
            range(len(scores)),
            key=lambda i: scores[i],
            reverse=True
        )[:top_k]
        
        # Build results
        results = []
        for idx in top_indices:
            if scores[idx] > 0:
                doc = self._documents[idx]
                metadata = doc.get("metadata", {}).copy()
                
                # Parse chunk type
                chunk_type_str = metadata.pop("chunk_type", "mixed")
                try:
                    chunk_type = ChunkType(chunk_type_str)
                except ValueError:
                    chunk_type = ChunkType.MIXED
                
                results.append(RetrievalResult(
                    content=doc.get("content", ""),
                    score=scores[idx],
                    chunk_type=chunk_type,
                    metadata=metadata
                ))
        
        return results


class HybridRetriever(Retriever):
    """
    Hybrid retriever combining dense and sparse retrieval.
    
    Uses reciprocal rank fusion (RRF) to combine results.
    """
    
    def __init__(
        self,
        dense_retriever: DenseRetriever,
        sparse_retriever: BM25Retriever,
        dense_weight: float = 0.5,
        sparse_weight: float = 0.5,
        rrf_k: int = 60
    ):
        """
        Initialize hybrid retriever.
        
        Args:
            dense_retriever: Dense retriever instance
            sparse_retriever: Sparse retriever instance
            dense_weight: Weight for dense results
            sparse_weight: Weight for sparse results
            rrf_k: RRF constant (higher = more emphasis on lower ranks)
        """
        self.dense_retriever = dense_retriever
        self.sparse_retriever = sparse_retriever
        self.dense_weight = dense_weight
        self.sparse_weight = sparse_weight
        self.rrf_k = rrf_k
    
    @trace_function(name="hybrid_retrieve")
    def retrieve(
        self,
        query: str,
        top_k: int = 10,
        filter_metadata: Optional[Dict[str, Any]] = None
    ) -> List[RetrievalResult]:
        """Retrieve using hybrid approach."""
        # Get results from both retrievers
        dense_results = self.dense_retriever.retrieve(query, top_k * 2, filter_metadata)
        sparse_results = self.sparse_retriever.retrieve(query, top_k * 2, filter_metadata)
        
        # Compute RRF scores
        content_to_result: Dict[str, RetrievalResult] = {}
        content_scores: Dict[str, float] = {}
        
        # Process dense results
        for rank, result in enumerate(dense_results):
            content_key = result.content[:500]  # Use first 500 chars as key
            rrf_score = self.dense_weight / (self.rrf_k + rank + 1)
            
            content_scores[content_key] = content_scores.get(content_key, 0) + rrf_score
            if content_key not in content_to_result:
                content_to_result[content_key] = result
        
        # Process sparse results
        for rank, result in enumerate(sparse_results):
            content_key = result.content[:500]
            rrf_score = self.sparse_weight / (self.rrf_k + rank + 1)
            
            content_scores[content_key] = content_scores.get(content_key, 0) + rrf_score
            if content_key not in content_to_result:
                content_to_result[content_key] = result
        
        # Sort by combined score and return top-k
        sorted_keys = sorted(
            content_scores.keys(),
            key=lambda k: content_scores[k],
            reverse=True
        )[:top_k]
        
        results = []
        for key in sorted_keys:
            result = content_to_result[key]
            result.score = content_scores[key]
            results.append(result)
        
        return results


class Reranker:
    """
    Cross-encoder reranker for improving retrieval quality.
    """
    
    def __init__(
        self,
        model_name: str = "cross-encoder/ms-marco-MiniLM-L-6-v2",
        device: str = "auto"
    ):
        """
        Initialize reranker.
        
        Args:
            model_name: Cross-encoder model name
            device: Device to use
        """
        self.model_name = model_name
        self.device = self._resolve_device(device)
        self._model = None
    
    def _resolve_device(self, device: str) -> str:
        """Resolve device string."""
        if device != "auto":
            return device
        
        try:
            import torch
            if torch.cuda.is_available():
                return "cuda"
            elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
                return "mps"
        except ImportError:
            pass
        
        return "cpu"
    
    def _load_model(self):
        """Lazy load the reranker model."""
        if self._model is not None:
            return
        
        try:
            from sentence_transformers import CrossEncoder
            self._model = CrossEncoder(self.model_name, device=self.device)
            logger.info(f"Loaded reranker model: {self.model_name}")
        except ImportError:
            raise ImportError(
                "sentence-transformers is required for reranking. "
                "Install with: pip install sentence-transformers"
            )
    
    @trace_function(name="rerank")
    def rerank(
        self,
        query: str,
        results: List[RetrievalResult],
        top_k: Optional[int] = None
    ) -> List[RetrievalResult]:
        """
        Rerank results using cross-encoder.
        
        Args:
            query: Original query
            results: Results to rerank
            top_k: Number of results to return (all if None)
            
        Returns:
            Reranked results
        """
        if not results:
            return []
        
        self._load_model()
        
        # Create query-document pairs
        pairs = [(query, result.content) for result in results]
        
        # Get reranking scores
        scores = self._model.predict(pairs)
        
        # Update scores and sort
        for result, score in zip(results, scores):
            result.score = float(score)
        
        sorted_results = sorted(results, key=lambda r: r.score, reverse=True)
        
        if top_k:
            return sorted_results[:top_k]
        return sorted_results


class RetrievalService:
    """
    Main retrieval service orchestrating different retrieval strategies.
    """
    
    def __init__(
        self,
        vector_store: VectorStore,
        embedding_model: EmbeddingModel,
        settings=None
    ):
        """
        Initialize retrieval service.
        
        Args:
            vector_store: Vector store instance
            embedding_model: Embedding model instance
            settings: Optional RetrievalSettings
        """
        self.vector_store = vector_store
        self.embedding_model = embedding_model
        self.settings = settings
        
        # Initialize retrievers
        self.dense_retriever = DenseRetriever(vector_store, embedding_model)
        self.sparse_retriever = None  # Initialized when documents are added
        self.hybrid_retriever = None
        
        # Initialize reranker if enabled
        self.reranker = None
        if settings and settings.rerank:
            self.reranker = Reranker(
                model_name=settings.rerank_model,
                device=embedding_model.device if embedding_model else "auto"
            )
        
        # Settings
        self.strategy = settings.strategy if settings else "dense"
        self.top_k = settings.top_k if settings else 10
        self.rerank_top_k = settings.rerank_top_k if settings else 5
    
    def build_sparse_index(self, documents: List[Dict[str, Any]]) -> None:
        """
        Build sparse index for BM25 retrieval.
        
        Args:
            documents: List of documents with 'content' and 'metadata'
        """
        bm25_settings = self.settings.bm25 if self.settings else None
        
        self.sparse_retriever = BM25Retriever(
            documents=documents,
            k1=bm25_settings.k1 if bm25_settings else 1.5,
            b=bm25_settings.b if bm25_settings else 0.75
        )
        
        # Create hybrid retriever
        self.hybrid_retriever = HybridRetriever(
            dense_retriever=self.dense_retriever,
            sparse_retriever=self.sparse_retriever
        )
    
    @trace_function(name="retrieve")
    def retrieve(
        self,
        query: str,
        top_k: Optional[int] = None,
        strategy: Optional[str] = None,
        filter_metadata: Optional[Dict[str, Any]] = None,
        rerank: Optional[bool] = None
    ) -> List[RetrievalResult]:
        """
        Retrieve relevant documents for a query.
        
        Args:
            query: Search query
            top_k: Number of results (uses default if not specified)
            strategy: Retrieval strategy ("dense", "sparse", "hybrid")
            filter_metadata: Optional metadata filters
            rerank: Whether to rerank results
            
        Returns:
            List of retrieval results
        """
        top_k = top_k or self.top_k
        strategy = strategy or self.strategy
        rerank = rerank if rerank is not None else (self.reranker is not None)
        
        # Select retriever
        if strategy == "hybrid" and self.hybrid_retriever:
            retriever = self.hybrid_retriever
        elif strategy == "sparse" and self.sparse_retriever:
            retriever = self.sparse_retriever
        else:
            retriever = self.dense_retriever
        
        # Retrieve
        results = retriever.retrieve(
            query=query,
            top_k=top_k * 2 if rerank else top_k,  # Get more for reranking
            filter_metadata=filter_metadata
        )
        
        # Rerank if enabled
        if rerank and self.reranker and results:
            results = self.reranker.rerank(
                query=query,
                results=results,
                top_k=self.rerank_top_k
            )
        
        return results[:top_k]
    
    def retrieve_by_type(
        self,
        query: str,
        chunk_type: ChunkType,
        top_k: Optional[int] = None
    ) -> List[RetrievalResult]:
        """
        Retrieve documents of a specific type.
        
        Args:
            query: Search query
            chunk_type: Type of chunks to retrieve
            top_k: Number of results
            
        Returns:
            Filtered retrieval results
        """
        results = self.retrieve(
            query=query,
            top_k=(top_k or self.top_k) * 2,  # Get more to filter
            filter_metadata={"chunk_type": chunk_type.value}
        )
        
        # Filter by type (in case filter not supported)
        filtered = [r for r in results if r.chunk_type == chunk_type]
        return filtered[:top_k or self.top_k]
    
    def retrieve_math(self, query: str, top_k: int = 10) -> List[RetrievalResult]:
        """Retrieve math-related content."""
        return self.retrieve_by_type(query, ChunkType.MATH, top_k)
    
    def retrieve_code(self, query: str, top_k: int = 10) -> List[RetrievalResult]:
        """Retrieve code-related content."""
        return self.retrieve_by_type(query, ChunkType.CODE, top_k)
    
    def retrieve_narrative(self, query: str, top_k: int = 10) -> List[RetrievalResult]:
        """Retrieve narrative/explanatory content."""
        return self.retrieve_by_type(query, ChunkType.NARRATIVE, top_k)
