#!/usr/bin/env python3
"""Mirror the ACE-Step production model packages from the upstream R2 bucket.

Downloads the three content-addressed packages (direct reference, DiT rev7,
VAE rev7), verifies every payload's SHA-256 and byte length against its
manifest, and checks the aggregate inventory against the counts documented in
the upstream demo README (113 payloads + 3 manifests, 5,749,459,255 bytes).

Files the reference manifest lists but the direct R2 deployment deliberately
excludes (planner/semantic material, VAE-only shards, superseded DiT layers)
404 and are recorded as `excluded` — the success criterion is byte-exact
match of what IS served, not of the full manifest.
"""
import concurrent.futures as cf
import hashlib
import json
import os
import sys
import urllib.request

ORIGIN = "https://ace-step-wgsl-models.narcotic.sh"
DEST = os.path.dirname(os.path.abspath(__file__))
PACKAGES = [
    {
        "id": "reference-direct",
        "prefix": "v1/reference/18f36c6420976475af65ecd833ca56c6119706322ce54120389d4915d8e80db6",
        "manifest": "direct-manifest-b44a3d157009d035a8f20aa752db4ceef2fac5bd140eff13be8f7488bc978089.json",
        "manifest_sha256": "b44a3d157009d035a8f20aa752db4ceef2fac5bd140eff13be8f7488bc978089",
        "expect_files": 58,
        "expect_bytes": 2_558_130_431,
    },
    {
        "id": "dit-revision7",
        "prefix": "v1/dit-revision7/d3fc0020efcf60702db411da2fd4b93e9bb84f1437ed310aef01c892727e452f",
        "manifest": "manifest.json",
        "manifest_sha256": "d3fc0020efcf60702db411da2fd4b93e9bb84f1437ed310aef01c892727e452f",
        "expect_files": 48,
        "expect_bytes": 3_020_808_192,
    },
    {
        "id": "vae-revision7",
        "prefix": "v1/vae-revision7/36a54d79777d6826088095ba6ebc028fb4bea546368c0f0a29cd0eee8d656da7",
        "manifest": "manifest.json",
        "manifest_sha256": "36a54d79777d6826088095ba6ebc028fb4bea546368c0f0a29cd0eee8d656da7",
        "expect_files": 7,
        "expect_bytes": 168_791_552,
    },
]


def fetch(url: str) -> bytes:
    req = urllib.request.Request(url, headers={"User-Agent": "fluidaudio-mirror/1"})
    with urllib.request.urlopen(req, timeout=120) as r:
        return r.read()


def mirror_file(prefix: str, name: str, sha256: str, byte_length: int):
    """Download one payload; returns (name, status, bytes)."""
    out = os.path.join(DEST, prefix, name)
    if os.path.exists(out):
        data = open(out, "rb").read()
        if len(data) == byte_length and hashlib.sha256(data).hexdigest() == sha256:
            return name, "cached", byte_length
    try:
        data = fetch(f"{ORIGIN}/{prefix}/{name}")
    except urllib.error.HTTPError as e:
        if e.code == 404:
            return name, "excluded", 0
        raise
    digest = hashlib.sha256(data).hexdigest()
    if digest != sha256:
        return name, f"HASH MISMATCH {digest}", len(data)
    if len(data) != byte_length:
        return name, f"SIZE MISMATCH {len(data)}", len(data)
    os.makedirs(os.path.dirname(out), exist_ok=True)
    with open(out + ".tmp", "wb") as f:
        f.write(data)
    os.replace(out + ".tmp", out)
    return name, "ok", byte_length


def main():
    grand_files = 0
    grand_bytes = 0
    failures = []
    for pkg in PACKAGES:
        murl = f"{ORIGIN}/{pkg['prefix']}/{pkg['manifest']}"
        mbytes = fetch(murl)
        mdigest = hashlib.sha256(mbytes).hexdigest()
        if mdigest != pkg["manifest_sha256"]:
            print(f"FATAL: manifest hash mismatch for {pkg['id']}: {mdigest}")
            sys.exit(1)
        mpath = os.path.join(DEST, pkg["prefix"], pkg["manifest"])
        os.makedirs(os.path.dirname(mpath), exist_ok=True)
        open(mpath, "wb").write(mbytes)
        manifest = json.loads(mbytes)
        files = manifest["files"]
        print(f"[{pkg['id']}] manifest verified, {len(files)} file records", flush=True)

        got_files = 0
        got_bytes = 0
        excluded = 0
        with cf.ThreadPoolExecutor(max_workers=6) as ex:
            futures = [
                ex.submit(mirror_file, pkg["prefix"], f["name"], f["sha256"], f["byteLength"])
                for f in files
            ]
            for i, fut in enumerate(cf.as_completed(futures)):
                name, status, nbytes = fut.result()
                if status in ("ok", "cached"):
                    got_files += 1
                    got_bytes += nbytes
                elif status == "excluded":
                    excluded += 1
                else:
                    failures.append((pkg["id"], name, status))
                if (i + 1) % 20 == 0:
                    print(f"[{pkg['id']}] {i + 1}/{len(files)} processed, {got_bytes/1e9:.2f} GB", flush=True)

        ok_count = got_files == pkg["expect_files"] and got_bytes == pkg["expect_bytes"]
        print(
            f"[{pkg['id']}] {'PASS' if ok_count else 'MISMATCH'}: "
            f"{got_files} files / {got_bytes} bytes "
            f"(expected {pkg['expect_files']} / {pkg['expect_bytes']}), {excluded} excluded-by-design",
            flush=True,
        )
        if not ok_count:
            failures.append((pkg["id"], "<inventory>", f"{got_files} files / {got_bytes} bytes"))
        grand_files += got_files
        grand_bytes += got_bytes

    print(f"TOTAL: {grand_files} payloads, {grand_bytes} bytes (expected 113 / 5747730175)")
    if failures:
        print("FAILURES:")
        for pkg_id, name, status in failures:
            print(f"  {pkg_id}: {name}: {status}")
        sys.exit(1)
    print("MIRROR COMPLETE AND VERIFIED")


if __name__ == "__main__":
    main()