"""Local embeddings via fastembed (bge-small), L2-normalized for cosine = dot."""

from __future__ import annotations

from collections.abc import Iterable
from typing import Protocol

import numpy as np

from . import config


class EmbeddingModel(Protocol):
    """Minimal fastembed-compatible surface (also satisfied by test fakes)."""

    def embed(self, texts: Iterable[str], **kwargs: object) -> Iterable[np.ndarray]: ...
    def query_embed(self, texts: Iterable[str], **kwargs: object) -> Iterable[np.ndarray]: ...


_model: EmbeddingModel | None = None


def get_model() -> EmbeddingModel:
    """Lazily construct the (heavy) fastembed model once."""
    global _model
    if _model is None:  # pragma: no cover - exercised only with the real model
        from fastembed import TextEmbedding

        # threads caps onnxruntime's intra-op parallelism (the main CPU driver);
        # CPUExecutionProvider avoids grabbing any accelerator.
        _model = TextEmbedding(
            model_name=config.EMBED_MODEL,
            threads=config.EMBED_THREADS,
            providers=["CPUExecutionProvider"],
        )
    return _model


def _normalize(matrix: np.ndarray) -> np.ndarray:
    if matrix.size == 0:
        return matrix.reshape(0, config.EMBED_DIM).astype(np.float32)
    norms = np.linalg.norm(matrix, axis=1, keepdims=True)
    norms[norms == 0] = 1.0
    return (matrix / norms).astype(np.float32)


def embed_passages(texts: list[str], *, model: EmbeddingModel | None = None) -> np.ndarray:
    """Embed corpus passages -> (N, dim) normalized float32."""
    model = model or get_model()
    vectors = list(
        model.embed(texts, batch_size=config.EMBED_BATCH_SIZE, parallel=config.EMBED_PARALLEL)
    )
    return _normalize(np.array(vectors, dtype=np.float32))


def embed_query(text: str, *, model: EmbeddingModel | None = None) -> np.ndarray:
    """Embed a search query -> (dim,) normalized float32 (with bge query prefix)."""
    model = model or get_model()
    vector = np.array(next(iter(model.query_embed([text]))), dtype=np.float32)
    return _normalize(vector.reshape(1, -1))[0]