#!/usr/bin/env python3
"""WARDROBEGOD — one skeleton, infinite fits. The wardrobe generator bench.

Stdlib server + three.js page (no build step, house style — sibling of NPCFACTORY/imagelab).

    python3 server.py                 → http://localhost:8150
    WG_HOST=0.0.0.0 python3 server.py → expose on the tailnet

What it does:
  · library panes over bodies / garments / generated images / finished outfits
  · loads FBX or GLB on a 3D stage (FBX auto-converts to GLB via local headless Blender)
  · Blender jobs: convert · scale-to-height · decimate-keep-weights · FIT (weight-transfer
    a garment onto a rigged body) · ASSEMBLE (body + fitted garments → one dressed GLB)
  · the garment GENERATOR pipeline (MODELBEAST, $0): flux image → bg_remove →
    either a hanging-garment TEXTURE (racks/paper-doll) or trellis_mac → rigid 3D
    wearable (hats/shoes/bags — bone-parent those; cloth that must bend gets FIT instead)

Env: MB_HOST (default m3ultra :8777) · MB_TOKEN (bearer; generator disabled without it)
     WG_PORT (8150) · WG_HOST (127.0.0.1)
"""
import json, mimetypes, os, re, shutil, subprocess, threading, time, urllib.parse, urllib.request, uuid
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

ROOT = os.path.dirname(os.path.abspath(__file__))
LIB = os.path.join(ROOT, 'library')
DIRS = {'bodies': os.path.join(LIB, 'bodies'), 'garments': os.path.join(LIB, 'garments'),
        'gen': os.path.join(LIB, 'gen'), 'out': os.path.join(ROOT, 'out')}
EXTRA_BODIES = [os.path.expanduser('~/Documents/anatomy'),
                os.path.expanduser('~/Documents/thriftgod/web/assets/models')]
CACHE = os.path.join(ROOT, '.glbcache')
BLENDER = os.environ.get('WG_BLENDER', '/Applications/Blender.app/Contents/MacOS/Blender')
OPS = os.path.join(ROOT, 'blender_ops.py')
MB = os.environ.get('MB_HOST', 'http://100.89.131.57:8777')
MB_TOKEN = os.environ.get('MB_TOKEN', '')
PORT = int(os.environ.get('WG_PORT', 8150))
HOST = os.environ.get('WG_HOST', '127.0.0.1')
for d in list(DIRS.values()) + [CACHE]:
    os.makedirs(d, exist_ok=True)

JOBS = {}                       # id → {status, note, out, log}
MODEL_EXT = ('.glb', '.gltf', '.fbx', '.obj')


def slug(s):
    return re.sub(r'[^a-z0-9]+', '-', (s or '').lower()).strip('-')[:60] or 'x'


def allowed(path):
    p = os.path.realpath(path)
    roots = list(DIRS.values()) + EXTRA_BODIES + [CACHE]
    return any(p.startswith(os.path.realpath(r) + os.sep) or p == os.path.realpath(r) for r in roots)


def scan():
    out = {}
    for key, d in DIRS.items():
        exts = MODEL_EXT if key in ('bodies', 'garments', 'out') else ('.png', '.jpg', '.webp')
        rows = []
        dirs = [d] + (EXTRA_BODIES if key == 'bodies' else [])
        for dd in dirs:
            if not os.path.isdir(dd):
                continue
            for f in sorted(os.listdir(dd)):
                if f.lower().endswith(exts):
                    p = os.path.join(dd, f)
                    rows.append({'name': f, 'path': p, 'kb': os.path.getsize(p) // 1024,
                                 'home': os.path.basename(dd)})
        out[key] = rows
    return out


def run_blender(args, jid):
    """One headless Blender op; stdout tail lands in the job log."""
    cmd = [BLENDER, '-b', '--python', OPS, '--'] + [str(a) for a in args]
    r = subprocess.run(cmd, capture_output=True, text=True, timeout=1800)
    tail = '\n'.join((r.stdout + r.stderr).strip().splitlines()[-12:])
    JOBS[jid]['log'] = tail
    if r.returncode != 0 or 'Error' in r.stderr:
        raise RuntimeError(tail[-400:])


def glb_of(path, jid=None):
    """FBX/OBJ → cached GLB (stage + ops always speak GLB); GLB passes through."""
    if path.lower().endswith(('.glb', '.gltf')):
        return path
    tgt = os.path.join(CACHE, slug(os.path.basename(path)) + '.glb')
    if not os.path.exists(tgt) or os.path.getmtime(tgt) < os.path.getmtime(path):
        run_blender(['convert', path, tgt], jid or new_job('convert', 'convert ' + os.path.basename(path)))
    return tgt


def new_job(kind, note):
    jid = uuid.uuid4().hex[:10]
    JOBS[jid] = {'status': 'running', 'kind': kind, 'note': note, 'out': None, 'log': '', 't': time.time()}
    return jid


def job_thread(jid, fn):
    def go():
        try:
            JOBS[jid]['out'] = fn(jid)
            JOBS[jid]['status'] = 'done'
        except Exception as e:
            JOBS[jid]['status'] = 'error'
            JOBS[jid]['log'] = (JOBS[jid]['log'] + '\n' + str(e)).strip()[-1500:]
    threading.Thread(target=go, daemon=True).start()


# ---------- MODELBEAST client (token stays server-side, never in the page) ----------
def mb_req(path, data=None, raw=None, ctype='application/json'):
    req = urllib.request.Request(MB + path, data=raw if raw is not None else (json.dumps(data).encode() if data else None))
    req.add_header('Authorization', 'Bearer ' + MB_TOKEN)
    if data is not None or raw is not None:
        req.add_header('Content-Type', ctype)
    with urllib.request.urlopen(req, timeout=120) as r:
        body = r.read()
    try:
        return json.loads(body)
    except ValueError:
        return body


def mb_wait(job_id, jid, timeout=900):
    t0 = time.time()
    while time.time() - t0 < timeout:
        j = mb_req(f'/api/jobs/{job_id}')
        st = j.get('status')
        JOBS[jid]['note'] = f"farm: {st}"
        if st in ('done', 'completed', 'succeeded'):
            return j
        if st in ('error', 'failed', 'cancelled'):
            raise RuntimeError('farm job failed: ' + str(j.get('error') or st))
        time.sleep(3)
    raise RuntimeError('farm job timed out')


def mb_outputs(job_id):
    assets = mb_req(f'/api/assets?parent_job={job_id}')
    rows = assets.get('assets', assets) if isinstance(assets, dict) else assets
    return rows or []


def mb_download(asset, tgt):
    aid = asset.get('id') or asset.get('asset_id')
    data = mb_req(f'/api/assets/{aid}/file')
    with open(tgt, 'wb') as f:
        f.write(data if isinstance(data, bytes) else json.dumps(data).encode())
    return tgt


def mb_upload(path):
    import mimetypes as mt
    boundary = uuid.uuid4().hex
    name = os.path.basename(path)
    ctype = mt.guess_type(path)[0] or 'application/octet-stream'
    with open(path, 'rb') as f:
        payload = f.read()
    body = (f'--{boundary}\r\nContent-Disposition: form-data; name="file"; filename="{name}"\r\n'
            f'Content-Type: {ctype}\r\n\r\n').encode() + payload + f'\r\n--{boundary}--\r\n'.encode()
    out = mb_req('/api/assets', raw=body, ctype=f'multipart/form-data; boundary={boundary}')
    return out.get('id') or out.get('asset_id') or (out.get('asset') or {}).get('id')


# ---------- HTTP ----------
class H(BaseHTTPRequestHandler):
    def log_message(self, *a):
        pass

    def send(self, code, body, ctype='application/json'):
        self.send_response(code)
        self.send_header('Content-Type', ctype)
        self.send_header('Content-Length', str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def j(self, obj, code=200):
        self.send(code, json.dumps(obj).encode())

    def do_GET(self):
        u = urllib.parse.urlparse(self.path)
        q = dict(urllib.parse.parse_qsl(u.query))
        if u.path == '/':
            return self.send(200, open(os.path.join(ROOT, 'web', 'index.html'), 'rb').read(), 'text/html')
        if u.path == '/api/lib':
            return self.j({'lib': scan(), 'mb': bool(MB_TOKEN), 'blender': os.path.exists(BLENDER)})
        if u.path.startswith('/api/job/'):
            jid = u.path.rsplit('/', 1)[1]
            return self.j(JOBS.get(jid) or {'status': 'unknown'})
        if u.path == '/file':
            p = q.get('p', '')
            if not allowed(p) or not os.path.isfile(p):
                return self.j({'error': 'nope'}, 403)
            return self.send(200, open(p, 'rb').read(), mimetypes.guess_type(p)[0] or 'application/octet-stream')
        if u.path == '/glb':                      # stage loader: any model file → GLB (convert+cache)
            p = q.get('p', '')
            if not allowed(p) or not os.path.isfile(p):
                return self.j({'error': 'nope'}, 403)
            try:
                g = glb_of(p)
            except Exception as e:
                return self.j({'error': str(e)[-300:]}, 500)
            return self.send(200, open(g, 'rb').read(), 'model/gltf-binary')
        return self.j({'error': 'not found'}, 404)

    def do_POST(self):
        u = urllib.parse.urlparse(self.path)
        q = dict(urllib.parse.parse_qsl(u.query))
        n = int(self.headers.get('Content-Length') or 0)
        raw = self.rfile.read(n) if n else b''

        if u.path == '/api/upload':               # file input → library dir
            to, name = q.get('to', 'bodies'), os.path.basename(q.get('name', 'upload.glb'))
            if to not in DIRS or not name.lower().endswith(MODEL_EXT + ('.png', '.jpg', '.webp')):
                return self.j({'error': 'bad target'}, 400)
            p = os.path.join(DIRS[to], name)
            with open(p, 'wb') as f:
                f.write(raw)
            return self.j({'ok': True, 'path': p})

        body = json.loads(raw or b'{}')

        if u.path == '/api/blender':              # {op, args:{...}} → background Blender job
            op = body.get('op')
            a = body.get('args', {})
            jid = new_job(op, body.get('note', op))

            def fx(jid):
                if op == 'scale':
                    src = glb_of(a['path'], jid)
                    out = os.path.join(DIRS['bodies'], slug(a.get('name') or os.path.basename(a['path']).rsplit('.', 1)[0]) + '.glb')
                    run_blender(['scale', src, out, a['height']], jid)
                elif op == 'decimate':
                    src = glb_of(a['path'], jid)
                    out = os.path.join(DIRS['bodies'], slug(os.path.basename(a['path']).rsplit('.', 1)[0]) + f"-{int(a['tris'])//1000}k.glb")
                    run_blender(['decimate', src, out, int(a['tris'])], jid)
                elif op == 'fit':
                    b, g = glb_of(a['body'], jid), glb_of(a['garment'], jid)
                    mode = a.get('mode', 'garment')
                    base = slug(os.path.basename(a['garment']).rsplit('.', 1)[0])
                    out = os.path.join(DIRS['garments'] if mode == 'garment' else DIRS['out'],
                                       base + ('-fitted.glb' if mode == 'garment' else '-dressed.glb'))
                    run_blender(['fit', b, g, out, mode, a.get('inflate', 3)], jid)
                elif op == 'assemble':
                    b = glb_of(a['body'], jid)
                    gs = [glb_of(g, jid) for g in a['garments']]
                    out = os.path.join(DIRS['out'], slug(a.get('name') or 'outfit') + '.glb')
                    run_blender(['assemble', b, out] + gs, jid)
                elif op == 'convert':
                    src = a['path']
                    out = os.path.join(DIRS['bodies'], slug(os.path.basename(src).rsplit('.', 1)[0]) + '.glb')
                    run_blender(['convert', src, out], jid)
                else:
                    raise RuntimeError('unknown op')
                return out
            job_thread(jid, fx)
            return self.j({'job': jid})

        if u.path == '/api/gen':                  # the garment pipeline, step 1: flux image
            if not MB_TOKEN:
                return self.j({'error': 'MB_TOKEN not set on this box — export MB_TOKEN=… and restart'}, 400)
            prompt = body.get('prompt', '')
            style = ('product photo of a single {} laid flat on a plain white background, '
                     'front view, no mannequin, no person, no text, soft even lighting, '
                     '1990s Australian op-shop garment, slightly worn').format(prompt)
            jid = new_job('gen', 'flux: ' + prompt[:40])

            def fx(jid):
                j = mb_req('/api/jobs', {'operator': 'flux_local', 'asset_id': None,
                                         'params': {'prompt': style, 'model': 'flux2-klein-4b',
                                                    'steps': 4, 'width': 768, 'height': 768,
                                                    'seed': int(time.time()) % 99991}})
                done = mb_wait(j.get('id') or j.get('job_id'), jid)
                outs = mb_outputs(j.get('id') or j.get('job_id'))
                if not outs:
                    raise RuntimeError('no image came back')
                tgt = os.path.join(DIRS['gen'], slug(prompt) + '.png')
                return mb_download(outs[0], tgt)
            job_thread(jid, fx)
            return self.j({'job': jid})

        if u.path == '/api/rmbg':                 # step 2: cut it out
            if not MB_TOKEN:
                return self.j({'error': 'MB_TOKEN not set'}, 400)
            p = body.get('path', '')
            if not allowed(p):
                return self.j({'error': 'nope'}, 403)
            jid = new_job('rmbg', 'cutout ' + os.path.basename(p))

            def fx(jid):
                aid = mb_upload(p)
                j = mb_req('/api/jobs', {'operator': 'bg_remove_local', 'asset_id': aid, 'params': {}})
                mb_wait(j.get('id') or j.get('job_id'), jid)
                outs = mb_outputs(j.get('id') or j.get('job_id'))
                if not outs:
                    raise RuntimeError('no cutout came back')
                tgt = p.rsplit('.', 1)[0] + '-cut.png'
                return mb_download(outs[0], tgt)
            job_thread(jid, fx)
            return self.j({'job': jid})

        if u.path == '/api/to3d':                 # step 3a: rigid wearable via TRELLIS (hats/shoes/bags)
            if not MB_TOKEN:
                return self.j({'error': 'MB_TOKEN not set'}, 400)
            p = body.get('path', '')
            if not allowed(p):
                return self.j({'error': 'nope'}, 403)
            jid = new_job('to3d', '3D: ' + os.path.basename(p) + ' (~5 min on the farm)')

            def fx(jid):
                aid = mb_upload(p)
                j = mb_req('/api/jobs', {'operator': 'trellis_mac', 'asset_id': aid, 'params': {}})
                mb_wait(j.get('id') or j.get('job_id'), jid, timeout=1200)
                outs = mb_outputs(j.get('id') or j.get('job_id'))
                glbs = [o for o in outs if str(o.get('name', o.get('filename', ''))).lower().endswith('.glb')] or outs
                if not glbs:
                    raise RuntimeError('no mesh came back')
                tgt = os.path.join(DIRS['garments'], slug(os.path.basename(p).rsplit('.', 1)[0]) + '-rigid.glb')
                return mb_download(glbs[0], tgt)
            job_thread(jid, fx)
            return self.j({'job': jid})

        if u.path == '/api/hangable':             # step 3b: cutout → hanging-garment texture pack
            p = body.get('path', '')
            if not allowed(p):
                return self.j({'error': 'nope'}, 403)
            tgt = os.path.join(DIRS['garments'], os.path.basename(p))
            shutil.copyfile(p, tgt)
            return self.j({'ok': True, 'path': tgt,
                           'note': 'texture saved to garments — use on rack planes (fittings.js garment({image}))'})

        return self.j({'error': 'not found'}, 404)


if __name__ == '__main__':
    print(f'WARDROBEGOD on http://{HOST}:{PORT}  ·  blender={os.path.exists(BLENDER)}  ·  farm token={"yes" if MB_TOKEN else "NO (generator off)"}')
    ThreadingHTTPServer((HOST, PORT), H).serve_forever()