#!/usr/bin/env python3
"""mbgen.py — batch text→image on the MODELBEAST farm (flux_local / comfyui_sd).

Reads a JSON job spec, submits with a concurrency cap, polls, and pulls each
output image back by `parent_job` (NEVER newest-asset: two concurrent jobs
raced and crossed outputs on 2026-07-16).

Spec: [{"out": "path.png", "operator": "flux_local", "params": {...}}, ...]
Usage: mbgen.py spec.json [--max-active 4]
Token: MB_TOKEN env, else ~/Documents/fluxgod-work/.env
"""
import json, os, sys, time, urllib.request, urllib.error

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


def token():
    t = os.environ.get('MB_TOKEN')
    if t:
        return t
    for e in ENVS:
        p = os.path.expanduser(e)
        if os.path.exists(p):
            for line in open(p):
                if line.startswith('MB_TOKEN='):
                    return line.split('=', 1)[1].strip().strip('"\'')
    sys.exit('no MB_TOKEN')


TOK = token()


def req(path, data=None, headers=None, raw=False, timeout=180):
    h = {'Authorization': f'Bearer {TOK}'}
    h.update(headers or {})
    r = urllib.request.Request(HOST + path, data=data, headers=h)
    body = urllib.request.urlopen(r, timeout=timeout).read()
    if raw:
        return body
    s = body.decode('utf-8', 'replace')
    # job logs carry raw control chars — strip before parsing
    return json.loads(''.join(c if c >= ' ' or c in '\t' else ' ' for c in s))


def submit(item):
    payload = {'operator': item['operator'], 'params': item['params']}
    j = req('/api/jobs', data=json.dumps(payload).encode(),
            headers={'Content-Type': 'application/json'})
    return j['id']


def fetch_output(jid, out):
    """Pull the image asset belonging to THIS job (parent_job match only)."""
    a = req('/api/assets?limit=200')
    items = a if isinstance(a, list) else a.get('items', [])
    mine = [x for x in items
            if x.get('parent_job') == jid
            and str(x.get('name', x.get('filename', ''))).lower().endswith(('.png', '.jpg', '.jpeg', '.webp'))]
    if not mine:
        return False
    blob = req(f"/api/assets/{mine[0]['id']}/file", raw=True)
    os.makedirs(os.path.dirname(out) or '.', exist_ok=True)
    with open(out, 'wb') as f:
        f.write(blob)
    return True


def main():
    spec = json.load(open(sys.argv[1]))
    cap = 4
    if '--max-active' in sys.argv:
        cap = int(sys.argv[sys.argv.index('--max-active') + 1])
    todo = [dict(i, _state='pending') for i in spec
            if not (os.path.exists(i['out']) and os.path.getsize(i['out']) > 1000)]
    skipped = len(spec) - len(todo)
    if skipped:
        print(f'[mb] {skipped} already done, skipping')
    active, done, failed = {}, 0, []
    t0 = time.time()
    while todo or active:
        while todo and len(active) < cap:
            it = todo.pop(0)
            try:
                jid = submit(it)
                active[jid] = it
                print(f'[mb] + {os.path.basename(it["out"])} → {jid[:8]}', flush=True)
            except urllib.error.HTTPError as e:
                if e.code in (429, 409):      # queue full — retry this one later
                    todo.insert(0, it)
                    time.sleep(10)
                    break
                failed.append((it['out'], f'submit {e.code}'))
        if not active:
            continue
        time.sleep(6)
        for jid in list(active):
            try:
                st = req(f'/api/jobs/{jid}').get('status')
            except Exception:
                continue
            if st in ('done', 'error', 'cancelled'):
                it = active.pop(jid)
                if st == 'done' and fetch_output(jid, it['out']):
                    done += 1
                    print(f'[mb] ✓ {os.path.basename(it["out"])}  ({done} done, {len(todo)} left, {int(time.time()-t0)}s)', flush=True)
                else:
                    failed.append((it['out'], st))
                    print(f'[mb] ✗ {os.path.basename(it["out"])} — {st}', flush=True)
    print(f'\n[mb] {done} generated, {len(failed)} failed in {int(time.time()-t0)}s')
    for f, why in failed:
        print(f'  ✗ {os.path.basename(f)} — {why}')
    return 1 if failed else 0


if __name__ == '__main__':
    sys.exit(main())