#!/usr/bin/env python3
"""Batch-generate GigaSlop art on MODELBEAST (flux_local + bg_remove_local).

Contract per ~/Documents/MESHGOD/scripts/mb_recon.py and thriftgod/gen_assets.py:
multipart upload to /api/assets (field "file"); jobs take {operator, asset_id, params};
outputs found by scanning /api/assets for parent_job (jobs don't back-link).

Usage: python3 tools/mb_gen.py [--only room|objects|slop]
Writes PNGs to public/assets/{gen,slop}/, manifest to tools/mb_manifest.json.
Heartbeats to ~/.jobs/gigaslop-artgen.status.
"""
import hashlib, json, mimetypes, os, pathlib, sys, time, urllib.request, uuid

HOST = "http://100.89.131.57:8777"
ROOT = pathlib.Path(__file__).resolve().parent.parent
HEART = pathlib.Path.home() / ".jobs" / "gigaslop-artgen.status"
MAX_ACTIVE = 3  # guest cap is 4; leave headroom


def token():
    for line in (pathlib.Path.home() / "Documents/backnforth/.env").read_text().splitlines():
        if line.startswith("MB_TOKEN="):
            return line.split("=", 1)[1].strip()
    raise SystemExit("MB_TOKEN not found")


TOKEN = token()


def api(path, payload=None, data=None, headers=None, raw=False):
    h = {"Authorization": "Bearer " + TOKEN}
    if payload is not None:
        data = json.dumps(payload).encode()
        h["Content-Type"] = "application/json"
    h.update(headers or {})
    r = urllib.request.Request(HOST + path, data=data, headers=h)
    body = urllib.request.urlopen(r, timeout=180).read()
    if raw:
        return body
    s = body.decode("utf-8", "replace")
    return json.loads("".join(c if c >= " " or c in "\t" else " " for c in s))


def upload(path):
    boundary = uuid.uuid4().hex
    ctype = mimetypes.guess_type(str(path))[0] or "application/octet-stream"
    name = os.path.basename(path)
    body = (
        f'--{boundary}\r\nContent-Disposition: form-data; name="file"; filename="{name}"\r\n'
        f"Content-Type: {ctype}\r\n\r\n"
    ).encode() + pathlib.Path(path).read_bytes() + f"\r\n--{boundary}--\r\n".encode()
    a = api("/api/assets", data=body,
            headers={"Content-Type": f"multipart/form-data; boundary={boundary}"})
    return a.get("id") or (a.get("items") or [a])[0].get("id")


def seed_for(slug):
    return int(hashlib.sha1(slug.encode()).hexdigest()[:7], 16)


def beat(msg):
    HEART.parent.mkdir(exist_ok=True)
    HEART.write_text(json.dumps({"ts": time.strftime("%F %T"), "msg": msg}) + "\n")
    print(msg, flush=True)


def asset_for_job(jid):
    items = api("/api/assets")
    if isinstance(items, dict):
        items = items.get("items", [])
    mine = [a for a in items if a.get("parent_job") == jid]
    return mine[0] if mine else None


def run_batch(specs, outdir):
    """specs: list of (label, operator, asset_id, params). Throttled, unordered completion."""
    outdir.mkdir(parents=True, exist_ok=True)
    results, queue, active = {}, list(specs), {}
    while queue or active:
        while queue and len(active) < MAX_ACTIVE:
            label, op, aid, params = queue.pop(0)
            jid = api("/api/jobs", {"operator": op, "asset_id": aid, "params": params})["id"]
            active[jid] = label
            beat(f"submitted {label} -> job {jid} ({len(queue)} queued)")
        time.sleep(3)
        for jid in list(active):
            st = api(f"/api/jobs/{jid}").get("status")
            if st in ("done", "error", "cancelled"):
                label = active.pop(jid)
                if st != "done":
                    beat(f"FAILED {label}: {st}")
                    continue
                a = asset_for_job(jid)
                if not a:
                    beat(f"done {label} but no asset for job {jid}")
                    continue
                data = api(f"/api/assets/{a['id']}/file", raw=True)
                path = outdir / f"{label}.png"
                path.write_bytes(data)
                results[label] = {"job": jid, "asset": a["id"], "path": str(path)}
                beat(f"done {label} ({len(data)//1024} KB)")
    return results


STYLE = (
    "isometric video game asset, 2:1 isometric angle viewed from above at 45 degrees, "
    "dark moody lighting with neon accent glow, detailed painterly style, crisp edges, "
    "centered single object on plain solid dark grey background, no text, "
)

OBJECTS = [
    ("desk_pc", "a cheap beige-and-black budget gaming PC tower with one small GPU, side panel off, dusty, one RGB fan"),
    ("desk", "a cluttered small wooden computer desk with monitor, keyboard, energy drink cans, tangled cables"),
    ("router", "a consumer wifi router with antennas, blinking green LEDs"),
    ("desk_fan", "a cheap plastic oscillating desk fan"),
    ("gpu_rig", "a DIY open-air mining rig frame made of milk crates holding four mismatched GPUs, tangled riser cables"),
    ("window_ac", "a battered window air conditioning unit dripping slightly"),
    ("power_strip", "an overloaded power strip surge protector with too many plugs, slightly scorched"),
    ("server_rack_mini", "a small half-height 19 inch server rack with a few rack servers, status LEDs"),
]

SLOP = [
    "hyperrealistic cat wearing a business suit crying in the rain, shocked face, red arrow, oversaturated",
    "muscular baby bodybuilder lifting a car, impossible anatomy, oversaturated colors",
    "two podcast hosts yelling at each other across a table, exaggerated shocked expressions, red arrows",
    "ancient roman emperor playing a video game on a glowing gaming PC, dramatic lighting",
    "infinite spiral of golden retrievers wearing sunglasses on a beach, uncanny, oversaturated",
    "a submarine made of watermelons in the ocean, shocked scuba diver pointing, oversaturated",
    "grandma arm wrestling a robot in a kitchen, sparks flying, exaggerated expressions",
    "hyperrealistic shark bursting out of a swimming pool at a birthday party, red circle",
    "medieval knight reviewing fast food fries, uncanny smile, bright arrows",
    "city skyline made entirely of pasta at sunset, tiny cars, oversaturated dreamlike",
]
SLOP_STYLE = "AI generated clickbait youtube thumbnail, "

ROOM = (
    "empty isometric bedroom interior for a video game, 2:1 isometric projection, two visible walls "
    "meeting at back corner, dark wooden floor with subtle grid tiles, night time, window with city "
    "lights, moody blue-purple lighting, one neon strip, no furniture, no people, no text, "
    "clean detailed painterly game art"
)


def flux_spec(label, prompt, w, h):
    return (label, "flux_local", None,
            {"prompt": prompt, "model": "flux2-klein-4b", "steps": 4,
             "width": w, "height": h, "seed": seed_for(label)})


def main():
    only = sys.argv[sys.argv.index("--only") + 1] if "--only" in sys.argv else None
    manifest = {}
    gen, slop = ROOT / "public/assets/gen", ROOT / "public/assets/slop"

    if only in (None, "room"):
        beat("phase: room shell")
        manifest.update(run_batch([flux_spec("room", ROOM, 1024, 768)], gen))

    if only in (None, "objects"):
        beat("phase: object sprites")
        raw = run_batch([flux_spec(n, STYLE + d, 768, 768) for n, d in OBJECTS], gen)
        manifest.update(raw)
        beat("phase: bg removal")
        cut_specs = []
        for name, info in raw.items():
            aid = upload(info["path"])
            cut_specs.append((f"{name}_cut", "bg_remove_local", aid,
                              {"resolution": 1024, "background": "transparent"}))
        manifest.update(run_batch(cut_specs, gen))

    if only in (None, "slop"):
        beat("phase: slop thumbnails")
        manifest.update(run_batch(
            [flux_spec(f"slop_{i:02d}", SLOP_STYLE + p, 640, 384) for i, p in enumerate(SLOP)],
            slop))

    mpath = ROOT / "tools/mb_manifest.json"
    old = json.loads(mpath.read_text()) if mpath.exists() else {}
    old.update(manifest)
    mpath.write_text(json.dumps(old, indent=2))
    beat(f"ALL DONE — {len(manifest)} assets this run")


if __name__ == "__main__":
    main()