#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json

from huggingface_hub import model_info, snapshot_download

from alpamayo_lite.config import load_config


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", default="configs/local_overfit.yaml")
    args = parser.parse_args()
    cfg = load_config(args.config)
    output_dir = cfg.resolve(cfg.paths.model_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    info = model_info(cfg.vlm.repo_id, files_metadata=True)
    total = info.safetensors.total if info.safetensors is not None else None
    if total is not None and not (700_000_000 <= total <= 900_000_000):
        raise RuntimeError(f"Unexpected model size for {cfg.vlm.repo_id}: {total:,} parameters")

    snapshot_download(
        repo_id=cfg.vlm.repo_id,
        local_dir=output_dir,
        allow_patterns=[
            "*.json",
            "*.txt",
            "*.safetensors",
        ],
    )
    provenance = {
        "repo_id": cfg.vlm.repo_id,
        "revision": info.sha,
        "safetensors_parameters": total,
        "pipeline_tag": info.pipeline_tag,
        "library_name": info.library_name,
    }
    (output_dir / "download_provenance.json").write_text(
        json.dumps(provenance, indent=2, ensure_ascii=False) + "\n", encoding="utf-8"
    )
    print(json.dumps(provenance, indent=2, ensure_ascii=False))


if __name__ == "__main__":
    main()