"""Download the deployable PyTorch weights from Hugging Face Hub."""

from __future__ import annotations

import argparse
from pathlib import Path

from huggingface_hub import hf_hub_download


REPO_ID = "MilanMagik/AstroStarDetect"
ROOT = Path(__file__).resolve().parent

MODEL_FILES = {
    "denoise": ["StarDenoise/models/deep_denoise_cnn_AI3_6.pth"],
    "starnet": [
        "StarNetPyTorch/models/weights_G_RGB.pth",
        "StarNetPyTorch/models/weights_G_Greyscale.pth",
    ],
}


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Download AstroColorProcess model weights from Hugging Face."
    )
    parser.add_argument(
        "--only",
        choices=("all", "denoise", "starnet"),
        default="all",
        help="Select which model weights to download (default: all).",
    )
    parser.add_argument(
        "--revision",
        default="main",
        help="Hugging Face branch, tag, or commit (default: main).",
    )
    parser.add_argument(
        "--force",
        action="store_true",
        help="Download again even when a cached file is available.",
    )
    args = parser.parse_args()

    groups = MODEL_FILES if args.only == "all" else {args.only: MODEL_FILES[args.only]}
    for filenames in groups.values():
        for filename in filenames:
            local_path = hf_hub_download(
                repo_id=REPO_ID,
                filename=filename,
                revision=args.revision,
                local_dir=ROOT,
                force_download=args.force,
            )
            print(f"Downloaded: {local_path}")


if __name__ == "__main__":
    main()