"""Dense text embeddings via BAAI/bge-m3.

bge-m3 maps Hindi passages and an English query into the **same** vector space, so an
English sentence retrieves Hindi content natively — no query-time translation. It needs
no "query:"/"passage:" prefix (symmetric), and its dense vectors are unit-normalized, so
cosine == dot product.
"""
from __future__ import annotations

from typing import List, Optional

import numpy as np

from app.config import Config, get_config


class Embedder:
    def __init__(self, cfg: Optional[Config] = None):
        self.cfg = cfg or get_config()
        self._model = None

    def _load(self):
        if self._model is not None:
            return self._model
        from FlagEmbedding import BGEM3FlagModel

        e = self.cfg.embedding
        # FlagEmbedding auto-selects CUDA when available; use_fp16 ~halves VRAM and ~2x speed.
        self._model = BGEM3FlagModel(e["model"], use_fp16=bool(e["use_fp16"]))
        return self._model

    def embed_passages(self, texts: List[str]) -> np.ndarray:
        """Return an (N, dim) float32 array of dense embeddings."""
        if not texts:
            return np.zeros((0, self.cfg.embedding["dim"]), dtype=np.float32)
        model = self._load()
        e = self.cfg.embedding
        out = model.encode(
            texts,
            batch_size=e["batch_size"],
            max_length=e["max_length"],
            return_dense=True,
            return_sparse=False,
            return_colbert_vecs=False,
        )
        vecs = np.asarray(out["dense_vecs"], dtype=np.float32)
        if vecs.ndim == 1:
            vecs = vecs.reshape(1, -1)
        return np.ascontiguousarray(vecs)

    def embed_query(self, text: str) -> np.ndarray:
        """Return a single (dim,) float32 vector for a query string."""
        return self.embed_passages([text])[0]