import os
from contextlib import asynccontextmanager

import numpy as np
import onnxruntime as ort
from fastapi import FastAPI, Header, HTTPException
from huggingface_hub import hf_hub_download
from pydantic import BaseModel
from transformers import AutoTokenizer

# int8 ONNX bge-m3, dense-only. Run raw via onnxruntime — no torch/optimum, so
# no export-path version conflicts. ~3x faster + ~half RAM vs fp32.
MODEL_NAME = os.getenv("EMBED_MODEL", "libryo-ai/BAAI-bge-m3-int8")
ONNX_FILE = os.getenv("EMBED_ONNX_FILE", "model.onnx")
EMBED_TOKEN = os.getenv("EMBED_TOKEN")  # optional shared secret
MAX_TOKENS = int(os.getenv("EMBED_MAX_TOKENS", "512"))
BATCH = int(os.getenv("EMBED_BATCH", "32"))
# HF free CPU exposes more cores than it gives you; let onnxruntime's default
# (0 = pick all) be overridable. Tune via env if 2 isn't fastest.
THREADS = int(os.getenv("EMBED_THREADS", "0"))

_session: ort.InferenceSession | None = None
_tokenizer = None
_input_names: set[str] = set()


@asynccontextmanager
async def lifespan(_: FastAPI):
    global _session, _tokenizer, _input_names
    _tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    path = hf_hub_download(MODEL_NAME, ONNX_FILE)
    opts = ort.SessionOptions()
    opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
    if THREADS:
        opts.intra_op_num_threads = THREADS
    _session = ort.InferenceSession(
        path, sess_options=opts, providers=["CPUExecutionProvider"]
    )
    _input_names = {i.name for i in _session.get_inputs()}
    _encode(["warmup"])  # pay first-run graph optimization now, not on a user request
    yield


app = FastAPI(lifespan=lifespan)


class EmbedRequest(BaseModel):
    texts: list[str]


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


@app.get("/")
def health():
    return {"status": "ok", "model": MODEL_NAME, "ready": _session is not None}


def _encode(texts: list[str]) -> list[list[float]]:
    out: list[list[float]] = []
    for i in range(0, len(texts), BATCH):
        chunk = texts[i : i + BATCH]
        enc = _tokenizer(
            chunk, padding=True, truncation=True,
            max_length=MAX_TOKENS, return_tensors="np",
        )
        # The ONNX graph requires token_type_ids but the xlm-roberta tokenizer
        # doesn't emit them (bge-m3 ignores them) → feed zeros for any required
        # input the tokenizer didn't produce.
        ids = enc["input_ids"]
        feed = {n: enc[n] if n in enc else np.zeros_like(ids) for n in _input_names}
        hidden = _session.run(None, feed)[0]  # (B, T, 1024) last_hidden_state
        # dense embedding = CLS token (position 0), then L2-normalize so cosine == dot.
        cls = hidden[:, 0]
        cls = cls / np.clip(np.linalg.norm(cls, axis=1, keepdims=True), 1e-12, None)
        out.extend(cls.astype(np.float32).tolist())
    return out


# sync `def` → FastAPI runs it in a threadpool, so the blocking ONNX inference
# does not stall the event loop.
@app.post("/embed", response_model=EmbedResponse)
def embed(req: EmbedRequest, authorization: str | None = Header(default=None)):
    if EMBED_TOKEN and authorization != f"Bearer {EMBED_TOKEN}":
        raise HTTPException(status_code=401, detail="unauthorized")
    if _session is None:
        raise HTTPException(status_code=503, detail="model loading")
    if not req.texts:
        raise HTTPException(status_code=400, detail="texts required")

    vecs = _encode(req.texts)
    return EmbedResponse(embeddings=vecs, dim=len(vecs[0]), model=MODEL_NAME)