#!/usr/bin/env python3
"""flux.mjs drop-in for the M3 Ultra MODELBEAST box — same TSV manifests, $0, no daily cap.
usage: python3 tools/flux_local.py manifest.tsv outdir [variants=1]
env:   FLUX_W/FLUX_H (default 1024x1024), MB_HOST (default http://100.89.131.57:8777),
       FLUX_MODEL (default flux2-klein-4b), FLUX_STEPS (default 4)
Writes outdir/<name>_<v>.jpg, skips existing — identical contract to flux.mjs, so every
existing art_*.tsv runs unchanged. Seed = stable hash of name+variant → reproducible.
"""
import json, os, sys, time, urllib.request, hashlib

mani, outdir = sys.argv[1], sys.argv[2]
variants = int(sys.argv[3]) if len(sys.argv) > 3 else 1
HOST = os.environ.get('MB_HOST', 'http://100.89.131.57:8777')
W, H = int(os.environ.get('FLUX_W', 1024)), int(os.environ.get('FLUX_H', 1024))
MODEL = os.environ.get('FLUX_MODEL', 'flux2-klein-4b')
STEPS = int(os.environ.get('FLUX_STEPS', 4))
os.makedirs(outdir, exist_ok=True)

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={'Content-Type': 'application/json'} if data is not None else {})
    with urllib.request.urlopen(req, timeout=120) as r:
        return r.read() if raw else json.loads(r.read())

rows = [l.rstrip('\n').split('\t', 1) for l in open(mani)
        if l.strip() and not l.startswith('#')]
made = failed = 0
for name, prompt in rows:
    for v in range(1, variants + 1):
        out = os.path.join(outdir, f'{name}_{v}.jpg')
        if os.path.exists(out): print('skip', out); continue
        seed = int(hashlib.sha1(f'{name}_{v}'.encode()).hexdigest()[:7], 16)
        try:
            job = api('/api/jobs', {'operator': 'flux_local', 'asset_id': None,
                'params': {'prompt': prompt, 'model': MODEL, 'steps': STEPS,
                           'width': W, 'height': H, 'seed': seed}})['id']
            while True:
                j = api(f'/api/jobs/{job}')
                if j['status'] in ('done', 'error', 'cancelled'): break
                time.sleep(2)
            if j['status'] != 'done':
                raise RuntimeError(j.get('log', '')[-200:])
            outs = [a for a in api('/api/assets') if a.get('parent_job') == job]
            img = api(f"/api/assets/{outs[0]['id']}/file", raw=True)
            open(out, 'wb').write(img)
            made += 1; print('made', out, flush=True)
        except Exception as e:
            failed += 1; print('FAIL', f'{name}_{v}:', str(e)[:180], flush=True)
print(f'done: {made} made, {failed} failed, {len(rows)} prompts')