#!/usr/bin/env python3
"""Generate MRPGI sprites on the MODELBEAST farm.

    python3 tools/sprites.py cases/case-01-substantial-gift     # all missing
    python3 tools/sprites.py cases/... --only ego,johnny        # named
    python3 tools/sprites.py cases/... --force                  # redo existing

Chain per sprite: flux_local (text->image) -> bg_remove_local (RMBG-2.0 cutout)
-> Pillow nearest-neighbour downscale to the sprite's target height. MRPGI
quantizes to EGA itself on load (MANUAL SS14), so we only owe it a small
transparent PNG with the subject's feet on the bottom edge.

Sprite prompts live in sprites.json next to the game's game.json.
Token: MB_TOKEN env, else ~/Documents/backnforth/.env.
"""
import io
import json
import os
import sys
import time
import urllib.request
import uuid
from concurrent.futures import ThreadPoolExecutor

from PIL import Image

HOST = os.environ.get("MB_HOST", "http://100.89.131.57:8777")
# ponytail: guest token allows 4 active jobs; 3 leaves headroom for other users.
# Drop to MB_POOL=1 if the farm starts returning 401s (that's the cap, not auth).
POOL = int(os.environ.get("MB_POOL", "3"))

STYLE = (
    "1980s Sierra AGI adventure game pixel art sprite, EGA 16-colour palette, "
    "chunky low-resolution pixels, flat solid colour blocks, hard black outline, "
    "single subject centred and complete, full body head to feet, no cropping, "
    "plain flat white background, no text, no logo, no shadow, no border"
)


def token():
    t = os.environ.get("MB_TOKEN")
    if t:
        return t
    for line in open(os.path.expanduser("~/Documents/backnforth/.env")):
        if line.startswith("MB_TOKEN="):
            return line.split("=", 1)[1].strip()
    sys.exit("no MB_TOKEN")


def req(path, data=None, headers=None, raw=False):
    h = {"Authorization": f"Bearer {token()}"}
    h.update(headers or {})
    # The guest token caps concurrent jobs; over it the farm returns 429 (and
    # occasionally a transient 401/5xx). Back off rather than losing the sprite.
    for attempt in range(6):
        try:
            body = urllib.request.urlopen(
                urllib.request.Request(HOST + path, data=data, headers=h), timeout=300
            ).read()
            break
        except urllib.error.HTTPError as e:
            if e.code not in (401, 429, 500, 502, 503) or attempt == 5:
                raise
            time.sleep(5 * (attempt + 1))
    if raw:
        return body
    # Job logs carry raw control chars that break json.loads (known farm quirk).
    s = body.decode("utf-8", "replace")
    return json.loads("".join(c if c >= " " or c == "\t" else " " for c in s))


def upload(name, blob):
    b = uuid.uuid4().hex
    body = (
        f'--{b}\r\nContent-Disposition: form-data; name="file"; filename="{name}"\r\n'
        f"Content-Type: image/png\r\n\r\n".encode()
        + blob
        + f"\r\n--{b}--\r\n".encode()
    )
    a = req("/api/assets", data=body,
            headers={"Content-Type": f"multipart/form-data; boundary={b}"})
    return a.get("id") or (a.get("items") or [a])[0].get("id")


def run_job(payload):
    """Submit, poll to completion, return the output asset's bytes."""
    j = req("/api/jobs", data=json.dumps(payload).encode(),
            headers={"Content-Type": "application/json"})
    jid = j["id"]
    for _ in range(150):  # 150 * 4s = 10 min ceiling
        time.sleep(4)
        st = req(f"/api/jobs/{jid}").get("status")
        if st in ("done", "error", "cancelled"):
            break
    if st != "done":
        raise RuntimeError(f"job {jid} ended {st}")
    items = req("/api/assets?limit=100")
    items = items if isinstance(items, list) else items.get("items", [])
    # parent_job is the only reliable link back — newest-asset races other users.
    mine = [x for x in items if x.get("parent_job") == jid]
    if not mine:
        raise RuntimeError(f"job {jid} produced no linked asset")
    return req(f"/api/assets/{mine[0]['id']}/file", raw=True)


def trim_and_scale(blob, height):
    """Crop to the subject's alpha bounds, then nearest-neighbour to `height`.

    Trimming first is what puts the feet on the bottom edge, which is how MRPGI
    positions and depth-sorts a sprite (MANUAL SS14).
    """
    img = Image.open(io.BytesIO(blob)).convert("RGBA")
    box = img.getbbox()
    if box:
        img = img.crop(box)
    w = max(1, round(img.width * height / img.height))
    img = img.resize((w, height), Image.NEAREST)
    out = io.BytesIO()
    img.save(out, "PNG")
    return out.getvalue()


def make(name, spec, outdir):
    prompt = f"{spec['prompt']}. {STYLE}"
    png = run_job({
        "operator": "flux_local",
        "params": {"prompt": prompt, "model": "flux2-klein-4b",
                   "steps": 4, "width": 512, "height": 512,
                   # Stable seed per sprite name => reruns are reproducible.
                   "seed": abs(hash(name)) % 100000},
    })
    aid = upload(f"{name}_raw.png", png)
    cut = run_job({"operator": "bg_remove_local", "asset_id": aid,
                   "params": {"resolution": 1024, "background": "transparent"}})
    path = os.path.join(outdir, f"{name}.png")
    with open(path, "wb") as f:
        f.write(trim_and_scale(cut, spec.get("h", 22)))
    return path


def main():
    game = sys.argv[1].rstrip("/")
    only = None
    if "--only" in sys.argv:
        only = set(sys.argv[sys.argv.index("--only") + 1].split(","))
    force = "--force" in sys.argv

    specs = json.load(open(os.path.join(game, "sprites.json")))
    outdir = os.path.join(game, "sprites")
    os.makedirs(outdir, exist_ok=True)

    todo = {
        n: s for n, s in specs.items()
        if not n.startswith("_")  # "_comment" etc. are notes, not sprites
        and (only is None or n in only)
        and (force or not os.path.exists(os.path.join(outdir, f"{n}.png")))
    }
    if not todo:
        print("nothing to do")
        return
    print(f"[mb] {len(todo)} sprites -> {outdir}", flush=True)

    def one(item):
        name, spec = item
        try:
            make(name, spec, outdir)
            print(f"[mb] ok   {name}", flush=True)
        except Exception as e:  # keep the batch alive; report at the end
            print(f"[mb] FAIL {name}: {e}", flush=True)
            return name

    with ThreadPoolExecutor(POOL) as ex:
        failed = [f for f in ex.map(one, todo.items()) if f]
    print(f"[mb] done. failed: {failed or 'none'}")


if __name__ == "__main__":
    main()