#!/usr/bin/env python3
"""wardrobegod batch factory — build-first catalogue items, serially, via the
live wardrobegod API. Heartbeat at ~/.jobs/wardrobe-batch1.json.

Per item: /api/gen (klein on the farm) -> /api/rmbg (RMBG-2.0 cutout) ->
[3D items] /api/to3d (TRELLIS, ~5 min serial GPU lane).
Cutout PNGs and GLBs land in library/ automatically — that is the point of
driving the product API instead of the farm directly.
"""
import json, os, time, urllib.request

API = "http://100.91.239.7:8150"
HB = os.path.expanduser("~/.jobs/wardrobe-batch1.json")
os.makedirs(os.path.dirname(HB), exist_ok=True)

# id, phrase, template, view, to3d?
ITEMS = [
    # --- R2 rigid props/hats -> full 3D chain ---
    ("cap-baseball-fwd",     "faded navy blue baseball cap", "object3d", "hat", True),
    ("hardhat-construction", "yellow construction hard hat", "object3d", "hat", True),
    ("headphones-dj-overear","black over-ear DJ headphones with a coiled cable", "object3d", "hat", True),
    ("record-vinyl-12in",    "12 inch black vinyl record with a plain orange centre label", "object3d", "hat", True),
    ("sunglasses-wayfarer-generic", "black chunky rectangular-frame sunglasses", "object3d", "hat", True),
    ("bag-record-crossbody", "black canvas record bag with a wide shoulder strap", "object3d", "bag", True),
    ("shopping-basket",      "red plastic shopping basket with two metal handles", "object3d", "bag", True),
    ("sneaker-low-canvas",   "pair of plain white low-top canvas sneakers", "object3d", "shoes", True),
    ("boot-work-steelcap",   "pair of tan leather steel-cap work boots", "object3d", "shoes", True),
    ("chain-necklace-curb",  "chunky silver curb chain necklace laid in a circle", "object3d", "hat", True),
    # --- R1 garment cutouts (gen + rmbg only) ---
    ("tee-plain-white",      "plain white crew-neck t-shirt", "flat3d", "top", False),
    ("tee-print-smiley",     "black t-shirt with a large orange smiley face print", "flat3d", "top", False),
    ("tee-print-sunset",     "white t-shirt with a retro sunset stripes print", "flat3d", "top", False),
    ("hivis-vest-yellow",    "yellow hi-vis safety vest with silver reflective stripes", "flat3d", "top", False),
    ("tights-opaque-black",  "pair of black opaque tights", "flat3d", "bottom", False),
    ("shorts-denim-cutoff",  "light blue denim cutoff shorts", "flat3d", "bottom", False),
]


def api(path, payload=None, timeout=30):
    if payload is None:
        req = urllib.request.Request(API + path)
    else:
        req = urllib.request.Request(API + path, data=json.dumps(payload).encode(),
                                     headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.load(r)


def wait(jid, timeout=1500):
    t0 = time.time()
    while time.time() - t0 < timeout:
        j = api(f"/api/job/{jid}")
        if j.get("status") in ("done", "error"):
            return j
        time.sleep(6)
    return {"status": "error", "log": "poll timeout"}


def hb(state):
    state["updated"] = time.strftime("%Y-%m-%d %H:%M:%S")
    json.dump(state, open(HB, "w"), indent=1)


results = []
state = {"job": "wardrobe-batch1", "status": "running", "total": len(ITEMS),
         "done": 0, "failed": 0, "current": "", "results": results}
hb(state)

for i, (iid, phrase, tpl, view, want3d) in enumerate(ITEMS):
    row = {"id": iid, "steps": {}}
    results.append(row)
    state["current"] = iid
    hb(state)
    try:
        g = api("/api/gen", {"phrase": phrase, "template": tpl, "view": view,
                             "backend": "klein"})
        j = wait(g["job"], 300)
        row["steps"]["gen"] = j.get("status")
        if j.get("status") != "done":
            raise RuntimeError("gen: " + str(j.get("log", ""))[:200])
        png = j["out"]

        c = api("/api/rmbg", {"path": png})
        j = wait(c["job"], 300)
        row["steps"]["rmbg"] = j.get("status")
        if j.get("status") != "done":
            raise RuntimeError("rmbg: " + str(j.get("log", ""))[:200])
        cut = j["out"]
        row["cut"] = cut

        if want3d:
            t = api("/api/to3d", {"path": cut})
            j = wait(t["job"], 1500)
            row["steps"]["to3d"] = j.get("status")
            if j.get("status") != "done":
                raise RuntimeError("to3d: " + str(j.get("log", ""))[:200])
            row["glb"] = j.get("out")
        state["done"] += 1
    except Exception as e:
        row["error"] = str(e)[:300]
        state["failed"] += 1
    hb(state)

state["status"] = "done"
state["current"] = ""
hb(state)
print(json.dumps(state, indent=1))