#!/usr/bin/env python3
"""Generate a room background on the MODELBEAST farm (flux_local).

Reads MB_TOKEN from ~/Documents/backnforth/.env (never printed). The engine
quantizes whatever lands in pics/ to 256 colors at load, so the prompt aims
for painterly SCI1 material, not pixel-perfection.

  python3 tools/mb_background.py roskilde "prompt text..."
"""

import json
import os
import re
import sys
import time
import urllib.request

HOST = "http://100.89.131.57:8777"


def token():
    # MB_TOKEN env var, else the usual .env spots (ultra keeps it in
    # backnforth, the m3ultra primary in MODELBEAST/.env as MB_LOCAL_TOKEN).
    if os.environ.get("MB_TOKEN"):
        return os.environ["MB_TOKEN"]
    for env, keys in [
        (os.path.expanduser("~/Documents/backnforth/.env"), ("MB_TOKEN",)),
        (os.path.expanduser("~/Documents/MODELBEAST/.env"), ("MB_TOKEN", "MB_LOCAL_TOKEN")),
    ]:
        if not os.path.exists(env):
            continue
        with open(env) as f:
            for line in f:
                for k in keys:
                    m = re.match(rf"\s*{k}\s*=\s*(\S+)", line)
                    if m:
                        return m.group(1).strip().strip('"')
    raise SystemExit("no MODELBEAST token found in known .env locations")


def api(path, data=None, raw=False):
    req = urllib.request.Request(
        HOST + path,
        data=json.dumps(data).encode() if data is not None else None,
        headers={"Authorization": f"Bearer {token()}", "Content-Type": "application/json"},
        method="POST" if data is not None else "GET",
    )
    with urllib.request.urlopen(req, timeout=60) as r:
        body = r.read()
    if raw:
        return body
    # job logs can carry raw control chars — strip before parsing
    text = re.sub(rb"[\x00-\x08\x0b-\x1f]", b"", body).decode("utf-8", "replace")
    return json.loads(text)


def main():
    name, prompt = sys.argv[1], sys.argv[2]
    job = api("/api/jobs", {"operator": "flux_local", "params": {"prompt": prompt, "width": 1024, "height": 608}})
    jid = job.get("id") or job.get("job_id")
    print(f"submitted flux_local job {jid}")
    for _ in range(120):
        time.sleep(3)
        j = api(f"/api/jobs/{jid}")
        status = j.get("status", "?")
        if status in ("done", "error"):
            print(f"job {jid}: {status}")
            if status == "error":
                raise SystemExit(1)
            break
    else:
        raise SystemExit("timed out")
    # newest own asset; jobs don't back-link, so match parent_job
    assets = api("/api/assets")
    items = assets if isinstance(assets, list) else assets.get("assets", assets.get("items", []))
    mine = [a for a in items if str(a.get("parent_job", "")) == str(jid)]
    pick = (mine or items)[0]
    data = api(f"/api/assets/{pick['id']}/file", raw=True)
    out = os.path.join(os.path.dirname(__file__), "..", "pics", f"{name}.png")
    os.makedirs(os.path.dirname(out), exist_ok=True)
    with open(out, "wb") as f:
        f.write(data)
    print(f"wrote pics/{name}.png ({len(data)} bytes, asset {pick.get('name', pick['id'])})")


if __name__ == "__main__":
    main()