#!/usr/bin/env python3
"""Generate a texture on the local MODELBEAST farm (operator: flux_local).

Usage: flux_local.py "<prompt>" <out.jpg> [size_px] [target_luma] [WxH]

size_px is the output's max dimension; WxH overrides the GENERATION aspect
(e.g. 1024x512 for a 2:1 sky panorama). Output keeps the generation aspect.

The fallback for tools/flux_texture.sh when Cloudflare Workers AI returns 429 --
its free tier is 10,000 neurons/day and a full texture-set regen burns through
it. Same output contract: a downscaled, luma-clamped jpg at <out>.

Queue + token per ~/.claude/skills/fleet/SKILL.md. The token is read from disk
and never printed. Job log JSON carries raw control chars, hence the scrub.
"""
import json
import os
import re
import sys
import time
import urllib.request

HOST = os.environ.get("MB_HOST", "http://100.89.131.57:8777")
ENVS = ["~/Documents/fluxgod-work/.env", "~/Documents/backnforth/.env"]


def token():
    for p in ENVS:
        p = os.path.expanduser(p)
        if not os.path.exists(p):
            continue
        for line in open(p):
            if line.startswith("MB_TOKEN"):
                return line.split("=", 1)[1].strip().strip("'\"")
    sys.exit("no MB_TOKEN found on disk")


def call(path, tok, data=None):
    req = urllib.request.Request(
        HOST + path,
        data=json.dumps(data).encode() if data is not None else None,
        headers={"Authorization": "Bearer " + tok, "Content-Type": "application/json"},
    )
    raw = urllib.request.urlopen(req, timeout=60).read().decode("utf-8", "replace")
    return json.loads(re.sub(r"[\x00-\x1f]", "", raw))


def main():
    prompt, out = sys.argv[1], sys.argv[2]
    size = int(sys.argv[3]) if len(sys.argv) > 3 else 512
    luma = float(sys.argv[4]) if len(sys.argv) > 4 else 0.0
    tok = token()

    if len(sys.argv) > 5 and "x" in sys.argv[5]:
        gw, gh = (int(v) for v in sys.argv[5].split("x"))
    else:
        gw = gh = max(size, 768)  # generate at >=768 then downscale, like the CF path
    job = call("/api/jobs", tok, {"operator": "flux_local", "params": {
        "prompt": prompt, "width": gw, "height": gh, "steps": 6}})
    jid = job["id"]
    for _ in range(120):
        time.sleep(3)
        st = call("/api/jobs/%s" % jid, tok)
        if st["status"] in ("done", "error"):
            break
    else:
        sys.exit("timed out waiting on job %s" % jid)
    if st["status"] != "done":
        sys.exit("job %s failed: %s" % (jid, (st.get("error") or "")[:200]))

    assets = call("/api/assets?limit=40", tok)
    rows = assets if isinstance(assets, list) else assets.get("assets", assets.get("items", []))
    mine = [a for a in rows if a.get("parent_job") == jid]
    if not mine:
        sys.exit("job %s produced no asset" % jid)

    req = urllib.request.Request("%s/api/assets/%s/file" % (HOST, mine[0]["id"]),
                                 headers={"Authorization": "Bearer " + tok})
    blob = urllib.request.urlopen(req, timeout=120).read()
    tmp = out + ".src"
    open(tmp, "wb").write(blob)

    from PIL import Image, ImageStat
    im = Image.open(tmp).convert("RGB")
    scale = size / max(im.size)
    im = im.resize((max(1, int(im.width * scale)), max(1, int(im.height * scale))), Image.LANCZOS)
    if luma:
        mean = ImageStat.Stat(im.convert("L")).mean[0]
        if mean > luma:
            im = im.point(lambda v, k=luma / mean: int(v * k))
    im.save(out, quality=88)
    os.remove(tmp)
    print("OK %s (local)" % out)


if __name__ == "__main__":
    main()