"""FastEmbed wrapper for ONNX-based local embeddings.

Uses BAAI/bge-small-en-v1.5 (384 dimensions, ~80MB) which runs
on CPU with no API key required. Batch-embeds symbols during indexing
for semantic search via sqlite-vec.
"""

import logging
from typing import Optional

log = logging.getLogger(__name__)

# Lazy-loaded singleton
_embedder = None
MODEL_NAME = 'BAAI/bge-small-en-v1.5'
EMBEDDING_DIM = 384


class Embedder:
    """Lazy-loading FastEmbed wrapper."""

    def __init__(self, model_name: str = MODEL_NAME):
        self.model_name = model_name
        self._model = None

    @property
    def model(self):
        if self._model is None:
            try:
                from fastembed import TextEmbedding
                log.info('Loading embedding model %s...', self.model_name)
                self._model = TextEmbedding(model_name=self.model_name)
                log.info('Embedding model loaded.')
            except ImportError:
                log.warning(
                    'fastembed not installed. Semantic search disabled. '
                    'Install with: pip install fastembed',
                )
                raise
        return self._model

    def embed(self, text: str) -> list[float]:
        """Embed a single text string."""
        results = list(self.model.embed([text]))
        return results[0].tolist()

    def embed_batch(self, texts: list[str]) -> list[list[float]]:
        """Embed a batch of texts efficiently."""
        if not texts:
            return []
        results = list(self.model.embed(texts))
        return [r.tolist() for r in results]

    def embed_symbol(self, symbol_dict: dict) -> str:
        """Build a text representation of a symbol for embedding.

        Concatenates name, kind, signature, and docstring into a single
        string that captures the symbol's semantics.
        """
        parts = []
        if symbol_dict.get('kind'):
            parts.append(symbol_dict['kind'])
        if symbol_dict.get('qualified_name'):
            parts.append(symbol_dict['qualified_name'])
        if symbol_dict.get('signature'):
            parts.append(symbol_dict['signature'][:200])
        if symbol_dict.get('docstring'):
            parts.append(symbol_dict['docstring'][:300])
        return ' '.join(parts)


def get_embedder() -> Embedder:
    """Get the global embedder singleton."""
    global _embedder
    if _embedder is None:
        _embedder = Embedder()
    return _embedder


def hybrid_search(
    db,
    query: str,
    embedder: Optional[Embedder] = None,
    kind: str | None = None,
    limit: int = 20,
    k: int = 60,
) -> list[dict]:
    """Hybrid search combining FTS5 BM25 + vector similarity via RRF.

    Reciprocal Rank Fusion (RRF) merges ranked lists from different
    retrieval methods. Score = sum(1 / (k + rank)) across all lists.

    Args:
        db: Database instance.
        query: Search query string.
        embedder: Embedder instance (if None, falls back to FTS-only).
        kind: Optional symbol kind filter.
        limit: Max results to return.
        k: RRF constant (default 60).

    Returns:
        List of symbol dicts sorted by fused score.
    """
    fetch_limit = limit * 3

    # 1. FTS5 BM25 search
    fts_results = db.search_symbols_fts(query, kind=kind, limit=fetch_limit)

    # 2. Vector similarity search (if embedder available)
    vec_results = []
    if embedder:
        try:
            query_vec = embedder.embed(query)
            vec_results = db.search_symbols_vec(query_vec, limit=fetch_limit)
        except Exception:
            log.warning('Vector search failed, using FTS only', exc_info=True)

    # 3. Reciprocal Rank Fusion
    scores: dict[str, float] = {}
    symbol_data: dict[str, dict] = {}

    for rank, result in enumerate(fts_results):
        sid = result['symbol_id']
        scores[sid] = scores.get(sid, 0) + 1 / (k + rank)
        symbol_data[sid] = result

    for rank, result in enumerate(vec_results):
        sid = result['symbol_id']
        scores[sid] = scores.get(sid, 0) + 1 / (k + rank)
        if sid not in symbol_data:
            symbol_data[sid] = result

    # 4. Sort by fused score, return top-k
    ranked = sorted(scores.items(), key=lambda x: -x[1])[:limit]

    results = []
    for sid, score in ranked:
        entry = symbol_data[sid].copy()
        entry['score'] = round(score, 6)
        results.append(entry)

    return results
