from __future__ import annotations

import argparse
import json
from pathlib import Path
import sys

MODEL_ID = "Wan-AI/Wan2.2-Animate-2-14B-Distilled-Diffusers"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--repo", default=MODEL_ID)
    ap.add_argument("--output", default="/content/Wan2_2_Animate_Repo")
    ap.add_argument("--max-workers", type=int, default=8)
    args = ap.parse_args()

    from huggingface_hub import HfApi, snapshot_download

    info = HfApi().model_info(args.repo)
    revision = info.sha
    print("Model   :", args.repo)
    print("Revision:", revision)
    print("Output  :", args.output)
    print("Downloading complete Diffusers model (~45.9 GB; transformer ~32.8 GB)...")

    snapshot_download(
        repo_id=args.repo,
        revision=revision,
        local_dir=args.output,
        max_workers=args.max_workers,
    )

    root = Path(args.output)
    required = [
        root / "model_index.json",
        root / "transformer" / "config.json",
        root / "transformer" / "diffusion_pytorch_model.safetensors.index.json",
        root / "text_encoder",
        root / "tokenizer",
        root / "image_encoder",
        root / "vae",
    ]
    for p in required:
        if not p.exists():
            raise FileNotFoundError(p)

    index = json.loads((root / "transformer" / "diffusion_pytorch_model.safetensors.index.json").read_text())
    wm = index["weight_map"]
    shards = sorted(set(wm.values()))
    if len(wm) != 1303:
        raise RuntimeError(f"expected 1303 transformer tensors, got {len(wm)}")
    if len(shards) != 4:
        raise RuntimeError(f"expected four transformer shards, got {shards}")
    for shard in shards:
        if not (root / "transformer" / shard).is_file():
            raise FileNotFoundError(root / "transformer" / shard)

    mi = json.loads((root / "model_index.json").read_text())
    if mi.get("_class_name") != "WanAnimate2Pipeline":
        raise RuntimeError(f"wrong pipeline class: {mi.get('_class_name')!r}")

    (root / "ORBITQUANT_SOURCE_REVISION.txt").write_text(revision + "\n")
    print("\nMODEL DOWNLOAD PASS")
    print("  1303 transformer tensors")
    print("  4/4 transformer shards")
    print("  full T5/CLIP/VAE pipeline components present")


if __name__ == "__main__":
    main()