"""本地嵌入(embedding)服务。

对外提供 OpenAI 兼容的 `POST /v1/embeddings`,任何支持自定义 OpenAI 兼容端点
的框架都能直接对接,不需要改 SDK。

与 local-remarker 的区别:
  * 端点是 OpenAI 标准(而非 SiliconFlow 私有格式),调用方不会拼错路径;
  * 后处理是 CLS 池化 + L2 归一化(bge 系列官方做法),而不是取 logit;
  * 完全不依赖 Ollama —— 模型以 ONNX 形式直接跑在 onnxruntime 上。
"""

import os
import time
from typing import Any

import numpy as np
import onnxruntime as ort
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
from tokenizers import Tokenizer

# ---------------------------------------------------------------- 配置
MODEL_DIR = os.environ.get("MODEL_DIR", "/models")
ONNX_FILE = os.environ.get("ONNX_FILE", "model_int8.onnx")
TOKENIZER_FILE = os.environ.get("TOKENIZER_FILE", "tokenizer.json")
MODEL_NAME = os.environ.get("MODEL_NAME", "bge-m3")
# bge-m3 原生支持 8192 上下文,但 CPU 上长序列非常慢,默认按 512 截断
MAX_LEN = int(os.environ.get("MAX_LEN", "512"))
BATCH_SIZE = int(os.environ.get("BATCH_SIZE", "16"))
PAD_ID = int(os.environ.get("PAD_ID", "1"))
NUM_THREADS = os.environ.get("NUM_THREADS", "") or None

app = FastAPI(title="local-embedding", version="1.0")


def _build_session() -> tuple[ort.InferenceSession, Tokenizer]:
    onnx_path = os.path.join(MODEL_DIR, ONNX_FILE)
    tok_path = os.path.join(MODEL_DIR, TOKENIZER_FILE)
    for path in (onnx_path, tok_path):
        if not os.path.isfile(path):
            raise FileNotFoundError(
                f"缺少模型文件: {path}(entrypoint 应先执行 download_model.py)"
            )

    options = ort.SessionOptions()
    options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
    options.intra_op_num_threads = int(NUM_THREADS) if NUM_THREADS else (os.cpu_count() or 4)

    session = ort.InferenceSession(
        onnx_path, sess_options=options, providers=["CPUExecutionProvider"]
    )

    tokenizer = Tokenizer.from_file(tok_path)
    tokenizer.enable_truncation(max_length=MAX_LEN)
    return session, tokenizer


SESSION, TOKENIZER = _build_session()
INPUT_NAMES = {i.name for i in SESSION.get_inputs()}
OUTPUT_NAMES = [o.name for o in SESSION.get_outputs()]


def _pool(outputs: list[np.ndarray]) -> np.ndarray:
    """从 ONNX 输出取句向量并做 L2 归一化。

    不同导出方式的第一个输出形状不一致,两种都兼容:
      * [batch, hidden]      —— 已池化好的句向量
      * [batch, seq, hidden] —— token 级输出,取 [CLS](第 0 个 token)

    bge 系列官方推荐 CLS 池化 + L2 归一化,归一化后余弦相似度与内积等价。
    """
    hidden = outputs[0]
    if hidden.ndim == 3:
        hidden = hidden[:, 0, :]  # [CLS]
    norms = np.linalg.norm(hidden, axis=1, keepdims=True)
    norms[norms == 0] = 1.0
    return hidden / norms


def _forward(rows: list[tuple[list[int], list[int], list[int]]]) -> np.ndarray:
    """把 (ids, attention_mask, type_ids) 组成 batch 跑一次推理,返回归一化向量。"""
    width = max(len(ids) for ids, _, _ in rows)
    size = len(rows)

    input_ids = np.full((size, width), PAD_ID, dtype=np.int64)
    attention_mask = np.zeros((size, width), dtype=np.int64)
    token_type_ids = np.zeros((size, width), dtype=np.int64)

    for row, (ids, mask, types) in enumerate(rows):
        length = len(ids)
        input_ids[row, :length] = ids
        attention_mask[row, :len(mask)] = mask
        token_type_ids[row, :len(types)] = types

    feeds = {"input_ids": input_ids, "attention_mask": attention_mask}
    if "token_type_ids" in INPUT_NAMES:
        feeds["token_type_ids"] = token_type_ids

    return _pool(SESSION.run(None, feeds)).astype(np.float32)


def _encode(texts: list[str], token_ids: dict[int, list[int]]) -> list[tuple]:
    """统一编码入口。token_ids 里的下标表示该项已是 token id 序列。"""
    rows, tokens = [], 0
    for index, text in enumerate(texts):
        if index in token_ids:
            ids = token_ids[index][:MAX_LEN]
            mask = [1] * len(ids)
            types = [0] * len(ids)
        else:
            encoding = TOKENIZER.encode(text)
            ids, mask, types = encoding.ids, encoding.attention_mask, encoding.type_ids
        rows.append((ids, mask, types))
        tokens += len(ids)
    return rows, tokens


def _invalid(message: str) -> JSONResponse:
    return JSONResponse(
        {"error": {"message": message, "type": "invalid_request_error"}},
        status_code=400,
    )


async def _handle(request: Request) -> JSONResponse:
    try:
        body = await request.json()
    except Exception:  # noqa: BLE001
        return _invalid("invalid json body")

    if "input" not in body:
        return _invalid("input is required")

    raw = body["input"]
    items = [raw] if isinstance(raw, str) else raw
    if not isinstance(items, list) or not items:
        return _invalid("input must be a non-empty string or array")

    # 允许混排字符串与 token id 数组(OpenAI 规范允许后者)
    texts, token_ids = [], {}
    for position, item in enumerate(items):
        if isinstance(item, str):
            texts.append(item)
        elif isinstance(item, list) and item and all(isinstance(x, int) for x in item):
            token_ids[position] = item
            texts.append("")
        else:
            return _invalid("input items must be strings or token id arrays")

    started = time.time()
    try:
        vectors, tokens = [], 0
        for start in range(0, len(texts), BATCH_SIZE):
            chunk = texts[start:start + BATCH_SIZE]
            offset = {k - start: v for k, v in token_ids.items() if start <= k < start + len(chunk)}
            rows, chunk_tokens = _encode(chunk, offset)
            vectors.append(_forward(rows))
            tokens += chunk_tokens
        matrix = np.concatenate(vectors, axis=0)
    except Exception as exc:  # noqa: BLE001
        return JSONResponse(
            {"error": {"message": str(exc), "type": "server_error"}}, status_code=500
        )

    return JSONResponse({
        "object": "list",
        "data": [
            {"object": "embedding", "index": i, "embedding": matrix[i].tolist()}
            for i in range(matrix.shape[0])
        ],
        "model": body.get("model") or MODEL_NAME,
        "usage": {"prompt_tokens": tokens, "total_tokens": tokens},
        "_meta": {
            "latency_ms": int((time.time() - started) * 1000),
            "dimension": int(matrix.shape[1]),
        },
    })


# OpenAI 标准路径 + 常见别名
for _path in ("/v1/embeddings", "/embeddings", "/v1/embed", "/embed"):
    app.post(_path)(_handle)


@app.get("/health")
def health() -> dict[str, Any]:
    return {"status": "ok", "model": MODEL_NAME}


@app.get("/info")
def info() -> dict[str, Any]:
    return {
        "model": MODEL_NAME,
        "onnx_file": ONNX_FILE,
        "model_dir": MODEL_DIR,
        "max_len": MAX_LEN,
        "batch_size": BATCH_SIZE,
        "input_names": sorted(INPUT_NAMES),
        "output_names": OUTPUT_NAMES,
    }


# 部分客户端启动时会先探测 /v1/models
@app.get("/v1/models")
def list_models() -> dict[str, Any]:
    return {
        "object": "list",
        "data": [{
            "id": MODEL_NAME,
            "object": "model",
            "created": int(time.time()),
            "owned_by": "local",
        }],
    }