"""Internal RAG: End-to-end differentiable retrieval inside the model.

Extends KeyValueMemory with:
- Scalable top-k retrieval (sparse attention)
- Document ingestion pipeline
- Optional FAISS integration for millions of entries
"""

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Optional, List, Tuple, Callable, Any
from .memory import KeyValueMemory


class ScalableMemory(nn.Module):
    """Memory module that scales to millions of entries via top-k retrieval.
    
    Unlike KeyValueMemory which uses full softmax attention, this uses
    sparse top-k retrieval for efficiency, then soft attention over the
    top-k for differentiability.
    
    Args:
        dim: Embedding dimension.
        initial_size: Initial number of memory slots.
        top_k: Number of entries to retrieve.
        use_faiss: Use FAISS for fast approximate search (requires faiss-cpu/gpu).
        learnable: If True, keys/values are learnable parameters.
        
    Example:
        >>> memory = ScalableMemory(dim=768, initial_size=10000, top_k=32)
        >>> retrieved = memory.retrieve(query, k=16)
    """
    
    def __init__(
        self,
        dim: int,
        initial_size: int = 1024,
        top_k: int = 32,
        use_faiss: bool = False,
        learnable: bool = False,
    ):
        super().__init__()
        self.dim = dim
        self.memory_size = initial_size
        self.top_k = top_k
        self.use_faiss = use_faiss
        self.scale = math.sqrt(dim)
        
        # Memory storage
        if learnable:
            self.keys = nn.Parameter(torch.randn(initial_size, dim) * 0.02)
            self.values = nn.Parameter(torch.randn(initial_size, dim) * 0.02)
        else:
            self.register_buffer("keys", torch.randn(initial_size, dim) * 0.02)
            self.register_buffer("values", torch.randn(initial_size, dim) * 0.02)
        
        self.register_buffer("usage", torch.zeros(initial_size))
        self.write_pointer = 0
        
        # FAISS index (initialized lazily)
        self._faiss_index = None
    
    def _build_faiss_index(self):
        """Build FAISS index for fast search."""
        try:
            import faiss
        except ImportError:
            raise ImportError("FAISS required. Install with: pip install faiss-cpu")
        
        # Use inner product (MIPS) index
        self._faiss_index = faiss.IndexFlatIP(self.dim)
        
        # Add current keys
        keys_np = self.keys.detach().cpu().numpy()
        self._faiss_index.add(keys_np)
    
    def add(self, keys: torch.Tensor, values: torch.Tensor):
        """Add entries to memory.
        
        Args:
            keys: Key vectors (batch, dim).
            values: Value vectors (batch, dim).
        """
        batch_size = keys.size(0)
        
        # Extend if needed
        if self.write_pointer + batch_size > self.memory_size:
            new_size = max(self.memory_size * 2, self.write_pointer + batch_size)
            self._extend(new_size)
        
        # Write entries
        start = self.write_pointer
        end = start + batch_size
        
        with torch.no_grad():
            self.keys[start:end] = keys.detach()
            self.values[start:end] = values.detach()
            self.usage[start:end] = 1.0
        
        self.write_pointer = end
        
        # Invalidate FAISS index
        if self._faiss_index is not None:
            self._faiss_index = None
    
    def _extend(self, new_size: int):
        """Grow memory capacity."""
        device = self.keys.device
        
        new_keys = torch.randn(new_size, self.dim, device=device) * 0.02
        new_values = torch.randn(new_size, self.dim, device=device) * 0.02
        new_usage = torch.zeros(new_size, device=device)
        
        new_keys[:self.memory_size] = self.keys
        new_values[:self.memory_size] = self.values
        new_usage[:self.memory_size] = self.usage
        
        if isinstance(self.keys, nn.Parameter):
            self.keys = nn.Parameter(new_keys)
            self.values = nn.Parameter(new_values)
        else:
            self.register_buffer("keys", new_keys)
            self.register_buffer("values", new_values)
        self.register_buffer("usage", new_usage)
        self.memory_size = new_size
    
    def top_k_indices(self, query: torch.Tensor, k: Optional[int] = None) -> torch.Tensor:
        """Get indices of top-k most similar entries.
        
        Args:
            query: Query vectors (batch, dim) or (batch, seq, dim).
            k: Number of entries to retrieve.
            
        Returns:
            Indices tensor (batch, k) or (batch, seq, k).
        """
        k = k or self.top_k
        
        if self.use_faiss and self._faiss_index is None:
            self._build_faiss_index()
        
        if self.use_faiss:
            # Use FAISS for approximate search
            query_np = query.detach().cpu().numpy()
            if query.dim() == 3:
                orig_shape = query.shape[:2]
                query_np = query_np.reshape(-1, self.dim)
            
            _, indices = self._faiss_index.search(query_np, k)
            indices = torch.from_numpy(indices).to(query.device)
            
            if query.dim() == 3:
                indices = indices.view(*orig_shape, k)
        else:
            # Exact search
            scores = torch.matmul(query, self.keys.t()) / self.scale
            _, indices = scores.topk(k, dim=-1)
        
        return indices
    
    def retrieve(
        self,
        query: torch.Tensor,
        k: Optional[int] = None,
        return_weights: bool = False,
    ) -> Tuple[torch.Tensor, ...]:
        """Retrieve top-k entries with differentiable soft attention.
        
        Args:
            query: Query vectors (batch, dim) or (batch, seq, dim).
            k: Number of entries to retrieve.
            return_weights: Return attention weights.
            
        Returns:
            retrieved: Retrieved values (batch, dim) or (batch, seq, dim).
            weights: (optional) Attention weights over top-k.
        """
        k = k or self.top_k
        
        # Get top-k indices
        indices = self.top_k_indices(query, k)
        
        # Gather top-k keys and values
        if query.dim() == 2:  # (batch, dim)
            top_k_keys = self.keys[indices]  # (batch, k, dim)
            top_k_values = self.values[indices]  # (batch, k, dim)
            
            # Soft attention over top-k
            scores = torch.einsum('bd,bkd->bk', query, top_k_keys) / self.scale
            weights = F.softmax(scores, dim=-1)
            retrieved = torch.einsum('bk,bkd->bd', weights, top_k_values)
            
        else:  # (batch, seq, dim)
            batch, seq, _ = query.shape
            
            # Flatten for indexing
            flat_indices = indices.view(-1, k)  # (batch*seq, k)
            top_k_keys = self.keys[flat_indices].view(batch, seq, k, -1)
            top_k_values = self.values[flat_indices].view(batch, seq, k, -1)
            
            # Soft attention
            scores = torch.einsum('bsd,bskd->bsk', query, top_k_keys) / self.scale
            weights = F.softmax(scores, dim=-1)
            retrieved = torch.einsum('bsk,bskd->bsd', weights, top_k_values)
        
        if return_weights:
            return retrieved, weights
        return retrieved
    
    def save(self, path: str):
        """Save memory to disk."""
        torch.save({
            "keys": self.keys,
            "values": self.values,
            "usage": self.usage,
            "write_pointer": self.write_pointer,
            "dim": self.dim,
            "memory_size": self.memory_size,
        }, path)
    
    def load(self, path: str, device: str = None):
        """Load memory from disk."""
        if device is None:
            device = self.keys.device
        
        data = torch.load(path, map_location=device)
        
        if data["memory_size"] != self.memory_size:
            self._extend(data["memory_size"])
        
        with torch.no_grad():
            self.keys.copy_(data["keys"])
            self.values.copy_(data["values"])
            self.usage.copy_(data["usage"])
        self.write_pointer = data["write_pointer"]


class DocumentStore:
    """Pipeline for ingesting documents into ScalableMemory.
    
    Converts text documents into key-value pairs using an encoder model,
    then stores them in memory for retrieval.
    
    Args:
        memory: ScalableMemory instance.
        encoder: Model to encode text chunks (must have forward() returning embeddings).
        tokenizer: Tokenizer for the encoder.
        chunk_size: Max tokens per chunk.
        chunk_overlap: Overlap between chunks.
        
    Example:
        >>> memory = ScalableMemory(dim=768, initial_size=100000)
        >>> store = DocumentStore(memory, encoder_model, tokenizer)
        >>> store.ingest(["Document 1 text", "Document 2 text"])
    """
    
    def __init__(
        self,
        memory: ScalableMemory,
        encoder: nn.Module,
        tokenizer,
        chunk_size: int = 256,
        chunk_overlap: int = 32,
        device: str = "cuda",
    ):
        self.memory = memory
        self.encoder = encoder.to(device)
        self.tokenizer = tokenizer
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        self.device = device
        
        self.encoder.eval()
    
    def _chunk_text(self, text: str) -> List[str]:
        """Split text into overlapping chunks."""
        tokens = self.tokenizer.encode(text)
        chunks = []
        
        start = 0
        while start < len(tokens):
            end = min(start + self.chunk_size, len(tokens))
            chunk_tokens = tokens[start:end]
            chunks.append(self.tokenizer.decode(chunk_tokens))
            
            if end >= len(tokens):
                break
            start = end - self.chunk_overlap
        
        return chunks
    
    def _encode_chunk(self, chunk: str) -> torch.Tensor:
        """Encode a text chunk to embedding."""
        tokens = self.tokenizer.encode(chunk)
        input_ids = torch.tensor([tokens], device=self.device)
        
        with torch.no_grad():
            output = self.encoder(input_ids)
            
            # Handle different output formats
            if isinstance(output, dict):
                hidden = output.get("last_hidden_state", output.get("hidden_states", output))
            elif isinstance(output, tuple):
                hidden = output[0]
            else:
                hidden = output
            
            # Mean pooling
            embedding = hidden.mean(dim=1).squeeze(0)
        
        return embedding
    
    def ingest(
        self,
        documents: List[str],
        batch_size: int = 32,
        show_progress: bool = True,
    ):
        """Ingest documents into memory.
        
        Args:
            documents: List of document texts.
            batch_size: Encoding batch size.
            show_progress: Show progress bar.
        """
        all_chunks = []
        for doc in documents:
            chunks = self._chunk_text(doc)
            all_chunks.extend(chunks)
        
        if show_progress:
            try:
                from tqdm.auto import tqdm
                iterator = tqdm(range(0, len(all_chunks), batch_size), desc="Ingesting")
            except ImportError:
                iterator = range(0, len(all_chunks), batch_size)
        else:
            iterator = range(0, len(all_chunks), batch_size)
        
        for i in iterator:
            batch_chunks = all_chunks[i:i + batch_size]
            
            embeddings = []
            for chunk in batch_chunks:
                emb = self._encode_chunk(chunk)
                embeddings.append(emb)
            
            if embeddings:
                keys = torch.stack(embeddings)
                values = keys.clone()  # For now, key == value
                self.memory.add(keys, values)
        
        print(f"Ingested {len(all_chunks)} chunks from {len(documents)} documents")
    
    def query(
        self,
        text: str,
        k: int = 5,
    ) -> torch.Tensor:
        """Query memory with text.
        
        Args:
            text: Query text.
            k: Number of results.
            
        Returns:
            Retrieved embeddings.
        """
        query_emb = self._encode_chunk(text).unsqueeze(0)
        retrieved = self.memory.retrieve(query_emb, k=k)
        return retrieved


class InternalRAGBlock(nn.Module):
    """Transformer block with integrated RAG retrieval.
    
    Replaces standard attention with retrieval-augmented attention:
    1. Query memory for relevant context
    2. Concatenate retrieved context with self-attention output
    3. Gate to control how much retrieval contributes
    
    Args:
        dim: Model dimension.
        memory: ScalableMemory instance.
        num_heads: Number of attention heads.
        top_k: Number of memory entries to retrieve.
        dropout: Dropout probability.
    """
    
    def __init__(
        self,
        dim: int,
        memory: ScalableMemory,
        num_heads: int = 8,
        top_k: int = 16,
        dropout: float = 0.1,
    ):
        super().__init__()
        self.dim = dim
        self.memory = memory
        self.top_k = top_k
        
        # Query projection for memory lookup
        self.query_proj = nn.Linear(dim, dim)
        
        # Integration layers
        self.retrieved_proj = nn.Linear(dim, dim)
        self.gate = nn.Sequential(
            nn.Linear(dim * 2, dim),
            nn.Sigmoid(),
        )
        nn.init.constant_(self.gate[0].bias, -2.0)
        
        self.norm = nn.LayerNorm(dim)
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Apply retrieval-augmented processing.
        
        Args:
            x: Input hidden states (batch, seq, dim).
            
        Returns:
            Output with retrieval integration.
        """
        normed = self.norm(x)
        
        # Generate query
        query = self.query_proj(normed)
        
        # Retrieve from memory per-token (preserves local context sensitivity)
        # ScalableMemory.retrieve() supports 3D input (batch, seq, dim)
        retrieved = self.memory.retrieve(query, k=self.top_k)  # (batch, seq, dim)
        
        # Project retrieved values
        retrieved = self.retrieved_proj(retrieved)
        retrieved = self.dropout(retrieved)
        
        # Gated integration
        combined = torch.cat([normed, retrieved], dim=-1)
        gate = self.gate(combined)
        
        output = x + gate * retrieved
        
        return output
