#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Download/cache customer demo models for offline service startup.
Run this once on a PC with network access. After it succeeds,
start_ai_service_customer_demo.bat will load models from ./models only.
"""
from __future__ import annotations

import os
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent
MODEL_ROOT = ROOT / "models"
ASR_DIR = MODEL_ROOT / "faster-whisper-base"
TRANSLATOR_DIR = MODEL_ROOT / "opus-mt-en-zh"


def _ensure_package() -> None:
    try:
        import huggingface_hub  # noqa: F401
    except Exception:
        print("[AI][ERROR] huggingface_hub is not installed.")
        print("[AI][FIX] Run install_customer_demo_deps.bat first.")
        raise SystemExit(2)


def _download(repo_id: str, local_dir: Path) -> None:
    from huggingface_hub import snapshot_download

    local_dir.mkdir(parents=True, exist_ok=True)
    print("[AI][MODEL] downloading {} -> {}".format(repo_id, local_dir))
    snapshot_download(
        repo_id=repo_id,
        local_dir=str(local_dir),
        local_dir_use_symlinks=False,
        resume_download=True,
    )
    print("[AI][MODEL] ready: {}".format(local_dir))


def main() -> int:
    _ensure_package()
    MODEL_ROOT.mkdir(parents=True, exist_ok=True)

    # ASR model used by faster-whisper --asr-model models\faster-whisper-base
    _download("Systran/faster-whisper-base", ASR_DIR)

    # English -> Chinese text translation used by Transformers AutoModelForSeq2SeqLM
    _download("Helsinki-NLP/opus-mt-en-zh", TRANSLATOR_DIR)

    print("")
    print("[AI][MODEL] All customer demo models are cached.")
    print("[AI][MODEL] You can now run start_ai_service_customer_demo.bat without internet access.")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
