#!/usr/bin/env python3
"""Download the whole Gaussian Splat Objects Dataset.

Every object, every level of detail, no exceptions. About 2.7 GB in all.
Standard library only, so there is nothing to install first.

    python download_all_splats.py                 # into ./splats
    python download_all_splats.py --out /somewhere
    python download_all_splats.py --lod 10k       # only the small tier, ~600 MB

Files are written straight to disk as they arrive, one at a time, so nothing has
to fit in memory and there is no zip to unpack afterwards. Every file is checked
against the SHA-256 recorded in the dataset inventory. Anything already on disk
and correct is skipped, so if the run is interrupted you can simply run it again
and it carries on where it stopped.

Data and licenses: https://github.com/marcelpadilla/splats
"""

import argparse
import hashlib
import json
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path

INVENTORY_URL = "https://raw.githubusercontent.com/marcelpadilla/splats/main/data/inventory.json"
AGENT = {"User-Agent": "splats-download-all/1"}
CHUNK = 1 << 20
RETRIES = 4
TIMEOUT = 60


def human(n):
    for unit, step in (("GB", 1e9), ("MB", 1e6), ("kB", 1e3)):
        if n >= step:
            return f"{n / step:.1f} {unit}"
    return f"{n:.0f} B"


def fetch(url):
    req = urllib.request.Request(url, headers=AGENT)
    with urllib.request.urlopen(req, timeout=TIMEOUT) as r:
        return r.read()


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


def download(url, dest, sha, size):
    """Stream one file to disk, verify it, then move it into place.

    It lands on a .part file first, so an interrupted run never leaves a
    truncated splat that looks finished on the next pass.
    """
    part = dest.with_name(dest.name + ".part")
    h = hashlib.sha256()
    got = 0
    req = urllib.request.Request(url, headers=AGENT)
    with urllib.request.urlopen(req, timeout=TIMEOUT) as r, open(part, "wb") as f:
        while True:
            chunk = r.read(CHUNK)
            if not chunk:
                break
            f.write(chunk)
            h.update(chunk)
            got += len(chunk)
            if size:
                print(f"\r    {100 * got / size:5.1f}%  {human(got)}", end="", flush=True)
    if size:
        print("\r" + " " * 32 + "\r", end="")
    if sha and h.hexdigest() != sha:
        part.unlink(missing_ok=True)
        raise ValueError("checksum did not match, the file was not kept")
    part.replace(dest)
    return got


def wanted_files(objects, lod):
    """Every file to fetch, as (relative path, sha256 or None, bytes)."""
    out = []
    for o in objects:
        tiers = o.get("lods") or []
        if not tiers:
            # Single-file objects have no tiers to choose between, so they come
            # along whatever --lod says. Leaving them out would quietly drop the
            # generated objects, which are most of the collection.
            out.append((o["file"], o.get("sha256"), o.get("bytes", 0)))
        else:
            for t in tiers:
                if lod in (None, "all") or t["lod"] == lod:
                    out.append((t["file"], t.get("sha256"), t.get("bytes", 0)))
        for extra in (o.get("meta"), o.get("thumbnail")):
            if extra:
                out.append((extra, None, 0))
    return out


def main():
    ap = argparse.ArgumentParser(
        description="Download every object in the Gaussian Splat Objects Dataset.")
    ap.add_argument("--out", default="splats",
                    help="where to put them (default ./splats)")
    ap.add_argument("--lod", default="all",
                    help="only this tier for the objects that have tiers "
                         "(10k, 100k, 500k, 1m). Default is all of them.")
    ap.add_argument("--inventory", default=INVENTORY_URL,
                    help="the inventory to read (default is the live one)")
    args = ap.parse_args()

    out = Path(args.out).expanduser().resolve()
    out.mkdir(parents=True, exist_ok=True)

    print(f"Reading the inventory from {args.inventory}")
    try:
        raw = fetch(args.inventory)
    except (urllib.error.URLError, OSError) as e:
        sys.exit(f"Could not read the inventory: {e}")
    doc = json.loads(raw)
    base = doc["base_url"]
    objects = doc["objects"]

    # A mistyped tier must not pass quietly. Single-file objects come along
    # whatever --lod says, so the run would look like it worked and hand back a
    # fraction of the collection.
    tiers = {}
    for o in objects:
        for t in o.get("lods") or []:
            tiers.setdefault(t["lod"], t.get("splats", 0))
    known = sorted(tiers, key=tiers.get)
    if args.lod not in ("all", *known):
        sys.exit(f"There is no {args.lod} tier. Try: {', '.join(known)} or all.")

    files = wanted_files(objects, args.lod)
    total = sum(b for _, _, b in files)

    print(f"{len(objects)} objects, {len(files)} files, about {human(total)}")
    print(f"Into {out}")
    print()

    (out / "inventory.json").write_bytes(raw)

    fetched = skipped = failed = 0
    for i, (rel, sha, size) in enumerate(files, 1):
        dest = out / rel
        dest.parent.mkdir(parents=True, exist_ok=True)
        head = f"[{i}/{len(files)}] {rel}"

        if dest.exists() and (sha256_of(dest) == sha if sha else dest.stat().st_size > 0):
            skipped += 1
            print(f"{head}  already here")
            continue

        print(f"{head}{'  ' + human(size) if size else ''}")
        for attempt in range(1, RETRIES + 1):
            try:
                download(base + rel, dest, sha, size)
                fetched += 1
                break
            except KeyboardInterrupt:
                raise
            except Exception as e:
                if attempt == RETRIES:
                    print(f"    gave up: {e}")
                    failed += 1
                else:
                    wait = 2 ** attempt
                    print(f"    {e}. Trying again in {wait}s")
                    time.sleep(wait)

    print()
    print(f"Done. {fetched} downloaded, {skipped} already here, {failed} failed.")
    print(f"They are in {out}")
    print("Each object's license is in its meta.json, and all of them are in "
          "inventory.json.")
    if failed:
        print("Run this again to retry the ones that failed.")
        return 1
    return 0


if __name__ == "__main__":
    try:
        sys.exit(main())
    except KeyboardInterrupt:
        print("\nStopped. Run this again to carry on where it left off.")
        sys.exit(130)