"""CPU-only, unnormalized embeddings from multilingual-e5-small."""

from __future__ import annotations

import math
import os
import threading
from dataclasses import dataclass
from pathlib import Path
from typing import Literal

MODEL_ID = "intfloat/multilingual-e5-small"
MODEL_DIR = Path(os.getenv("E5_MODEL_DIR", "/opt/models/multilingual-e5-small"))
EXPECTED_DIMENSIONS = 384
MAX_INPUT_TOKENS = 512


class TooManyTokens(ValueError):
    def __init__(self, count: int, limit: int, input_index: int | None = None) -> None:
        self.count = count
        self.limit = limit
        self.input_index = input_index
        item = f"Input item {input_index}" if input_index is not None else "Input"
        super().__init__(f"{item} contains {count} tokens; maximum is {limit} (including prefix and special tokens)")


@dataclass(frozen=True)
class EmbeddingResult:
    tokens_count: int
    embedding: list[float]


class EmbeddingEngine:
    """Run official E5 attention-mask-aware mean pooling without L2 normalization."""

    def __init__(self, tokenizer: object, model: object, torch_module: object) -> None:
        self.tokenizer = tokenizer
        self.model = model
        self.torch = torch_module
        self._lock = threading.Lock()  # Prevent multiple CPU-heavy inferences in parallel.
        model_limit = int(getattr(model.config, "max_position_embeddings", MAX_INPUT_TOKENS))
        self.max_tokens = min(MAX_INPUT_TOKENS, model_limit)
        self.dimensions = int(model.config.hidden_size)
        if self.dimensions != EXPECTED_DIMENSIONS:
            raise RuntimeError(f"Unexpected embedding dimensions: {self.dimensions}")

    @classmethod
    def load(cls) -> EmbeddingEngine:
        import torch
        from transformers import AutoModel, AutoTokenizer

        threads = max(1, min(2, int(os.getenv("CPU_THREADS", "2"))))
        torch.set_num_threads(threads)
        tokenizer = AutoTokenizer.from_pretrained(str(MODEL_DIR), local_files_only=True)
        model = AutoModel.from_pretrained(str(MODEL_DIR), local_files_only=True)
        model.to("cpu")
        model.eval()  # Turn off dropout for repeatable output.
        return cls(tokenizer, model, torch)

    def encode(self, text: str, kind: Literal["query", "passage"]) -> EmbeddingResult:
        # The prefix is required by E5. The caller's original text is never modified.
        input_text = f"{kind}: {text}"
        with self._lock:
            encoded = self.tokenizer(
                input_text,
                return_tensors="pt",
                add_special_tokens=True,
                truncation=False,
            )
            tokens_count = int(encoded["attention_mask"].sum().item())
            if tokens_count > self.max_tokens:
                raise TooManyTokens(tokens_count, self.max_tokens)

            with self.torch.inference_mode():
                output = self.model(**encoded)
                token_embeddings = output.last_hidden_state
                attention_mask = encoded["attention_mask"].unsqueeze(-1)
                pooled = (token_embeddings * attention_mask).sum(dim=1) / attention_mask.sum(dim=1).clamp(min=1)
                vector = pooled.squeeze(0).cpu().tolist()

        if len(vector) != self.dimensions or any(not math.isfinite(value) for value in vector):
            raise RuntimeError("Model produced an invalid embedding")
        return EmbeddingResult(tokens_count=tokens_count, embedding=vector)