"""
Document indexing service for building and managing the vector store.

Handles:
- Embedding generation
- Vector store population
- Index management
"""

from pathlib import Path
from typing import Any, Dict, List, Optional, Union

from .chunking import Chunk, ChunkType
from .ingestion import ParsedDocument
from ..infra.vectordb import VectorStore, Document, create_vector_store
from ..infra.logging import get_logger, LogContext
from ..infra.tracing import trace_function

logger = get_logger(__name__)


class IndexingService:
    """
    Service for indexing documents into the vector store.
    
    Manages the embedding generation and storage of document chunks.
    """
    
    def __init__(
        self,
        vector_store: VectorStore,
        embedding_model: Optional["EmbeddingModel"] = None,
        settings=None
    ):
        """
        Initialize the indexing service.
        
        Args:
            vector_store: VectorStore instance
            embedding_model: EmbeddingModel for generating embeddings
            settings: Optional Settings
        """
        self.vector_store = vector_store
        self.embedding_model = embedding_model
        self.settings = settings
    
    @trace_function(name="index_chunks")
    def index_chunks(
        self,
        chunks: List[Chunk],
        batch_size: int = 32,
        show_progress: bool = True
    ) -> List[str]:
        """
        Index chunks into the vector store.
        
        Args:
            chunks: List of chunks to index
            batch_size: Batch size for embedding generation
            show_progress: Whether to show progress
            
        Returns:
            List of document IDs
        """
        if not chunks:
            logger.warning("No chunks to index")
            return []
        
        logger.info(f"Indexing {len(chunks)} chunks")
        
        # Generate embeddings if model available
        if self.embedding_model:
            contents = [chunk.content for chunk in chunks]
            embeddings = self.embedding_model.embed_batch(contents, batch_size=batch_size)
        else:
            embeddings = [None] * len(chunks)
        
        # Create documents for vector store
        documents = []
        for chunk, embedding in zip(chunks, embeddings):
            metadata = {
                "chunk_type": chunk.chunk_type.value,
                "chunk_index": chunk.chunk_index,
                "document_id": chunk.document_id or "",
                "section_hierarchy": "|".join(chunk.section_hierarchy),
                "page_numbers": ",".join(map(str, chunk.page_numbers)),
                **chunk.metadata
            }
            
            doc = Document(
                id=chunk.id,
                content=chunk.content,
                embedding=embedding,
                metadata=metadata
            )
            documents.append(doc)
        
        # Add to vector store
        doc_ids = self.vector_store.add_documents(documents)
        
        logger.info(f"Indexed {len(doc_ids)} documents")
        return doc_ids
    
    @trace_function(name="index_document")
    def index_document(
        self,
        document: ParsedDocument,
        chunking_service: "ChunkingService",
        strategy: str = "hierarchical"
    ) -> List[str]:
        """
        Index a complete document.
        
        Args:
            document: ParsedDocument to index
            chunking_service: ChunkingService for chunking
            strategy: Chunking strategy to use
            
        Returns:
            List of chunk IDs
        """
        with LogContext(document_id=document.metadata.source_path):
            # Chunk the document
            chunks = chunking_service.chunk(document, strategy=strategy)
            
            # Index the chunks
            return self.index_chunks(chunks)
    
    def delete_document(self, document_id: str) -> None:
        """
        Delete all chunks for a document.
        
        Args:
            document_id: Document ID to delete
        """
        # Note: This requires filtering support in the vector store
        # For now, we'd need to track document-chunk mappings separately
        logger.warning("Document deletion not fully implemented - requires chunk tracking")
    
    def get_stats(self) -> Dict[str, Any]:
        """Get indexing statistics."""
        return {
            "total_documents": self.vector_store.count(),
            "vector_store_backend": type(self.vector_store).__name__
        }
    
    def clear(self) -> None:
        """Clear all indexed documents."""
        self.vector_store.clear()
        logger.info("Cleared all indexed documents")


class EmbeddingModel:
    """
    Wrapper for embedding models.
    
    Supports multiple backends:
    - Sentence Transformers
    - HuggingFace Transformers
    - Custom embedding functions
    """
    
    def __init__(
        self,
        model_name: str = "BAAI/bge-small-en-v1.5",
        model_type: str = "sentence-transformers",
        device: str = "auto",
        normalize: bool = True,
        batch_size: int = 32
    ):
        """
        Initialize the embedding model.
        
        Args:
            model_name: Model name/path
            model_type: Type of model ("sentence-transformers", "huggingface")
            device: Device to use ("auto", "cpu", "cuda", "mps")
            normalize: Whether to normalize embeddings
            batch_size: Default batch size
        """
        self.model_name = model_name
        self.model_type = model_type
        self.device = self._resolve_device(device)
        self.normalize = normalize
        self.batch_size = batch_size
        
        self._model = None
        self._dimension = None
    
    def _resolve_device(self, device: str) -> str:
        """Resolve device string to actual device."""
        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 embedding model."""
        if self._model is not None:
            return
        
        logger.info(f"Loading embedding model: {self.model_name}")
        
        if self.model_type == "sentence-transformers":
            try:
                from sentence_transformers import SentenceTransformer
                self._model = SentenceTransformer(self.model_name, device=self.device)
                self._dimension = self._model.get_sentence_embedding_dimension()
            except ImportError:
                raise ImportError(
                    "sentence-transformers is required. "
                    "Install with: pip install sentence-transformers"
                )
        elif self.model_type == "huggingface":
            try:
                from transformers import AutoModel, AutoTokenizer
                import torch
                
                self._tokenizer = AutoTokenizer.from_pretrained(self.model_name)
                self._model = AutoModel.from_pretrained(self.model_name).to(self.device)
                self._model.eval()
                
                # Get dimension from model config
                self._dimension = self._model.config.hidden_size
            except ImportError:
                raise ImportError(
                    "transformers and torch are required. "
                    "Install with: pip install transformers torch"
                )
        else:
            raise ValueError(f"Unknown model type: {self.model_type}")
        
        logger.info(f"Loaded embedding model with dimension {self._dimension}")
    
    @property
    def dimension(self) -> int:
        """Get embedding dimension."""
        self._load_model()
        return self._dimension
    
    @trace_function(name="embed_text")
    def embed(self, text: str) -> List[float]:
        """
        Embed a single text.
        
        Args:
            text: Text to embed
            
        Returns:
            Embedding vector as list of floats
        """
        self._load_model()
        
        if self.model_type == "sentence-transformers":
            embedding = self._model.encode(
                text,
                normalize_embeddings=self.normalize,
                convert_to_numpy=True
            )
            return embedding.tolist()
        else:
            return self._embed_with_transformers([text])[0]
    
    @trace_function(name="embed_batch")
    def embed_batch(
        self,
        texts: List[str],
        batch_size: Optional[int] = None,
        show_progress: bool = False
    ) -> List[List[float]]:
        """
        Embed a batch of texts.
        
        Args:
            texts: List of texts to embed
            batch_size: Batch size (uses default if not specified)
            show_progress: Whether to show progress bar
            
        Returns:
            List of embedding vectors
        """
        self._load_model()
        batch_size = batch_size or self.batch_size
        
        if self.model_type == "sentence-transformers":
            embeddings = self._model.encode(
                texts,
                batch_size=batch_size,
                normalize_embeddings=self.normalize,
                show_progress_bar=show_progress,
                convert_to_numpy=True
            )
            return embeddings.tolist()
        else:
            return self._embed_with_transformers(texts, batch_size)
    
    def _embed_with_transformers(
        self,
        texts: List[str],
        batch_size: int = 32
    ) -> List[List[float]]:
        """Embed using HuggingFace transformers."""
        import torch
        import numpy as np
        
        all_embeddings = []
        
        for i in range(0, len(texts), batch_size):
            batch_texts = texts[i:i + batch_size]
            
            # Tokenize
            inputs = self._tokenizer(
                batch_texts,
                padding=True,
                truncation=True,
                max_length=512,
                return_tensors="pt"
            ).to(self.device)
            
            # Generate embeddings
            with torch.no_grad():
                outputs = self._model(**inputs)
                
                # Mean pooling
                attention_mask = inputs["attention_mask"]
                token_embeddings = outputs.last_hidden_state
                
                input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
                embeddings = torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(
                    input_mask_expanded.sum(1), min=1e-9
                )
                
                if self.normalize:
                    embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
                
                all_embeddings.extend(embeddings.cpu().numpy().tolist())
        
        return all_embeddings


def create_embedding_model(settings) -> EmbeddingModel:
    """
    Factory function to create an embedding model from settings.
    
    Args:
        settings: EmbeddingSettings or Settings object
        
    Returns:
        EmbeddingModel instance
    """
    if hasattr(settings, "embedding"):
        emb_settings = settings.embedding
    else:
        emb_settings = settings
    
    return EmbeddingModel(
        model_name=emb_settings.model_name,
        model_type=emb_settings.model_type,
        device=emb_settings.device,
        normalize=emb_settings.normalize,
        batch_size=emb_settings.batch_size
    )
