#!/usr/bin/env python3
"""Fetch the conveyor art that scene/sorter.usd references.

    python scripts/fetch_assets.py

~270 MB, too big for git. Downloads in resumable Range chunks: the Omniverse bucket is
slow enough that a plain GET truncates, and a half-written USD fails to compose with a
confusing "could not open asset" rather than an obvious size error.
"""
from __future__ import annotations
import os, sys, time, urllib.request
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
DEST = ROOT / "assets" / "conveyors"
BASE = ("https://omniverse-content-production.s3-us-west-2.amazonaws.com/"
        "Assets/Isaac/6.0/Isaac/Props/Conveyors/")
CHUNK, TRIES = 1 << 21, 12

# the scene composes these; the Textures/ and Material Library/ folders come with them
FILES = ["ConveyorBelt_A06.usd", "ConveyorBelt_A24.usd"]
TEXTURES = [
    "Textures/M_ConveyorBelt_A01_Belt.usd", "Textures/M_ConveyorBelt_A01_Decal.usd",
    "Textures/MetalPainted_Blue_Glossy_A.usd", "Textures/Plastic_Orange_A.usd",
    "Textures/Plastic_Rough_Black_A.usd", "Textures/Steel_A.usd",
    "Textures/T_ConveyorBelt_A01_Belt_Albedo.png", "Textures/T_ConveyorBelt_A01_Belt_Normal.png",
    "Textures/T_ConveyorBelt_A01_Belt_ORM.png",
    "Textures/T_ConveyorsBelt_A01_Decal_Albedo.png", "Textures/T_ConveyorsBelt_A01_Decal_Alpha.png",
    "Textures/T_ConveyorsBelt_A01_Decal_ORM.png",
    "Material%20Library/physics_material.usd",
]

def size_of(url):
    try:
        req = urllib.request.Request(url, method="HEAD")
        with urllib.request.urlopen(req, timeout=60) as r:
            return int(r.headers.get("Content-Length", 0))
    except Exception:
        return 0

def grab(url, out: Path):
    total = size_of(url)
    have = out.stat().st_size if out.exists() else 0
    if total and have == total:
        print(f"  {out.name}: present"); return True
    out.parent.mkdir(parents=True, exist_ok=True)
    if not total:
        urllib.request.urlretrieve(url, out); return out.exists()
    pos = have if have < total else 0
    with open(out, "r+b" if pos else "wb") as fh:
        fh.seek(pos); t0 = time.time()
        while pos < total:
            end = min(pos + CHUNK - 1, total - 1)
            for a in range(TRIES):
                try:
                    req = urllib.request.Request(url)
                    req.add_header("Range", f"bytes={pos}-{end}")
                    with urllib.request.urlopen(req, timeout=120) as r:
                        data = r.read()
                    if len(data) != end - pos + 1:
                        raise IOError("short chunk")
                    fh.write(data); fh.flush(); pos += len(data)
                    print(f"\r  {out.name}: {100.0*pos/total:5.1f}%"
                          f"  {(pos-have)/max(time.time()-t0,1e-6)/1024:6.0f} KB/s", end="", flush=True)
                    break
                except Exception as exc:
                    if a == TRIES - 1:
                        print(f"\n  {out.name}: FAILED at {pos}: {exc}"); return False
                    time.sleep(min(2 ** a, 20))
    print()
    return out.stat().st_size == total

def main():
    print(f"fetching conveyor art into {DEST}")
    ok = True
    for rel in FILES + TEXTURES:
        local = DEST / rel.replace("%20", " ")
        ok &= grab(BASE + rel, local)
    print("\nall present" if ok else "\nincomplete - re-run, downloads resume")
    return 0 if ok else 1

if __name__ == "__main__":
    sys.exit(main())