"""本地 Sentence Transformers 模型加载与推理。"""
import asyncio
import json
import threading
from pathlib import Path
from typing import Any, Dict, List


class LocalEmbeddingService:
    """从 backend/models 安全加载本地向量模型，并复用已加载实例。"""

    MODELS_DIR = Path(__file__).resolve().parents[2] / "models"
    _models: Dict[str, Any] = {}
    _load_lock = threading.Lock()

    @classmethod
    def list_available_models(cls) -> List[Dict[str, Any]]:
        models_root = cls.MODELS_DIR.resolve()
        if not models_root.is_dir():
            return []

        models = []
        for model_path in sorted(models_root.iterdir(), key=lambda item: item.name.lower()):
            if not model_path.is_dir() or model_path.name.startswith("."):
                continue
            dimension = None
            pooling_config = model_path / "1_Pooling" / "config.json"
            try:
                with pooling_config.open("r", encoding="utf-8") as file:
                    dimension = json.load(file).get("word_embedding_dimension")
            except (OSError, ValueError, AttributeError):
                pass
            ready = any(
                (model_path / filename).is_file()
                for filename in ("model.safetensors", "pytorch_model.bin")
            )
            models.append({
                "name": model_path.name,
                "dimension": dimension,
                "ready": ready,
            })
        return models

    @classmethod
    def resolve_model_path(cls, model_name: str) -> Path:
        name = (model_name or "").strip()
        if not name:
            raise ValueError("缺少本地模型目录名称")

        models_root = cls.MODELS_DIR.resolve()
        requested = Path(name)
        candidate = requested.resolve() if requested.is_absolute() else (models_root / requested).resolve()
        if not candidate.is_relative_to(models_root):
            raise ValueError("本地模型必须位于 backend/models 目录中")
        if not candidate.is_dir():
            raise ValueError(f"本地模型目录不存在：{candidate.name}")
        return candidate

    @classmethod
    def _get_model(cls, model_path: Path):
        cache_key = str(model_path)
        with cls._load_lock:
            model = cls._models.get(cache_key)
            if model is not None:
                return model
            try:
                from sentence_transformers import SentenceTransformer
            except ImportError as exc:
                raise ValueError(
                    "本地向量模型依赖未安装，请执行 pip install -r requirements.txt"
                ) from exc
            model = SentenceTransformer(cache_key, local_files_only=True)
            cls._models[cache_key] = model
            return model

    @classmethod
    def _encode(cls, model_path: Path, text: str) -> List[float]:
        model = cls._get_model(model_path)
        vector = model.encode(
            text,
            normalize_embeddings=True,
            convert_to_numpy=True,
            show_progress_bar=False,
        )
        return [float(value) for value in vector.tolist()]

    @classmethod
    async def generate_embedding(cls, model_name: str, text: str) -> List[float]:
        model_path = cls.resolve_model_path(model_name)
        return await asyncio.to_thread(cls._encode, model_path, text)