"""
Lazy-loaded local embeddings via fastembed.
Model: BAAI/bge-small-en-v1.5 (~40MB, 384 dims, good quality/speed tradeoff)
Downloaded on first use to ~/.cache/fastembed/
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np

if TYPE_CHECKING:
    from fastembed import TextEmbedding

_model: "TextEmbedding | None" = None
MODEL_NAME = "BAAI/bge-small-en-v1.5"


def _get_model() -> "TextEmbedding":
    global _model
    if _model is None:
        from fastembed import TextEmbedding
        _model = TextEmbedding(MODEL_NAME)
    return _model


def embed(text: str) -> list[float]:
    model = _get_model()
    vectors = list(model.embed([text]))
    return vectors[0].tolist()


def cosine_similarity(a: list[float], b: list[float]) -> float:
    va = np.array(a, dtype=np.float32)
    vb = np.array(b, dtype=np.float32)
    denom = np.linalg.norm(va) * np.linalg.norm(vb)
    if denom == 0:
        return 0.0
    return float(np.dot(va, vb) / denom)


def rank_by_similarity(
    query: str,
    candidates: list[dict],
    top_k: int = 20,
    threshold: float = 0.3,
) -> list[dict]:
    """
    Rank a list of memory dicts by cosine similarity to query.
    Each candidate must have an 'embedding' key (list[float]).
    Returns top_k results above threshold, with a '_score' field added.
    """
    query_vec = embed(query)
    scored = []
    for c in candidates:
        emb = c.get("embedding")
        if not emb:
            continue
        score = cosine_similarity(query_vec, emb)
        if score >= threshold:
            result = {k: v for k, v in c.items() if k != "embedding"}
            result["_score"] = round(score, 4)
            scored.append(result)

    scored.sort(key=lambda x: x["_score"], reverse=True)
    return scored[:top_k]
