#!/usr/bin/env python3
"""mb.py — BLOBBO client for the MODELBEAST farm (m3ultra central queue :8777).

Contract per ~/Documents/MESHGOD/scripts/mb_recon.py (verified 2026-07-16):
Bearer MB_TOKEN (env, else ~/Documents/backnforth/.env), POST /api/jobs,
poll GET /api/jobs/{id}, retrieve own assets by parent_job (never newest-glob —
concurrent jobs race). Job JSON may carry raw control chars — strip pre-parse.
Guest tier: max 4 active jobs, enforced here with a submit window.

Usage:
  mb.py gen "<prompt>" out.png            one flux_local image
  mb.py batch prompts.json outdir/        [{"name":..,"prompt":..,"operator"?,"params"?}]
  mb.py img2glb in.png out.glb [operator] image -> mesh (default hunyuan3d_mlx)
  mb.py jobs                              list my recent jobs
"""
import json, mimetypes, os, sys, time, urllib.request, urllib.error, uuid

HOST = os.environ.get('MB_HOST', 'http://100.89.131.57:8777')
MAX_ACTIVE = 4


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 {})
    r = urllib.request.Request(HOST + path, data=data, headers=h)
    body = urllib.request.urlopen(r, timeout=180).read()
    if raw:
        return body
    s = body.decode('utf-8', 'replace')
    return json.loads(''.join(c if c >= ' ' or c in '\t' else ' ' for c in s))


def submit(operator, params=None, asset_id=None):
    """Returns job id, or None on 429 (farm throttles submits harder than the
    4-active contract — observed 429 on the 4th rapid submit, 2026-07-17)."""
    payload = {'operator': operator, 'params': params or {}}
    if asset_id:
        payload['asset_id'] = asset_id
    try:
        j = req('/api/jobs', data=json.dumps(payload).encode(),
                headers={'Content-Type': 'application/json'})
    except urllib.error.HTTPError as e:
        if e.code == 429:
            return None
        raise
    return j['id']


def status(jid):
    return req(f'/api/jobs/{jid}').get('status')


def upload(path):
    boundary = uuid.uuid4().hex
    ctype = mimetypes.guess_type(path)[0] or 'application/octet-stream'
    name = os.path.basename(path)
    body = (f'--{boundary}\r\nContent-Disposition: form-data; name="file"; filename="{name}"\r\n'
            f'Content-Type: {ctype}\r\n\r\n').encode() + open(path, 'rb').read() + f'\r\n--{boundary}--\r\n'.encode()
    a = req('/api/assets', data=body, headers={'Content-Type': f'multipart/form-data; boundary={boundary}'})
    return a.get('id') or (a.get('items') or [a])[0].get('id')


def outputs_for(jid, suffix=None):
    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]
    if suffix:
        mine = [x for x in mine if str(x.get('filename', x.get('name', ''))).endswith(suffix)]
    return mine


def fetch(asset_id, out):
    data = req(f'/api/assets/{asset_id}/file', raw=True)
    open(out, 'wb').write(data)
    return len(data)


def wait(jid, label='', interval=6):
    t0 = time.time()
    while True:
        st = status(jid)
        if st in ('done', 'error', 'cancelled'):
            print(f'[mb] {label or jid}: {st} ({int(time.time()-t0)}s)', flush=True)
            return st
        time.sleep(interval)


def cmd_gen(prompt, out):
    jid = submit('flux_local', {'prompt': prompt})
    if wait(jid, out) != 'done':
        sys.exit(1)
    outs = outputs_for(jid)
    if not outs:
        sys.exit('[mb] no output asset for job')
    n = fetch(outs[0]['id'], out)
    print(f'[mb] wrote {out} ({n//1024}KB)')


def cmd_batch(spec_path, outdir):
    jobs = json.load(open(spec_path))
    os.makedirs(outdir, exist_ok=True)
    done_names = {os.path.splitext(f)[0] for f in os.listdir(outdir)}
    jobs = [j for j in jobs if j['name'] not in done_names]
    if done_names:
        print(f'[mb] resume: skipping {len(done_names)} already in {outdir}', flush=True)
    pending, active, results = list(jobs), {}, {}
    while pending or active:
        while pending and len(active) < MAX_ACTIVE:
            job = pending[0]
            params = dict(job.get('params', {}))
            if 'prompt' in job:
                params.setdefault('prompt', job['prompt'])
            jid = submit(job.get('operator', 'flux_local'), params)
            if jid is None:
                print('[mb] throttled (429), backing off 15s', flush=True)
                time.sleep(15)
                break
            pending.pop(0)
            active[jid] = job
            print(f'[mb] submitted {job["name"]} -> {jid}', flush=True)
        time.sleep(6)
        for jid in list(active):
            st = status(jid)
            if st in ('done', 'error', 'cancelled'):
                job = active.pop(jid)
                results[job['name']] = st
                if st == 'done':
                    outs = outputs_for(jid)
                    if outs:
                        name = str(outs[0].get('filename', outs[0].get('name', 'out')))
                        ext = os.path.splitext(name)[1] or '.png'
                        out = os.path.join(outdir, job['name'] + ext)
                        n = fetch(outs[0]['id'], out)
                        print(f'[mb] {job["name"]}: done -> {out} ({n//1024}KB)', flush=True)
                    else:
                        print(f'[mb] {job["name"]}: done but no asset visible', flush=True)
                        results[job['name']] = 'done-no-asset'
                else:
                    print(f'[mb] {job["name"]}: {st}', flush=True)
    bad = {k: v for k, v in results.items() if v != 'done'}
    print(f'[mb] batch complete: {len(results)-len(bad)}/{len(results)} ok'
          + (f' — failed: {bad}' if bad else ''))


def cmd_img2glb(img, out, operator='hunyuan3d_mlx'):
    aid = upload(img)
    print(f'[mb] uploaded {os.path.basename(img)} -> asset {aid}')
    jid = submit(operator, {}, asset_id=aid)
    if wait(jid, out, interval=10) != 'done':
        sys.exit(1)
    outs = outputs_for(jid, '.glb')
    if not outs:
        sys.exit('[mb] no glb for job')
    n = fetch(outs[0]['id'], out)
    print(f'[mb] wrote {out} ({n//1024}KB)')


def cmd_jobs():
    j = req('/api/jobs?limit=20')
    items = j if isinstance(j, list) else j.get('items', [])
    for x in items:
        print(x.get('id'), x.get('operator'), x.get('status'))


if __name__ == '__main__':
    cmd = sys.argv[1] if len(sys.argv) > 1 else ''
    if cmd == 'gen':
        cmd_gen(sys.argv[2], sys.argv[3])
    elif cmd == 'batch':
        cmd_batch(sys.argv[2], sys.argv[3])
    elif cmd == 'img2glb':
        cmd_img2glb(*sys.argv[2:5])
    elif cmd == 'jobs':
        cmd_jobs()
    else:
        sys.exit(__doc__)
