"""预拉模型权重到本地缓存，避免容器启动时才下载（首次启动慢）。

在容器内执行（docker run + 挂载 ./models volume），复用镜像里的 huggingface_hub。
读 config.yaml 拿 asr.model / translation.model，下载到 HF_HOME（= /models/huggingface）。

幂等：已下过的模型跳过（HF cache 命中检测）。

用法（经 scripts/prefetch_models.sh 包装）：
    ./scripts/prefetch_models.sh                  # 读 config.yaml
    ./scripts/prefetch_models.sh config.gpu.yaml  # 读指定配置
"""

from __future__ import annotations

import os
import sys
import time
from pathlib import Path


def log(msg: str) -> None:
    """带时间戳的日志（容器内无项目 logger，直接 print）。"""
    ts = time.strftime("%H:%M:%S")
    print(f"[{ts}] {msg}", flush=True)


def load_config(config_path: str) -> dict:
    """读 YAML 配置，返回 asr.model / translation.model。"""
    import yaml
    with open(config_path, encoding="utf-8") as f:
        return yaml.safe_load(f)


def resolve_asr_repo(model_name: str) -> str:
    """faster-whisper 模型名 -> HF repo。

    预定义名（tiny.en 等）走 _MODELS 映射；已是 repo 路径（org/name）直接用。
    """
    from faster_whisper.utils import _MODELS
    if "/" in model_name:
        return model_name  # 已是完整 repo 路径
    repo = _MODELS.get(model_name)
    if repo is None:
        raise ValueError(f"未知 faster-whisper 模型名：{model_name}（不在 _MODELS 里）")
    return repo


def is_cached(repo_id: str, hf_home: str) -> bool:
    """检测 HF cache 是否已有该模型（snapshot 目录存在且非空）。"""
    # HF 缓存布局：<hf_home>/hub/models--<org>--<name>/snapshots/<hash>/
    cache_dir = Path(hf_home) / "hub" / f"models--{repo_id.replace('/', '--')}"
    snapshots = cache_dir / "snapshots"
    if not snapshots.is_dir():
        return False
    return any(snap.is_dir() and any(snap.iterdir()) for snap in snapshots.iterdir())


def download_model(repo_id: str, hf_home: str) -> None:
    """用 huggingface_hub 下载模型所有文件到 HF cache。"""
    from huggingface_hub import snapshot_download
    log(f"  下载 {repo_id} ...")
    snapshot_download(
        repo_id=repo_id,
        local_dir=None,  # 走标准 cache
        cache_dir=Path(hf_home) / "hub",
    )


def prefetch_asr(model_name: str, hf_home: str) -> None:
    """预拉 faster-whisper ASR 模型。"""
    repo = resolve_asr_repo(model_name)
    log(f"[ASR] model={model_name}  repo={repo}")
    if is_cached(repo, hf_home):
        log(f"  ✓ 已缓存，跳过")
        return
    download_model(repo, hf_home)
    log(f"  ✓ 完成")


def prefetch_translation(model_name: str, hf_home: str) -> None:
    """预拉翻译模型（transformers pipeline 用的 HF repo）。"""
    log(f"[翻译] model={model_name}")
    if is_cached(model_name, hf_home):
        log(f"  ✓ 已缓存，跳过")
        return
    download_model(model_name, hf_home)
    log(f"  ✓ 完成")


def main() -> int:
    config_path = os.environ.get("CONFIG_PATH", "/app/config.yaml")
    if len(sys.argv) > 1:
        config_path = sys.argv[1]
    hf_home = os.environ.get("HF_HOME", "/models/huggingface")

    log(f"配置文件: {config_path}")
    log(f"HF 缓存目录: {hf_home}")

    if not Path(config_path).is_file():
        log(f"✗ 配置文件不存在: {config_path}")
        return 1

    cfg = load_config(config_path)
    asr_model = cfg.get("asr", {}).get("model", "tiny.en")
    tr_model = cfg.get("translation", {}).get("model", "Helsinki-NLP/opus-mt-en-zh")

    log(f"ASR 模型: {asr_model}")
    log(f"翻译模型: {tr_model}")
    log("-" * 50)

    # 预拉 ASR
    try:
        prefetch_asr(asr_model, hf_home)
    except Exception as e:
        log(f"✗ ASR 模型预拉失败: {e}")
        return 2

    # 预拉翻译
    try:
        prefetch_translation(tr_model, hf_home)
    except Exception as e:
        log(f"✗ 翻译模型预拉失败: {e}")
        return 3

    log("-" * 50)
    log(f"全部完成。缓存位于: {hf_home}")
    # 打印缓存大小
    try:
        total = sum(f.stat().st_size for f in Path(hf_home).rglob("*") if f.is_file())
        log(f"缓存总大小: {total / 1024 / 1024 / 1024:.2f} GB")
    except Exception:
        pass
    return 0


if __name__ == "__main__":
    sys.exit(main())