# Trouve-KA — serveur d'embeddings
# Author: Simon-Pierre Boucher
# Contact: contact@spboucher.ai

"""Serveur HTTP d'embeddings multilingues.

Modèle : intfloat/multilingual-e5-small (384 dims) — espace vectoriel aligné
FR/EN : « thermopompe » et « heat pump » sont voisins sans traduction (§9).
e5 exige les préfixes "query: " / "passage: " — le paramètre `kind` les gère.

Volontairement autonome (fastapi + sentence-transformers seulement) : il tourne
dans son propre venv sur un node satellite, PAS dans l'image Docker principale
(torch pèse ~1 Go et n'a rien à faire dans le chemin de crawl).

Lancement :  uvicorn services.embedding.server:app --host 0.0.0.0 --port 8091
"""

import os
import threading
from typing import Literal

from fastapi import FastAPI
from pydantic import BaseModel, Field

MODEL_NAME = os.environ.get("EMBEDDING_MODEL", "intfloat/multilingual-e5-small")
RERANK_MODEL_NAME = os.environ.get(
    "RERANK_MODEL", "cross-encoder/mmarco-mMiniLMv2-L12-H384-v1"
)
MAX_TEXTS = int(os.environ.get("EMBEDDING_MAX_TEXTS", "256"))
MAX_CHARS = int(os.environ.get("EMBEDDING_MAX_CHARS", "1200"))

app = FastAPI(title="Trouve-KA Embeddings", docs_url=None, redoc_url=None)

_model = None
_reranker = None
_model_lock = threading.Lock()


def _get_model():
    global _model
    if _model is None:
        with _model_lock:
            if _model is None:
                from sentence_transformers import SentenceTransformer

                _model = SentenceTransformer(MODEL_NAME, device="cpu")
    return _model


def _get_reranker():
    global _reranker
    if _reranker is None:
        with _model_lock:
            if _reranker is None:
                from sentence_transformers import CrossEncoder

                _reranker = CrossEncoder(RERANK_MODEL_NAME, device="cpu")
    return _reranker


class EmbedRequest(BaseModel):
    texts: list[str] = Field(..., min_length=1, max_length=MAX_TEXTS)
    kind: Literal["query", "passage"] = "passage"


class EmbedResponse(BaseModel):
    vectors: list[list[float]]
    dim: int
    model: str


@app.get("/health")
def health() -> dict:
    return {"ok": True, "model": MODEL_NAME, "loaded": _model is not None}


@app.post("/embed", response_model=EmbedResponse)
def embed(req: EmbedRequest) -> EmbedResponse:
    model = _get_model()
    prefix = "query: " if req.kind == "query" else "passage: "
    texts = [prefix + t[:MAX_CHARS] for t in req.texts]
    vectors = model.encode(texts, normalize_embeddings=True, batch_size=64)
    return EmbedResponse(vectors=vectors.tolist(), dim=len(vectors[0]), model=MODEL_NAME)


class RerankRequest(BaseModel):
    query: str = Field(..., min_length=1, max_length=300)
    texts: list[str] = Field(..., min_length=1, max_length=64)


class RerankResponse(BaseModel):
    scores: list[float]
    model: str


@app.post("/rerank", response_model=RerankResponse)
def rerank(req: RerankRequest) -> RerankResponse:
    """Cross-encoder (requête, passage) → score de pertinence fin.

    Réservé au top-N (coût quadratique en attention) : l'API n'envoie que ~20
    candidats déjà filtrés par BM25+kNN.
    """
    model = _get_reranker()
    pairs = [(req.query, t[:MAX_CHARS]) for t in req.texts]
    scores = model.predict(pairs, batch_size=32)
    return RerankResponse(scores=[float(s) for s in scores], model=RERANK_MODEL_NAME)