from __future__ import annotations

import argparse
import subprocess
import sys
from pathlib import Path

from huggingface_hub import HfApi, hf_hub_url

from common import load_model_card


def selected_files(repo_id: str, allow_patterns: list[str] | None) -> list[tuple[str, int | None]]:
    api = HfApi()
    info = api.model_info(repo_id=repo_id, files_metadata=True)
    files = []
    for sibling in info.siblings:
        name = sibling.rfilename
        if name == ".gitattributes":
            continue
        if allow_patterns and name not in allow_patterns:
            continue
        files.append((name, sibling.size))
    return files


def run_curl(url: str, output: Path) -> int:
    output.parent.mkdir(parents=True, exist_ok=True)
    command = [
        "curl.exe",
        "-L",
        "-C",
        "-",
        "--retry",
        "10",
        "--retry-all-errors",
        "--retry-delay",
        "2",
        "--connect-timeout",
        "30",
        "--speed-time",
        "60",
        "--speed-limit",
        "1024",
        "-o",
        str(output),
        url,
    ]
    return subprocess.run(command, check=False).returncode


def main() -> int:
    parser = argparse.ArgumentParser(description="Download files from a Hugging Face model card using curl resume.")
    parser.add_argument("card_id")
    parser.add_argument("--local-dir", required=True)
    parser.add_argument("--allow-file", action="append", dest="allow_files")
    parser.add_argument("--max-files", type=int)
    args = parser.parse_args()

    card = load_model_card(args.card_id)
    repo_id = card.get("upstream_id")
    if not repo_id:
        print(f"{args.card_id}: upstream_id is required", file=sys.stderr)
        return 1

    files = selected_files(repo_id, args.allow_files)
    if args.max_files:
        files = files[: args.max_files]

    local_dir = Path(args.local_dir)
    failed = []
    for name, expected_size in files:
        output = local_dir / name
        current_size = output.stat().st_size if output.exists() else 0
        if expected_size and current_size == expected_size:
            print(f"OK {name}: {current_size}/{expected_size}")
            continue
        if expected_size and current_size > expected_size:
            failed.append((name, current_size, expected_size, "oversized"))
            continue

        print(f"GET {name}: {current_size}/{expected_size or '?'}")
        url = hf_hub_url(repo_id, name)
        code = run_curl(url, output)
        current_size = output.stat().st_size if output.exists() else 0
        if code != 0 or (expected_size and current_size != expected_size):
            failed.append((name, current_size, expected_size, code))

    if failed:
        print("Incomplete downloads:", file=sys.stderr)
        for name, current_size, expected_size, code in failed:
            print(f"- {name}: {current_size}/{expected_size or '?'} curl_code={code}", file=sys.stderr)
        return 1

    print(f"Downloaded {len(files)} file(s) for {args.card_id}.")
    return 0


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