from __future__ import annotations

import argparse
import os
import sys

from huggingface_hub import snapshot_download

from plan_hf_downloads import remote_plan
from common import iter_model_card_paths, read_yaml_mapping


def main() -> int:
    parser = argparse.ArgumentParser(description="Download missing Hugging Face model weights from model cards.")
    parser.add_argument("--model", action="append", dest="models", help="Model card id to download. Can be repeated.")
    parser.add_argument("--include-ready", action="store_true")
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()

    selected = set(args.models or [])
    token = os.environ.get("HF_TOKEN") or None
    plans = []
    for card_path in iter_model_card_paths():
        card = read_yaml_mapping(card_path)
        card_id = str(card.get("id") or "")
        if selected and card_id not in selected:
            continue
        plan = remote_plan(card_path, include_ready=args.include_ready)
        if plan and plan["missing_bytes"]:
            plans.append(plan)

    if not plans:
        print("No missing downloads.")
        return 0

    for plan in plans:
        missing_files = [item["name"] for item in plan["files"] if not item["complete"]]
        if not missing_files:
            continue
        print("")
        print(f"== {plan['id']} ==")
        print(f"Repository: {plan['repo_id']}")
        print(f"Destination: {plan['storage_path']}")
        print(f"Missing: {plan['missing_gb']} GB")
        for name in missing_files:
            print(f"- {name}")
        if args.dry_run:
            continue
        snapshot_download(
            repo_id=plan["repo_id"],
            local_dir=plan["storage_path"],
            token=token,
            allow_patterns=missing_files,
            ignore_patterns=["*.msgpack", "*.h5", "*.ot", "*.tflite", "*.onnx", "*.pb", "*.ckpt"],
        )

    return 0


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