app/plugins/model_probe.py

"""
模型 Provider 状态探测模块

为 CPU-only HuggingFace Space 环境提供模型可用性探测能力。
不强制加载完整大模型，仅报告 provider 在当前环境的状态。
"""

import logging
import time
import os
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional, Dict, Any, List
from enum import Enum

from pydantic import BaseModel, Field

from app.config.settings import settings

logger = logging.getLogger(__name__)


class ProviderAvailability(str, Enum):
    """Provider 可用状态"""
    AVAILABLE = "available"           # 模型依赖就绪，预期可运行
    UNAVAILABLE = "unavailable"       # 模型依赖缺失或环境不支持
    NOT_VERIFIED = "not_verified"     # 无法验证（如缺少凭证、资源限制等）


class ProviderStatus(BaseModel):
    """单个模型 Provider 的状态信息

    供前端和 API 展示模型可用性，不强制加载完整模型。
    """

    provider: str = Field(..., description="Provider 名称，如 whisper_small")
    model_id: str = Field(..., description="HuggingFace 模型 ID")
    status: ProviderAvailability = Field(
        default=ProviderAvailability.NOT_VERIFIED,
        description="可用状态：available/unavailable/not_verified"
    )
    cache_path: Optional[str] = Field(None, description="模型缓存路径")
    reason: Optional[str] = Field(None, description="状态原因说明")
    last_probe_at: Optional[datetime] = Field(None, description="最后探测时间")
    probe_duration_ms: Optional[float] = Field(None, description="探测耗时（毫秒）")
    extra: Dict[str, Any] = Field(default_factory=dict, description="额外信息")


class ModelProbeResult(BaseModel):
    """模型探测汇总结果

    包含所有 provider 的状态列表和探测摘要。
    """

    providers: List[ProviderStatus] = Field(default_factory=list)
    probe_time: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
    cpu_only: bool = True  # 当前环境是否为 CPU-only
    hf_home: str = ""
    hf_cache_dir: str = ""


def probe_cpu_only_environment() -> Dict[str, Any]:
    """探测 CPU-only 环境的基本信息

    不加载任何模型，仅检查环境变量、缓存目录和 GPU 可用性。

    Returns:
        Dict 包含 cpu_only 标志、hf_home、缓存路径等信息
    """
    result = {
        "cpu_only": True,
        "hf_home": str(settings.HF_HOME),
        "hf_cache_dir": str(settings.HF_CACHE_DIR),
        "hf_home_exists": settings.HF_HOME.exists(),
        "hf_cache_dir_exists": settings.HF_CACHE_DIR.exists(),
    }

    # 检查 GPU 可用性（不导入 torch，只检查环境变量）
    cuda_visible = os.getenv("CUDA_VISIBLE_DEVICES", "")
    if cuda_visible and cuda_visible != "-1":
        # 有 GPU 环境变量设置，但本项目目标是 CPU-only
        result["cuda_visible_devices"] = cuda_visible
        result["note"] = "检测到 CUDA 环境变量，但项目以 CPU-only 模式运行"

    # 检查 HuggingFace 环境变量
    hf_home_env = os.getenv("HF_HOME", "")
    if hf_home_env:
        result["hf_home_env"] = hf_home_env

    hf_token = os.getenv("HF_TOKEN", "") or os.getenv("HUGGINGFACE_HUB_TOKEN", "")
    result["hf_token_configured"] = bool(hf_token)

    # 检查 HuggingFace Hub 是否可用（尝试 import，不调用网络）
    try:
        import huggingface_hub  # noqa: F401
        result["huggingface_hub_available"] = True
        result["huggingface_hub_version"] = huggingface_hub.__version__
    except ImportError:
        result["huggingface_hub_available"] = False
        result["huggingface_hub_version"] = None

    return result


def probe_provider(
    provider: str,
    model_id: str,
    check_dependencies: List[str],
    check_cache: bool = True,
) -> ProviderStatus:
    """探测单个 Provider 的可用性

    执行轻量检查：依赖包是否可导入、模型缓存是否存在。
    不加载完整模型权重。

    Args:
        provider: Provider 名称
        model_id: HuggingFace 模型 ID
        check_dependencies: 需要检查的 Python 依赖包列表
        check_cache: 是否检查模型缓存

    Returns:
        ProviderStatus 包含探测结果
    """
    start_time = time.perf_counter()
    reasons = []
    status = ProviderAvailability.AVAILABLE

    # 检查依赖包
    missing_deps = []
    for dep in check_dependencies:
        try:
            __import__(dep)
        except ImportError:
            missing_deps.append(dep)

    if missing_deps:
        status = ProviderAvailability.UNAVAILABLE
        reasons.append(f"缺失依赖: {', '.join(missing_deps)}")

    # 检查模型缓存
    cache_dir_name = "models--" + model_id.replace("/", "--")
    cache_path = settings.HF_CACHE_DIR / cache_dir_name
    if check_cache:
        if cache_path.exists():
            # 检查缓存中是否有实际模型文件（blobs 或 snapshots）
            has_blobs = (cache_path / "blobs").exists() and any((cache_path / "blobs").iterdir())
            has_snapshots = (cache_path / "snapshots").exists() and any((cache_path / "snapshots").iterdir())
            if not has_blobs and not has_snapshots:
                # 缓存目录存在但无模型文件，标记为需要下载
                reasons.append("模型缓存目录存在但无模型文件，首次使用需下载")
        else:
            reasons.append(f"模型未缓存，首次使用将从 HuggingFace Hub 下载")

    # 检查 CPU-only 环境的额外限制
    cpu_env = probe_cpu_only_environment()
    if not cpu_env.get("huggingface_hub_available"):
        status = ProviderAvailability.UNAVAILABLE
        reasons.append("huggingface_hub 包不可用，无法访问 HuggingFace Hub")

    # 如果缺少依赖且无法下载，标记为 unavailable
    if missing_deps and not cpu_env.get("huggingface_hub_available"):
        status = ProviderAvailability.UNAVAILABLE

    duration_ms = (time.perf_counter() - start_time) * 1000

    return ProviderStatus(
        provider=provider,
        model_id=model_id,
        status=status,
        cache_path=str(cache_path) if check_cache else None,
        reason="; ".join(reasons) if reasons else "依赖和缓存检查通过",
        last_probe_at=datetime.now(timezone.utc),
        probe_duration_ms=round(duration_ms, 2),
        extra={
            "missing_dependencies": missing_deps,
            "cpu_only": cpu_env["cpu_only"],
        },
    )


def probe_all_providers() -> ModelProbeResult:
    """探测所有已知 Provider 的可用性

    汇总 ASR、OCR 等所有模型 provider 的状态。

    Returns:
        ModelProbeResult 包含所有 provider 状态
    """
    cpu_env = probe_cpu_only_environment()

    providers: List[ProviderStatus] = []

    # ASR providers（均基于 transformers pipeline）
    providers.append(probe_provider(
        provider="whisper_small",
        model_id="openai/whisper-small",
        check_dependencies=["transformers", "torch"],
    ))

    providers.append(probe_provider(
        provider="glm_asr_nano",
        model_id="zai-org/GLM-ASR-Nano-2512",
        check_dependencies=["transformers", "torch"],
    ))

    # OCR provider
    providers.append(probe_provider(
        provider="ppocrv6",
        model_id="PaddlePaddle/PP-OCRv6_medium_det_safetensors",
        check_dependencies=["paddleocr", "paddle"],
    ))

    # 补充 OCR rec 模型
    providers.append(probe_provider(
        provider="ppocrv6_rec",
        model_id="PaddlePaddle/PP-OCRv6_medium_rec_safetensors",
        check_dependencies=["paddleocr", "paddle"],
    ))

    return ModelProbeResult(
        providers=providers,
        cpu_only=cpu_env["cpu_only"],
        hf_home=cpu_env["hf_home"],
        hf_cache_dir=cpu_env["hf_cache_dir"],
    )