#!/usr/bin/env python3
"""archetype_batch.py — the 12-archetype factory run (all local, all free).

Per archetype: flux_local sheet (solid green, A-pose, FITTED clothing — open vests are
the hard rig case) -> bg_remove_local -> trellis_mac mesh -> MIRPAMO preclean+proxy rig
(on the box) -> retarget wan_walk_t1 -> library/banks/<name>_at_wanwalk.glb.

Idempotent: skips any archetype whose bank already exists. Progress to stdout (run it
under nohup/tee). MB flow per the house client pattern (control-char-safe, parent_job).
"""
import json
import os
import subprocess
import sys
import time
import urllib.request
import uuid

HERE = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
BANKS = os.path.join(HERE, 'library', 'banks')
SHEETS = os.path.join(HERE, 'library', 'sheets')
HOST = 'http://100.89.131.57:8777'
BOX = 'm3ultra@100.89.131.57'
os.makedirs(SHEETS, exist_ok=True)

ARCHETYPES = {
    'nanna':    'an elderly Australian woman in a floral dress and a buttoned lavender cardigan, grey permed hair, flat shoes',
    'postie':   'an Australia Post postal worker in 1990s uniform, red polo shirt and navy shorts, wide-brim hat, sturdy shoes',
    'goth':     'a young goth woman in 1990s style, long black fitted dress, black boots, dark bob haircut, pale makeup',
    'skater':   'a teenage skater boy in 1990s style, fitted graphic t-shirt, knee-length shorts, skate shoes, backwards cap',
    'surfie':   'a tanned Australian surfer man in 1990s boardshorts and a fitted rash vest, barefoot, sun-bleached hair',
    'schoolkid': 'an Australian school kid in 1990s public school uniform, fitted polo shirt and grey shorts, school shoes, backpack straps',
    'bizdad':   'a middle-aged businessman in a fitted 1990s grey suit, white shirt, tie, black leather shoes',
    'youngmum': 'a young Australian mum in 1990s high-waist jeans and a fitted striped t-shirt, white sneakers, hair in a scrunchie ponytail',
    'busdriver': 'an Australian bus driver in 1990s uniform, fitted light blue short-sleeve shirt with epaulettes, navy trousers, black shoes',
    'cop':      'an Australian police officer in 1990s uniform, fitted light blue shirt, navy trousers, peaked cap, black shoes',
    'paperboy': 'a paperboy in 1990s clothes, fitted t-shirt, jeans rolled at the ankle, canvas sneakers, satchel strap across the chest',
}
SUFFIX = (', standing straight facing camera, arms slightly away from sides in a relaxed A-pose, '
          'flat even lighting, solid green background, photorealistic, entire body in frame head to shoes')


def tok():
    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')


TOK = tok()


def req(p, data=None, headers=None, raw=False):
    h = {'Authorization': f'Bearer {TOK}'}
    h.update(headers or {})
    r = urllib.request.Request(HOST + p, data=data, headers=h)
    b = urllib.request.urlopen(r, timeout=180).read()
    if raw:
        return b
    s = b.decode('utf-8', 'replace')
    return json.loads(''.join(c if c >= ' ' or c in '\t' else ' ' for c in s))


def job(operator, params=None, asset_id=None):
    body = {'operator': operator, 'params': params or {}}
    if asset_id:
        body['asset_id'] = asset_id
    return req('/api/jobs', data=json.dumps(body).encode(),
               headers={'Content-Type': 'application/json'})['id']


def wait(jid, timeout=1800):
    t0 = time.time()
    while time.time() - t0 < timeout:
        time.sleep(8)
        st = req(f'/api/jobs/{jid}').get('status')
        if st in ('done', 'error', 'cancelled'):
            return st
    return 'timeout'


def output_asset(jid, skip=None):
    a = req('/api/assets?limit=60')
    items = a if isinstance(a, list) else a.get('items', [])
    mine = [x for x in items if x.get('parent_job') == jid and x.get('id') != skip]
    return mine[0]['id'] if mine else None


def sh(cmd, timeout=1200):
    return subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=timeout)


def run_one(name, desc):
    bank = os.path.join(BANKS, f'{name}_at_wanwalk.glb')
    if os.path.exists(bank):
        print(f'[{name}] SKIP (bank exists)', flush=True)
        return True
    print(f'[{name}] sheet…', flush=True)
    j1 = job('flux_local', {'prompt': 'full body character reference of ' + desc + SUFFIX,
                            'steps': 6, 'width': 768, 'height': 1344})
    if wait(j1, 300) != 'done':
        print(f'[{name}] FLUX FAILED', flush=True); return False
    sheet = output_asset(j1)
    open(os.path.join(SHEETS, f'{name}.png'), 'wb').write(req(f'/api/assets/{sheet}/file', raw=True))
    j2 = job('bg_remove_local', asset_id=sheet)
    if wait(j2, 300) != 'done':
        print(f'[{name}] KEY FAILED', flush=True); return False
    cut = output_asset(j2, skip=sheet)
    print(f'[{name}] mesh… (trellis)', flush=True)
    j3 = job('trellis_mac', asset_id=cut)
    if wait(j3, 1800) != 'done':
        print(f'[{name}] TRELLIS FAILED', flush=True); return False
    mesh = output_asset(j3, skip=cut)
    raw = req(f'/api/assets/{mesh}/file', raw=True)
    tmp_local = f'/tmp/{name}_mesh.glb'
    open(tmp_local, 'wb').write(raw)
    print(f'[{name}] rig+retarget… ({len(raw)//1048576}MB mesh)', flush=True)
    r = sh(f'scp -q {tmp_local} {BOX}:/tmp/{name}_mesh.glb && '
           f'ssh {BOX} "cd ~/Documents/MIRPAMO && ./bin/mirpamo rig /tmp/{name}_mesh.glb -o /tmp/{name}_rigged.glb && '
           f'./bin/mirpamo anim /tmp/{name}_rigged.glb /tmp/wan_walk_t1.fbx -o /tmp/{name}_bank.glb" && '
           f'scp -q {BOX}:/tmp/{name}_bank.glb {bank}')
    if r.returncode != 0 or not os.path.exists(bank):
        print(f'[{name}] RIG FAILED: {(r.stderr or r.stdout)[-200:]}', flush=True)
        return False
    print(f'[{name}] ✔ banked ({os.path.getsize(bank)//1048576}MB)', flush=True)
    return True


def main():
    only = sys.argv[1:] or list(ARCHETYPES)
    ok = fail = 0
    for name in only:
        try:
            ok += 1 if run_one(name, ARCHETYPES[name]) else 0
        except Exception as e:
            print(f'[{name}] EXCEPTION: {e}', flush=True)
            fail += 1
    print(f'=== batch done: {ok} banked, {len(only) - ok} failed/skipped-fail ===', flush=True)


if __name__ == '__main__':
    main()