#!/usr/bin/env python3
"""Tải model Gemma 4 QAT + mmproj + MTP từ Hugging Face."""

from __future__ import annotations

import argparse
from pathlib import Path

DEFAULT_SOURCE = "unsloth/gemma-4-12B-it-qat-GGUF"
DEFAULT_FILES = [
    "gemma-4-12B-it-qat-UD-Q4_K_XL.gguf",
    "mmproj-F16.gguf",
    "mtp-gemma-4-12B-it.gguf",
]


def download(
    dest: Path,
    source_repo: str = DEFAULT_SOURCE,
    files: list[str] | None = None,
    token: str | None = None,
) -> dict[str, Path]:
    from huggingface_hub import hf_hub_download

    dest.mkdir(parents=True, exist_ok=True)
    files = files or DEFAULT_FILES
    result: dict[str, Path] = {}

    for name in files:
        print(f"[download] {source_repo}/{name} -> {dest}/")
        path = hf_hub_download(
            repo_id=source_repo,
            filename=name,
            local_dir=str(dest),
            local_dir_use_symlinks=False,
            token=token,
        )
        result[name] = Path(path)
        print(f"  OK: {path}")

    return result


def main() -> int:
    p = argparse.ArgumentParser(description="Download Gemma 4 GGUF files for SRT translation")
    p.add_argument(
        "--dest",
        type=Path,
        default=Path("models"),
        help="Output directory (default: ./models)",
    )
    p.add_argument(
        "--source-repo",
        default=DEFAULT_SOURCE,
        help=f"Hugging Face model repo (default: {DEFAULT_SOURCE})",
    )
    p.add_argument(
        "--file",
        action="append",
        dest="files",
        help="Download specific file(s); repeat flag. Default: all 3 core files.",
    )
    p.add_argument("--token", default=None, help="HF token (optional, for gated models)")
    args = p.parse_args()

    download(args.dest, args.source_repo, args.files, args.token)
    print(f"\nDone. Models in: {args.dest.resolve()}")
    return 0


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