"""BGE 模型的本地快照解析与进程级复用。"""

from pathlib import Path
from threading import Lock
from typing import Any

from huggingface_hub.constants import HF_HUB_CACHE

_MODEL_CONFIG_FILES = ("config.json", "tokenizer_config.json")
_TOKENIZER_FILES = ("tokenizer.json", "sentencepiece.bpe.model")
_WEIGHT_FILES = (
    "model.safetensors",
    "model.safetensors.index.json",
    "pytorch_model.bin",
    "pytorch_model.bin.index.json",
)
_BGE_M3_MODELS: dict[tuple[str, str], Any] = {}
_BGE_M3_MODELS_LOCK = Lock()


def resolve_local_huggingface_snapshot(
    model_name_or_path: str,
    *,
    cache_dir: Path | None = None,
) -> str:
    """优先返回完整的本地模型快照，避免已下载模型仍发起联网探测。"""

    configured_path = Path(model_name_or_path).expanduser()
    if configured_path.is_dir():
        return str(configured_path.resolve())
    if "/" not in model_name_or_path:
        return model_name_or_path

    hub_cache = (cache_dir or Path(HF_HUB_CACHE)).expanduser()
    repository_dir = hub_cache / f"models--{model_name_or_path.replace('/', '--')}"
    candidates: list[Path] = []
    main_ref = repository_dir / "refs" / "main"
    if main_ref.is_file():
        revision = main_ref.read_text(encoding="utf-8").strip()
        if revision:
            candidates.append(repository_dir / "snapshots" / revision)

    snapshots_dir = repository_dir / "snapshots"
    if snapshots_dir.is_dir():
        remaining = sorted(
            (item for item in snapshots_dir.iterdir() if item.is_dir()),
            key=lambda item: item.stat().st_mtime,
            reverse=True,
        )
        candidates.extend(item for item in remaining if item not in candidates)

    for snapshot in candidates:
        if _is_complete_transformer_snapshot(snapshot):
            return str(snapshot.resolve())
    return model_name_or_path


def load_bge_m3_model(
    model_name_or_path: str,
    device: str,
    *,
    cache_dir: Path | None = None,
) -> Any:
    """加载并复用同一模型路径、同一设备上的 BGE-M3 实例。"""

    resolved_path = resolve_local_huggingface_snapshot(
        model_name_or_path,
        cache_dir=cache_dir,
    )
    cache_key = (resolved_path, device)
    with _BGE_M3_MODELS_LOCK:
        model = _BGE_M3_MODELS.get(cache_key)
        if model is None:
            from FlagEmbedding import BGEM3FlagModel  # type: ignore[import-untyped]

            model = BGEM3FlagModel(
                resolved_path,
                devices=device,
                use_fp16=device != "cpu",
                return_dense=True,
                return_sparse=False,
                return_colbert_vecs=False,
            )
            _BGE_M3_MODELS[cache_key] = model
        return model


def clear_bge_m3_model_cache() -> None:
    """清理进程级模型缓存，仅供测试和受控重载使用。"""

    with _BGE_M3_MODELS_LOCK:
        _BGE_M3_MODELS.clear()


def _is_complete_transformer_snapshot(snapshot: Path) -> bool:
    return (
        snapshot.is_dir()
        and all((snapshot / filename).is_file() for filename in _MODEL_CONFIG_FILES)
        and any((snapshot / filename).is_file() for filename in _TOKENIZER_FILES)
        and any((snapshot / filename).is_file() for filename in _WEIGHT_FILES)
    )
