#!/usr/bin/env python3
"""Download exactly the three pinned upstream DSpark shards and verify them."""

from __future__ import annotations

import argparse
import hashlib
import shutil
import subprocess
from pathlib import Path

from recipe import SOURCE_FILES, SOURCE_REPOSITORY, SOURCE_REVISION


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        while chunk := handle.read(8 * 1024 * 1024):
            digest.update(chunk)
    return digest.hexdigest()


def download(destination: Path, hf_executable: str) -> None:
    destination.mkdir(parents=True, exist_ok=True)
    command = [
        hf_executable,
        "download",
        SOURCE_REPOSITORY,
        "--revision",
        SOURCE_REVISION,
        "--local-dir",
        str(destination),
    ]
    for name in SOURCE_FILES:
        command.extend(["--include", name])
    subprocess.run(command, check=True)
    for name, expected in SOURCE_FILES.items():
        path = destination / name
        if not path.is_file():
            raise FileNotFoundError(f"download did not produce {path}")
        size = path.stat().st_size
        digest = _sha256(path)
        if size != expected["size"] or digest != expected["sha256"]:
            raise ValueError(
                f"{name}: expected {expected['size']} bytes/{expected['sha256']}, "
                f"found {size}/{digest}"
            )
        print(f"[download] verified {name}: {size:,} bytes {digest}", flush=True)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--destination", type=Path, default=Path("sources"))
    parser.add_argument(
        "--hf",
        default=shutil.which("hf"),
        help="Path to the hf executable (the pinned public source needs no token)",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if not args.hf:
        raise SystemExit("hf CLI not found on PATH; pass --hf /path/to/hf")
    download(args.destination, args.hf)


if __name__ == "__main__":
    main()