"""
ChunkHive Semantic Embedding Engine
====================================

Generate and manage vector embeddings for code chunks to enable semantic search,
similarity matching, and AI-powered code understanding.

Author: ChunkHive
License: Apache 2.0
Version: 0.3.0
"""

from typing import List, Optional, Tuple, Any, TYPE_CHECKING
from pathlib import Path
import logging
import pickle
from dataclasses import dataclass
import numpy as np

if TYPE_CHECKING:
    from sentence_transformers import SentenceTransformer as ST
else:
    ST = Any

try:
    from sentence_transformers import SentenceTransformer
    SENTENCE_TRANSFORMERS_AVAILABLE = True
except ImportError:
    SENTENCE_TRANSFORMERS_AVAILABLE = False
    SentenceTransformer = Any  # type: ignore

try:
    import faiss  # type: ignore
    FAISS_AVAILABLE = True
except ImportError:
    FAISS_AVAILABLE = False
    faiss = Any  # type: ignore

from ..chunkers.schema import SemanticChunk

logger = logging.getLogger(__name__)


# =============================================================================
# MODEL CONFIGURATIONS
# =============================================================================

@dataclass
class ModelConfig:
    """Configuration for an embedding model."""
    name: str
    model_id: str
    dimension: int
    description: str
    best_for: str


AVAILABLE_MODELS = {
    "codebert": ModelConfig(
        name="codebert",
        model_id="microsoft/codebert-base",
        dimension=768,
        description="Microsoft's CodeBERT - trained on code and documentation",
        best_for="General purpose code embedding"
    ),
    "unixcoder": ModelConfig(
        name="unixcoder",
        model_id="microsoft/unixcoder-base",
        dimension=768,
        description="Microsoft's UniXcoder - unified cross-modal pre-training",
        best_for="Code-to-code and code-to-text matching"
    ),
    "codesage": ModelConfig(
        name="codesage",
        model_id="codesage/codesage-small",
        dimension=1024,
        description="CodeSage - optimized for code understanding",
        best_for="Semantic code search and retrieval"
    ),
    "minilm": ModelConfig(
        name="minilm",
        model_id="sentence-transformers/all-MiniLM-L6-v2",
        dimension=384,
        description="Fast and lightweight general-purpose model",
        best_for="Quick prototyping and testing"
    ),
    
}


# =============================================================================
# TEXT PREPARATION (Key to search quality!)
# =============================================================================

def _prepare_chunk_text(chunk: SemanticChunk) -> str:
    """
    Prepare rich text representation of a chunk for embedding.
    
    This is THE most important function for search quality.
    More context = better semantic understanding = better results.
    
    Industry standard: include name, docstring, type context, AND code.
    Just embedding raw code gives poor results because models need
    natural language context to understand what code DOES.
    """
    parts = []
    
    # 1. File context (helps locate results)
    filename = Path(chunk.file_path).stem
    parts.append(f"File: {filename}")
    
    # 2. Type label (class vs function vs method matters!)
    type_label = {
        "class": "Class definition",
        "function": "Function",
        "method": "Method",
        "async_function": "Async function",
        "async_method": "Async method",
        "module": "Module",
    }.get(chunk.chunk_type, chunk.chunk_type)
    parts.append(f"Type: {type_label}")
    
    # 3. Name with semantic label
    if chunk.ast.name:
        parts.append(f"Name: {chunk.ast.name}")
    
    # 4. Parent class (critical for methods!)
    if chunk.ast.parent:
        parts.append(f"Belongs to class: {chunk.ast.parent}")
    
    # 5. Docstring — MOST IMPORTANT for semantic matching
    # Docstrings describe WHAT code does in natural language
    # This is what makes "authentication function" find auth code!
    if chunk.ast.docstring:
        parts.append(f"Description: {chunk.ast.docstring}")
    
    # 6. Function signature (for functions/methods)
    if chunk.ast.name and chunk.chunk_type in ["function", "method", "async_function", "async_method"]:
        sig = chunk.ast.get_signature()
        if sig:
            parts.append(f"Signature: {sig}")
    
    # 7. Code — truncate to avoid token limits (most models cap at 512 tokens)
    code = chunk.code
    if len(code) > 1000:
        # Keep first 800 chars (function signature + key logic usually there)
        code = code[:800] + "\n# ... (truncated)"
    parts.append(code)
    
    return "\n".join(parts)


def _prepare_query_text(query: str) -> str:
    """
    Prepare query text for embedding.
    
    Adds context to help the model understand it's a code search query.
    This improves matching against code embeddings.
    """
    # Adding "Code:" prefix helps align query embedding space
    # with code embedding space for better matching
    return f"Code search: {query}\nFind code that: {query}"


# =============================================================================
# SEMANTIC EMBEDDING ENGINE
# =============================================================================

class SemanticEmbeddingEngine:
    """
    Generate semantic vector embeddings for code chunks.
    
    Converts code chunks into dense vector representations that capture
    semantic meaning, enabling similarity search and AI-powered code
    understanding.
    
    Example:
        >>> engine = SemanticEmbeddingEngine(model_name="codebert")
        >>> embeddings = engine.generate_embeddings(chunks)
        >>> print(f"Shape: {embeddings.shape}")  # (n_chunks, 768)
    """
    
    def __init__(
        self,
        model_name: str = "minilm",
        device: str = "cpu",
        cache_dir: Optional[Path] = None
    ):
        if not SENTENCE_TRANSFORMERS_AVAILABLE:
            raise ImportError(
                "sentence-transformers not installed. "
                "Run: pip install sentence-transformers"
            )
        
        if model_name not in AVAILABLE_MODELS:
            raise ValueError(
                f"Unknown model: '{model_name}'. "
                f"Available: {list(AVAILABLE_MODELS.keys())}"
            )
        
        self.model_config = AVAILABLE_MODELS[model_name]
        self.model_name = model_name
        self.dimension = self.model_config.dimension
        self.device = device
        
        logger.info(f"Loading model: {self.model_config.model_id}")
        
        self.model: Any = SentenceTransformer(  # type: ignore
            self.model_config.model_id,
            device=device,
            cache_folder=str(cache_dir) if cache_dir else None
        )
        
        logger.info(f"Model ready on {device}")
    
    def generate_embeddings(
        self,
        chunks: List[SemanticChunk],
        batch_size: int = 32,
        show_progress: bool = True
    ) -> np.ndarray:
        """
        Generate embeddings for a list of code chunks.
        
        Uses rich text preparation (name + docstring + code) for better
        semantic understanding than raw code alone.
        
        Returns:
            Array of shape (len(chunks), dimension)
        """
        if not chunks:
            return np.array([])
        
        logger.info(f"Generating embeddings for {len(chunks)} chunks")
        
        # Rich text preparation — key to quality!
        texts = [_prepare_chunk_text(chunk) for chunk in chunks]
        
        embeddings = self.model.encode(  # type: ignore
            texts,
            batch_size=batch_size,
            show_progress_bar=show_progress,
            convert_to_numpy=True,
            normalize_embeddings=True  # L2 normalization for cosine similarity
        )
        
        logger.info(f"Generated embeddings: {embeddings.shape}")
        return embeddings  # type: ignore
    
    def generate_embedding_for_query(self, query: str) -> np.ndarray:
        """
        Generate embedding for a search query.
        
        Adds query context prefix to improve matching against code embeddings.
        """
        prepared_query = _prepare_query_text(query)
        
        embedding = self.model.encode(  # type: ignore
            [prepared_query],
            convert_to_numpy=True,
            normalize_embeddings=True
        )
        return embedding[0]  # type: ignore
    
    def attach_embeddings_to_chunks(
        self,
        chunks: List[SemanticChunk],
        embeddings: np.ndarray
    ) -> None:
        """Attach generated embeddings to chunk objects in-place."""
        if len(chunks) != len(embeddings):
            raise ValueError(
                f"Mismatch: {len(chunks)} chunks vs {len(embeddings)} embeddings"
            )
        
        for chunk, embedding in zip(chunks, embeddings):
            chunk.embedding = embedding.tolist()
            chunk.embedding_model = self.model_config.model_id
        
        logger.info(f"Attached embeddings to {len(chunks)} chunks")


# =============================================================================
# VECTOR SEARCH INDEX
# =============================================================================

class VectorSearchIndex:
    """
    Fast similarity search using FAISS.
    
    Provides k-nearest-neighbor search over code embeddings.
    Returns similarity percentages (0-100%) for human-readable results.
    
    Example:
        >>> index = VectorSearchIndex(dimension=768)
        >>> index.add_embeddings(embeddings, chunks)
        >>> results = index.search("authenticate user", engine, k=5)
        >>> for chunk, similarity in results:
        ...     print(f"{chunk.ast.name}: {similarity:.1f}% match")
    """
    
    def __init__(
        self,
        dimension: int,
        use_gpu: bool = False,
        index_type: str = "flat"
    ):
        if not FAISS_AVAILABLE:
            raise ImportError(
                "faiss not installed. Run: pip install faiss-cpu"
            )
        
        self.dimension = dimension
        self.use_gpu = use_gpu
        self.index_type = index_type
        self.chunks: List[SemanticChunk] = []
        
        if index_type == "flat":
            # Exact search — good for <100k vectors
            # Using Inner Product (IP) since embeddings are L2-normalized
            # IP on normalized vectors = cosine similarity!
            self.index: Any = faiss.IndexFlatIP(dimension)  # type: ignore
        elif index_type == "ivf":
            # Approximate search — good for >100k vectors
            quantizer: Any = faiss.IndexFlatIP(dimension)  # type: ignore
            self.index = faiss.IndexIVFFlat(quantizer, dimension, 100)  # type: ignore
        else:
            raise ValueError(f"Unknown index type: {index_type}. Use 'flat' or 'ivf'")
        
        logger.info(f"Created {index_type} FAISS index (dim={dimension})")
    
    def add_embeddings(
        self,
        embeddings: np.ndarray,
        chunks: List[SemanticChunk]
    ) -> None:
        """Add embeddings and chunks to the index."""
        if len(embeddings) != len(chunks):
            raise ValueError("Embeddings and chunks must have same length")
        
        if self.index_type == "ivf" and not self.index.is_trained:  # type: ignore
            logger.info("Training IVF index...")
            self.index.train(embeddings.astype('float32'))  # type: ignore
        
        self.index.add(embeddings.astype('float32'))  # type: ignore
        self.chunks.extend(chunks)
        
        logger.info(f"Index now has {self.index.ntotal} vectors")  # type: ignore
    
    def search(
        self,
        query: str,
        embedding_engine: SemanticEmbeddingEngine,
        k: int = 5,
        min_similarity: float = 0.0
    ) -> List[Tuple[SemanticChunk, float]]:
        """
        Search for similar code chunks.
        
        Args:
            query: Natural language query or code snippet
            embedding_engine: Engine to generate query embedding
            k: Number of results to return
            min_similarity: Minimum similarity % to include (0-100)
            
        Returns:
            List of (chunk, similarity_pct) tuples — higher % = better match
        """
        if self.index.ntotal == 0:  # type: ignore
            logger.warning("Index is empty")
            return []
        
        query_embedding = embedding_engine.generate_embedding_for_query(query)
        query_vec = query_embedding.reshape(1, -1).astype('float32')
        
        k = min(k, self.index.ntotal)  # type: ignore
        scores, indices = self.index.search(query_vec, k)  # type: ignore
        
        results = []
        for idx, score in zip(indices[0], scores[0]):
            if 0 <= idx < len(self.chunks):
                # Convert cosine similarity (-1 to 1) to percentage (0 to 100%)
                # score is dot product of normalized vectors = cosine similarity
                similarity_pct = float(score) * 100
                similarity_pct = max(0.0, min(100.0, similarity_pct))
                
                if similarity_pct >= min_similarity:
                    results.append((self.chunks[idx], similarity_pct))
        
        return results
    
    def search_by_embedding(
        self,
        embedding: np.ndarray,
        k: int = 5
    ) -> List[Tuple[SemanticChunk, float]]:
        """Search using a pre-computed embedding. Returns similarity percentages."""
        if self.index.ntotal == 0:  # type: ignore
            return []
        
        vec = embedding.reshape(1, -1).astype('float32')
        k = min(k, self.index.ntotal)  # type: ignore
        scores, indices = self.index.search(vec, k)  # type: ignore
        
        results = []
        for idx, score in zip(indices[0], scores[0]):
            if 0 <= idx < len(self.chunks):
                similarity_pct = max(0.0, min(100.0, float(score) * 100))
                results.append((self.chunks[idx], similarity_pct))
        
        return results
    
    def save(self, path: Path) -> None:
        """Save index and chunks to disk."""
        path.parent.mkdir(parents=True, exist_ok=True)
        faiss.write_index(self.index, str(path.with_suffix('.faiss')))  # type: ignore
        with open(path.with_suffix('.chunks'), 'wb') as f:
            pickle.dump(self.chunks, f)
        logger.info(f"Saved index to {path}")
    
    @classmethod
    def load(cls, path: Path, dimension: int) -> 'VectorSearchIndex':
        """Load index and chunks from disk."""
        index_obj = cls(dimension=dimension)
        index_obj.index = faiss.read_index(str(path.with_suffix('.faiss')))  # type: ignore
        with open(path.with_suffix('.chunks'), 'rb') as f:
            index_obj.chunks = pickle.load(f)
        logger.info(f"Loaded index from {path} ({index_obj.index.ntotal} vectors)")  # type: ignore
        return index_obj


# =============================================================================
# CONVENIENCE FUNCTIONS
# =============================================================================

def embed_chunks(
    chunks: List[SemanticChunk],
    model_name: str = "minilm",
    attach_to_chunks: bool = True
) -> np.ndarray:
    """Embed a list of chunks, optionally attaching embeddings to them."""
    engine = SemanticEmbeddingEngine(model_name=model_name)
    embeddings = engine.generate_embeddings(chunks)
    if attach_to_chunks:
        engine.attach_embeddings_to_chunks(chunks, embeddings)
    return embeddings


def create_search_index(
    chunks: List[SemanticChunk],
    model_name: str = "minilm",
    index_type: str = "flat"
) -> Tuple[VectorSearchIndex, SemanticEmbeddingEngine]:
    """Generate embeddings and build a searchable index in one step."""
    engine = SemanticEmbeddingEngine(model_name=model_name)
    embeddings = engine.generate_embeddings(chunks)
    index = VectorSearchIndex(dimension=engine.dimension, index_type=index_type)
    index.add_embeddings(embeddings, chunks)
    return index, engine


# =============================================================================
# EXPORTS
# =============================================================================

__all__ = [
    "SemanticEmbeddingEngine",
    "VectorSearchIndex",
    "ModelConfig",
    "AVAILABLE_MODELS",
    "embed_chunks",
    "create_search_index",
]