#!/usr/bin/env python3
"""mbgen.py — VINYLGOD asset gen via MODELBEAST (free, local farm).

Usage:
  python3 tools/mbgen.py "prompt text" out.png              # flux_local text->image
  python3 tools/mbgen.py --op bg_remove_local in.png out.png
Token: MB_TOKEN env, else ~/Documents/backnforth/.env (never print it).
Contract per ~/Documents/MESHGOD/scripts/mb_recon.py (verified 2026-07-16):
POST /api/jobs -> poll /api/jobs/{id} -> match own asset by parent_job.
Job logs carry raw control chars; strip before json.loads.
"""
import json, mimetypes, os, sys, time, urllib.request, uuid

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


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 {})
    body = urllib.request.urlopen(urllib.request.Request(HOST + path, data=data, headers=h), timeout=120).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 upload(path):
    boundary = uuid.uuid4().hex
    ctype = mimetypes.guess_type(path)[0] or 'application/octet-stream'
    body = (f'--{boundary}\r\nContent-Disposition: form-data; name="file"; filename="{os.path.basename(path)}"\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 run(op, params, out, asset_id=None):
    payload = {'operator': op, 'params': params}
    if asset_id:
        payload['asset_id'] = asset_id
    j = req('/api/jobs', data=json.dumps(payload).encode(), headers={'Content-Type': 'application/json'})
    jid = j['id']
    print(f'[mb] job {jid} ({op}) submitted', flush=True)
    t0 = time.time()
    while True:
        time.sleep(5)
        st = req(f'/api/jobs/{jid}').get('status')
        print(f'[mb] {st} ({int(time.time() - t0)}s)', flush=True)
        if st in ('done', 'error', 'cancelled'):
            break
    if st != 'done':
        sys.exit(f'[mb] job ended {st}')
    a = req('/api/assets?limit=100')
    items = a if isinstance(a, list) else a.get('items', [])
    mine = [x for x in items if x.get('parent_job') == jid]
    if not mine:
        sys.exit('[mb] no output asset with matching parent_job')
    # some operators emit a params .json alongside the real output — skip it
    media = [x for x in mine if not str(x.get('name', '')).endswith('.json')] or mine
    open(out, 'wb').write(req(f"/api/assets/{media[0]['id']}/file", raw=True))
    print(f'[mb] wrote {out} ({os.path.getsize(out) // 1024}KB)')


if __name__ == '__main__':
    args = sys.argv[1:]
    op = 'flux_local'
    if args and args[0] == '--op':
        op = args[1]
        args = args[2:]
    if op == 'flux_local':
        run(op, {'prompt': args[0]}, args[1])
    else:  # input-file operators
        run(op, {}, args[1], asset_id=upload(args[0]))