#!/usr/bin/env python3
"""LANE DOLL D3 — MAKE MORE wardrobe, free & forever, on MODELBEAST (M3 Ultra, zero dollars).

One command: prompt -> flux_local (text->image on solid green) -> bg_remove_local (keyed cutout) ->
doll_stage (704x1408 overlay) -> media/doll/<slot>_i_<itemkey>.png. Wire the item with apply_fashion.py.

  gen_wardrobe.py batch                stage the curated BATCH below (skips stems already on disk)
  gen_wardrobe.py one <slot> <key> "<garment desc>"   one piece

Solid-green law (memory): flux is prompted on SOLID GREEN — checkerboard/neutral-grey eat garments,
plaid survives. No baked text (flux can't spell — band tees are ART, no readable words). Token from
backnforth/.env; control-char-safe JSON (the localflux.mjs hardening). Mirrors MESHGOD mb_recon.py.
"""
import json, os, re, sys, time, urllib.request, uuid, mimetypes
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import doll_stage

HOST = os.environ.get('MB_HOST', 'http://100.89.131.57:8777')
FLUX_OP = os.environ.get('MB_FLUX_OP', 'flux_local')   # override if the box names it differently
STAGE = 'tools/_gen_stage'                              # raw flux + keyed intermediates land here

# curated Aussie-90s staples the wardrobe is thin on. (slot, itemkey, garment-phrase). Bags deferred.
BATCH = [
    ('top',    'flannelette_shirt',   'red and black checked flannelette shirt, open over a tee'),
    ('top',    'country_road_jumper', 'cream Country Road wool crew-neck jumper'),
    ('top',    'rugby_jumper',        'green and gold hooped rugby jumper with a white collar'),
    ('top',    'king_gee_work_shirt', 'blue King Gee cotton drill work shirt'),
    ('top',    'bonds_singlet',       'white ribbed cotton Bonds singlet'),
    ('top',    'band_tee_abstract',   'black band tee with an abstract screen-printed graphic, no words'),
    ('bottom', 'stubbies_shorts',     'short green Stubbies rugby shorts'),
    ('bottom', 'bike_shorts',         'black lycra bike shorts'),
    ('bottom', 'hard_yakka_work_pants','tan Hard Yakka cotton drill work trousers'),
    ('bottom', 'flannel_pyjama_pants','plaid flannel lounge pants'),
    ('hat',    'akubra_hat',          'brown Akubra felt bush hat with a wide brim'),
    ('hat',    'terry_towelling_hat', 'faded blue terry towelling bucket hat'),
    ('shoes',  'blundstone_boots',    'pair of tan Blundstone elastic-sided leather boots, side by side'),
    ('shoes',  'ugg_boots',           'pair of tan sheepskin ugg boots, side by side'),
    ('shoes',  'explorer_socks_sandals','pair of leather sandals worn with long socks, side by side'),
]


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 in env or backnforth/.env')


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=300).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))   # control-char-safe


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


def run_job(operator, params=None, asset_id=None, timeout=600):
    payload = {'operator': operator, 'params': params or {}}
    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']
    t0 = time.time()
    while True:
        time.sleep(5)
        st = req(f'/api/jobs/{jid}').get('status')
        if st in ('done', 'error', 'cancelled') or time.time() - t0 > timeout:
            break
    if st != 'done':
        raise RuntimeError(f'{operator} job {jid}: {st}')
    a = req('/api/assets')
    items = a if isinstance(a, list) else a.get('items', [])
    own = [x for x in items if x.get('parent_job') == jid and x.get('id') != asset_id]   # parent_job, not "newest"
    if not own:
        raise RuntimeError(f'{operator} job {jid}: no output asset')
    return own[0]['id']


def generate(slot, itemkey, phrase):
    """flux green-bg -> download -> bg_remove -> download keyed -> doll_stage. Returns the staged stem."""
    os.makedirs(STAGE, exist_ok=True)
    prompt = (f'1990s Australian {phrase}, product photo, flat lay, centered, solid bright green background, '
              f'no text, no words, no logos, no person, no hanger')
    raw = os.path.join(STAGE, f'{slot}_{itemkey}_flux.png')
    keyed = os.path.join(STAGE, f'{slot}_{itemkey}_keyed.png')
    gid = run_job(FLUX_OP, params={'prompt': prompt, 'width': 1024, 'height': 1024})
    open(raw, 'wb').write(req(f'/api/assets/{gid}/file', raw=True))
    kid = run_job('bg_remove_local', asset_id=upload(raw))
    open(keyed, 'wb').write(req(f'/api/assets/{kid}/file', raw=True))
    stem = doll_stage.stage(keyed, slot, itemkey)
    print(f'  generated {stem}  <- "{phrase}"', flush=True)
    return stem


def cmd_batch():
    done = skip = fail = 0
    for slot, key, phrase in BATCH:
        if os.path.exists(os.path.join(doll_stage.DOLL_DIR, f'{slot}_i_{key}.png')):
            skip += 1
            continue
        try:
            generate(slot, key, phrase); done += 1
        except Exception as e:
            print(f'  FAIL {slot}_i_{key}: {e}', flush=True); fail += 1
    print(f'=== {done} generated, {skip} skipped, {fail} failed. Wire them: add items via apply_fashion.py ===')


if __name__ == '__main__':
    c = sys.argv[1] if len(sys.argv) > 1 else 'batch'
    if c == 'one':          # one <slot> <itemkey> "<garment desc>"
        generate(sys.argv[2], sys.argv[3], sys.argv[4])
    else:
        cmd_batch()