"""文本向量化服务 - 只支持密集向量（由 Milvus 2.5+ 原生支持中文分词与 BM25 全文检索）"""
import os
import threading

import torch
import torch.nn.functional as F
from transformers import AutoModel, AutoTokenizer


def _mean_pooling(last_hidden_state: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
    """对 last_hidden_state 按 attention_mask 做 mean pooling。"""
    mask = attention_mask.unsqueeze(-1).expand(last_hidden_state.size()).float()
    sum_embeddings = torch.sum(last_hidden_state * mask, dim=1)
    sum_mask = torch.clamp(mask.sum(dim=1), min=1e-9)
    return sum_embeddings / sum_mask


def _create_dense_embedder() -> tuple[AutoTokenizer, AutoModel, str]:
    model_name = os.getenv("EMBEDDING_MODEL", "BAAI/bge-m3")
    device = os.getenv("EMBEDDING_DEVICE", "cpu")
    local_only = os.getenv("EMBEDDING_LOCAL_FILES_ONLY", "true").lower() in ("1", "true", "yes")

    tokenizer = AutoTokenizer.from_pretrained(model_name, local_files_only=local_only)
    model = AutoModel.from_pretrained(model_name, local_files_only=local_only)
    model.to(device)
    model.eval()
    return tokenizer, model, device


class EmbeddingService:
    """文本向量化服务 - 密集向量本地模型（基于 transformers，避免 sentence_transformers 拉入 pyarrow）"""

    # 与 BAAI/bge-m3 推荐保持一致；800 中文字符约 300~400 tokens，512 足够
    MAX_LENGTH = 512

    def __init__(self, state_path=None):
        self._tokenizer = None
        self._model = None
        self._device = None
        self._lock = threading.Lock()

    def _get_embedder(self):
        if self._model is None:
            with self._lock:
                if self._model is None:
                    self._tokenizer, self._model, self._device = _create_dense_embedder()
        return self._tokenizer, self._model, self._device

    def get_embeddings(self, texts: list[str]) -> list[list[float]]:
        if not texts:
            return []
        try:
            tokenizer, model, device = self._get_embedder()
            with torch.no_grad():
                inputs = tokenizer(
                    texts,
                    padding=True,
                    truncation=True,
                    max_length=self.MAX_LENGTH,
                    return_tensors="pt",
                )
                inputs = {k: v.to(device) for k, v in inputs.items()}
                outputs = model(**inputs)
                embeddings = _mean_pooling(outputs.last_hidden_state, inputs["attention_mask"])
                embeddings = F.normalize(embeddings, p=2, dim=1)
            return embeddings.cpu().numpy().tolist()
        except Exception as e:
            raise Exception(f"本地密集嵌入模型调用失败: {str(e)}") from e


# 全进程唯一实例（延迟加载模型，避免启动时立即加载）
embedding_service = EmbeddingService()
