"""
embedder.py — BGE-small-en-v1.5 via pure onnxruntime (no torch, no optimum).
ONNX weights downloaded from Xenova/bge-small-en-v1.5 at build time.
"""
from __future__ import annotations
import logging
import os
from typing import List
import numpy as np
from config import EMBED_BATCH_SIZE, EMBED_MODEL_NAME, EMBED_ONNX_DIR

logger = logging.getLogger(__name__)
_BGE_QUERY_PREFIX = "Represent this sentence: "


class Embedder:
    def __init__(self):
        self._session = None
        self._tokenizer = None
        self._ready = False

    def load(self) -> None:
        if self._ready:
            return
        import onnxruntime as ort
        from transformers import AutoTokenizer

        # model.onnx is placed flat in EMBED_ONNX_DIR by download_models.py
        model_path = os.path.join(EMBED_ONNX_DIR, "model.onnx")
        if not os.path.isfile(model_path):
            raise FileNotFoundError(
                f"ONNX model not found at {model_path}. "
                "Run download_models.py first."
            )

        logger.info("Loading ONNX embedder from %s", model_path)
        opts = ort.SessionOptions()
        opts.intra_op_num_threads = max(1, os.cpu_count() or 2)
        opts.inter_op_num_threads = max(1, os.cpu_count() or 2)
        self._session = ort.InferenceSession(
            model_path,
            sess_options=opts,
            providers=["CPUExecutionProvider"],
        )
        self._tokenizer = AutoTokenizer.from_pretrained(EMBED_ONNX_DIR)
        self._ready = True
        logger.info("Embedder ready (pure ONNX/onnxruntime, no torch)")

    def encode_chunks(self, texts: List[str]) -> np.ndarray:
        return self._encode(texts)

    def encode_query(self, text: str) -> np.ndarray:
        return self._encode([_BGE_QUERY_PREFIX + text])[0]

    def _encode(self, texts: List[str]) -> np.ndarray:
        if not self._ready:
            raise RuntimeError("Embedder.load() has not been called")
        all_embeddings = []
        for i in range(0, len(texts), EMBED_BATCH_SIZE):
            batch = texts[i: i + EMBED_BATCH_SIZE]
            enc = self._tokenizer(
                batch, padding=True, truncation=True,
                max_length=512, return_tensors="np",
            )
            inputs = {
                "input_ids": enc["input_ids"].astype(np.int64),
                "attention_mask": enc["attention_mask"].astype(np.int64),
            }
            if "token_type_ids" in enc:
                inputs["token_type_ids"] = enc["token_type_ids"].astype(np.int64)

            outputs = self._session.run(None, inputs)
            token_emb = outputs[0]  # (batch, seq, dim)
            mask = enc["attention_mask"][:, :, np.newaxis].astype(np.float32)
            summed = np.sum(token_emb * mask, axis=1)
            count = np.clip(mask.sum(axis=1), 1e-9, None)
            all_embeddings.append((summed / count).astype(np.float32))
        return np.vstack(all_embeddings)

    @property
    def is_ready(self) -> bool:
        return self._ready


embedder = Embedder()