#!/usr/bin/env python3
"""run_batch.py — generate every sprite in assets.py, unattended.

    python3 tools/gen/run_batch.py                 # everything still missing
    python3 tools/gen/run_batch.py pokie jukebox   # just these kinds
    python3 tools/gen/run_batch.py --force         # redo even if the png exists
    python3 tools/gen/run_batch.py --flat-only     # cheap pass, no mesh jobs

Resumable by design: a kind whose `public/props/<kind>.png` already exists is
skipped, so a run that dies at asset 30 of 47 costs 17 assets to finish, not 47.
Every intermediate (product shot, cutout, GLB) is kept in `art_incoming/` for
the same reason — re-angling a prop should never mean re-meshing it, which is
the [GLB] route's whole advantage (LANEHANDOVER, FABLE-SOLO-25).

Image tier order is Cloudflare FIRST — 10,000 free neurons a day, ~4s an image —
then MODELBEAST `flux_local`, which is slower but unmetered. The fallback is
automatic and one-way per run: once Cloudflare reports the daily allocation
spent, we stop asking it and go to the farm for the rest of the batch.

Heartbeat lands in ~/.jobs/not-tonight-assets.json per the fleet jobs
convention, so "where are we at" is answerable without reading the log.
"""
from __future__ import annotations

import concurrent.futures as cf
import json
import subprocess
import sys
import threading
import time
import traceback
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
import assets as manifest  # noqa: E402
import cf_flux  # noqa: E402
import mb  # noqa: E402

REPO = Path(__file__).resolve().parent.parent.parent
INCOMING = REPO / "art_incoming"
PROPS = REPO / "public" / "props"
HEARTBEAT = Path.home() / ".jobs" / "not-tonight-assets.json"

# Flipped the first time Cloudflare says the day's neurons are gone. One-way:
# re-probing per asset would waste a call each time to learn the same thing.
_cf_spent = False


def beat(**fields: object) -> None:
    HEARTBEAT.parent.mkdir(parents=True, exist_ok=True)
    payload = {"job": "not-tonight-assets", "updated": time.strftime("%Y-%m-%d %H:%M:%S"), **fields}
    HEARTBEAT.write_text(json.dumps(payload, indent=2) + "\n")


def out_png(kind: str) -> Path:
    return PROPS / f"{kind.replace(':', '_')}.png"


def draw(a: manifest.Asset, dest: Path) -> str:
    """One image from the cheapest tier that still has budget. Returns the tier."""
    global _cf_spent
    w, h = a.gen
    if not _cf_spent:
        try:
            # Workers AI schnell has no width/height knob — it returns square.
            # Fine for props; the wide plates below need the farm regardless.
            if w == h:
                cf_flux.generate(dest, a.prompt, steps=6, seed=a.seed)
                return "cloudflare"
        except cf_flux.Exhausted:
            _cf_spent = True
            print("[gen] cloudflare free tier spent — rest of the batch on the farm", flush=True)
        except Exception as exc:  # network wobble: fall through to the farm
            print(f"[gen] cloudflare failed ({exc.__class__.__name__}), using farm", flush=True)
    mb.image(dest, a.prompt, seed=a.seed, steps=4, width=w, height=h, label=a.kind)
    return "farm"


def do_flat(a: manifest.Asset) -> None:
    raw = INCOMING / f"{a.kind.replace(':', '_')}_raw.png"
    if not raw.exists():
        draw(a, raw)
    cmd = [sys.executable, str(REPO / "tools" / "props_import.py"), str(raw), a.kind]
    if not a.quantise:
        cmd.append("--no-quantise")
    if a.keep_bg:
        cmd.append("--keep-bg")
    subprocess.run(cmd, check=True)


def do_mesh(a: manifest.Asset) -> None:
    stem = a.kind.replace(":", "_")
    shot = INCOMING / f"{stem}_shot.png"
    cut = INCOMING / f"{stem}_cut.png"
    glb = INCOMING / f"{stem}.glb"

    if not shot.exists():
        draw(a, shot)
    if not cut.exists():
        # TRELLIS reconstructs markedly better from a clean cutout than from an
        # object sitting on a grey sweep — the operator's own docs say so, and
        # it costs one cheap local job.
        try:
            asset_id = mb.upload(shot)
            job = mb.submit("bg_remove_local", {"resolution": 1024}, asset_id=asset_id)
            if mb.wait(job, label=f"{a.kind}:cut") == "done":
                got = mb.output_asset(job, ".png")
                if got:
                    mb.fetch(got, cut)
        except Exception as exc:
            print(f"[gen] {a.kind}: bg_remove failed ({exc}) — meshing the raw shot", flush=True)
    src = cut if cut.exists() else shot

    if not glb.exists():
        mb.mesh(src, glb, operator="trellis2_mlx", label=a.kind)

    # Blender ortho render at the house rake, then the normal sprite import.
    subprocess.run(
        [sys.executable, str(REPO / "tools" / "glb_to_sprite.py"), str(glb), a.kind,
         "--rake", str(a.rake), "--px", "512"],
        check=True,
    )


def main() -> None:
    raw_args = sys.argv[1:]
    force = "--force" in raw_args
    flat_only = "--flat-only" in raw_args
    # Drop flags AND the value that follows a value-taking flag — otherwise
    # `--workers 3` leaves a bare "3" in the positional list and the batch
    # filters down to the kind named "3", i.e. nothing.
    argv: list[str] = []
    skip_next = False
    for i, tok in enumerate(raw_args):
        if skip_next:
            skip_next = False
            continue
        if tok in ("--workers",):
            skip_next = True
            continue
        if tok.startswith("--"):
            continue
        argv.append(tok)

    todo = [a for a in manifest.ALL if not argv or a.kind in argv]
    if flat_only:
        todo = [a for a in todo if a.route == "flat"]
    if not force:
        todo = [a for a in todo if not out_png(a.kind).exists()]

    workers = int(sys.argv[sys.argv.index("--workers") + 1]) if "--workers" in sys.argv else 3
    INCOMING.mkdir(exist_ok=True)
    done: list[str] = []
    failed: list[str] = []
    print(f"[gen] {len(todo)} assets to make, {workers} at a time", flush=True)
    beat(status="running", total=len(todo), done=0, failed=0)

    # The farm load-balances gpu jobs across nodes, so a strictly sequential
    # batch leaves most of the cluster idle — 37 meshes at ~7 min each is over
    # four hours in a queue of one. Three in flight stays under the guest
    # token's 4-active ceiling and keeps every node fed. Safe to overlap
    # because output retrieval matches on `parent_job`, never newest-of-type
    # (the 2026-07-16 race that handed one project's mesh to another).
    lock = threading.Lock()

    def run_one(a: manifest.Asset) -> None:
        t0 = time.time()
        try:
            (do_flat if a.route == "flat" else do_mesh)(a)
            with lock:
                done.append(a.kind)
            print(f"[gen] {a.kind} OK in {int(time.time() - t0)}s", flush=True)
        except Exception as exc:
            with lock:
                failed.append(a.kind)
            print(f"[gen] {a.kind} FAILED: {exc}", flush=True)
            traceback.print_exc()
        with lock:
            beat(status="running", total=len(todo), done=len(done), failed=len(failed),
                 last=a.kind, still_to_go=len(todo) - len(done) - len(failed))

    with cf.ThreadPoolExecutor(max_workers=workers) as pool:
        list(pool.map(run_one, todo))

    beat(status="finished", total=len(todo), done=len(done), failed=len(failed),
         failed_kinds=failed)
    print(f"\n[gen] DONE {len(done)}/{len(todo)}"
          + (f" — FAILED: {', '.join(failed)}" if failed else ""))
    # A partial batch is a normal outcome; the venue layouts fall back to
    # placeholders for anything missing, so this must not read as total failure.
    sys.exit(0 if not failed else 2)


if __name__ == "__main__":
    main()