"""Explicit, resumable, revision-pinned upstream download with hash verification."""
import argparse
import os
import hashlib
import json
from pathlib import Path
os.environ.setdefault("HF_HUB_DISABLE_XET", "1")
from huggingface_hub import HfApi, hf_hub_download

ROOT = Path(__file__).resolve().parents[1]
VARIANTS = ("5M", "28M", "123M", "209M")


def sha256(path):
    h = hashlib.sha256()
    with Path(path).open("rb") as f:
        for chunk in iter(lambda: f.read(8 * 1024 * 1024), b""):
            h.update(chunk)
    return h.hexdigest()


def download(variant):
    repo = f"mistralai/Timestral-{variant}"
    api = HfApi()
    pins_path = ROOT / "pins.json"
    pins = json.loads(pins_path.read_text()) if pins_path.exists() else {}
    revision = pins.get(variant, {}).get("revision")
    info = api.model_info(repo, revision=revision, files_metadata=True)
    path = ROOT / "models" / "source" / variant
    path.mkdir(parents=True, exist_ok=True)
    print(f"Downloading {repo}@{info.sha} to {path}", flush=True)
    inventory = {}
    for name in ("config.json", "model.safetensors", "LICENSE", "README.md"):
        dest = Path(hf_hub_download(repo, name, revision=info.sha, local_dir=path))
        meta = next(s for s in info.siblings if s.rfilename == name)
        if meta.lfs:
            expected = meta.lfs.sha256
            actual = sha256(dest)
        else:
            expected = meta.blob_id
            data = dest.read_bytes()
            actual = hashlib.sha1(f"blob {len(data)}\0".encode() + data).hexdigest()
        if actual != expected:
            raise ValueError(f"Hub blob hash mismatch for {repo}/{name}")
        inventory[name] = {"sha256": sha256(dest), "bytes": dest.stat().st_size, "hub_blob": expected}
        pinned = pins.get(variant, {}).get("files", {}).get(name)
        if pinned is not None and pinned != inventory[name]:
            raise ValueError(f"pinned source hash mismatch: {repo}/{name}")
        print(f"  {name}: {dest.stat().st_size:,} bytes, verified", flush=True)
    record = {"repo_id": repo, "revision": info.sha, "files": inventory}
    (path / "source.json").write_text(json.dumps(record, indent=2) + "\n")
    return path


if __name__ == "__main__":
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("variants", nargs="*", choices=VARIANTS, default=list(VARIANTS))
    args = p.parse_args()
    for size in args.variants:
        download(size)