#!/usr/bin/env python3
"""Post-processing follow-up on the 54 published benchmark outputs — PAID (small): how far is each raw engine output
from a game target, measured through the public PicoBerry API, no human scoring.

Steps per run (results.json → runs[]):
  1. remesh  : POST /v1/assets/{id}/remesh   {engine: 'pb-remesh', polycount: TARGET}  — retopology to a game-asset budget
  2. uv      : POST /v1/assets/{remeshed}/uv-unwrap {engine: 'pb-uv'}                  — editable UV layout + texture rebake
  3. rig     : POST /v1/assets/{id}/animate  {engine: 'tripo-rig'} and {engine: 'meshy-rig'} — character inputs only,
               on the ORIGINAL output (the question is "can this raw output be rigged"), rigging only, no clip.
Recorded: status, wall time (5 s polling), faces/vertices before → after (files.stats), list price, output asset id, error text.
Resumable: postprocess.json is rewritten after every step; already-recorded steps are skipped on restart.

Usage:
  PICOBERRY_API_KEY=pb_live_... python3 postprocess.py [--target 20000] [--concurrency 6] [--only remesh,uv,rig]
  python3 postprocess.py --summarize [postprocess.json]      # FREE: re-aggregate, never calls the API
"""
import json, os, sys, time, threading, urllib.request, urllib.error
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime, timezone

BASE = os.environ.get('PICOBERRY_API_BASE', 'https://api.picoberry.ai').rstrip('/')
KEY = os.environ.get('PICOBERRY_API_KEY', '')
HERE = os.path.dirname(os.path.abspath(__file__))
RESULTS = os.path.join(HERE, 'results.json'); OUT = os.path.join(HERE, 'postprocess.json')
lock = threading.Lock()
now = lambda: datetime.now(timezone.utc).isoformat(timespec='seconds').replace('+00:00', 'Z')

def arg(name, default=None):
    if name in sys.argv:
        i = sys.argv.index(name); return sys.argv[i + 1] if len(sys.argv) > i + 1 else default
    return default

def req(method, path, body=None, timeout=120):
    headers = {'Authorization': f'Bearer {KEY}'}; data = None
    if body is not None:
        data = json.dumps(body).encode(); headers['Content-Type'] = 'application/json'
    r = urllib.request.Request(BASE + path, data=data, method=method, headers=headers)
    try:
        with urllib.request.urlopen(r, timeout=timeout) as res:
            return res.status, json.loads(res.read().decode() or '{}')
    except urllib.error.HTTPError as e:
        raw = e.read().decode(errors='replace')
        try: return e.code, json.loads(raw)
        except Exception: return e.code, {'raw': raw[:300]}

def poll(asset_id, max_s=1800, every=5):
    t0 = time.time()
    while time.time() - t0 < max_s:
        st, d = req('GET', f'/v1/assets/{asset_id}', timeout=60)
        a = d.get('data') if st == 200 else None
        if a and a.get('taskStatus') in (2, 3):
            return a
        time.sleep(every)
    return None

def stats_of(asset):
    url = ((asset or {}).get('files') or {}).get('stats')
    if not url: return None
    try:
        with urllib.request.urlopen(url, timeout=60) as res: return json.loads(res.read().decode())
    except Exception as e:
        return {'error': str(e)[:100]}

def catalog():
    """List prices at run time for the post-processing engines (remesh / uv-unwrap / animate categories)."""
    out = {}
    for cat in ('remesh', 'uv-unwrap', 'animate'):
        st, d = req('GET', f'/v1/models?category={cat}')
        for m in (d.get('data') or []):
            out[m.get('label') or m.get('id')] = {'category': cat, 'cost': m.get('cost'), 'id': m.get('id')}
    return out

def step(kind, source_id, body, path):
    """Submit one post-processing task and wait. Returns a record; never raises."""
    t0 = time.time(); submitted = now()
    st, d = req('POST', path, body)
    rec = {'kind': kind, 'sourceAssetId': source_id, 'request': body, 'submittedAt': submitted}
    if st not in (200, 201) or not d.get('data'):
        rec.update(status='rejected', httpStatus=st, error=json.dumps(d)[:300], wallSeconds=round(time.time() - t0))
        return rec
    new_id = d['data']['id']
    a = poll(new_id)
    rec.update(assetId=new_id, completedAt=now(), wallSeconds=round(time.time() - t0))
    if not a:
        rec.update(status='timeout'); return rec
    rec['status'] = 'succeeded' if a.get('taskStatus') == 2 else 'failed'
    if rec['status'] == 'failed':
        rec['error'] = (a.get('failureReason') or a.get('error') or a.get('status') or 'failed')
    s = stats_of(a)
    if s: rec['stats'] = {k: s.get(k) for k in ('vertices', 'faces', 'topology') if k in s}
    files = a.get('files') or {}
    rec['hasTextures'] = bool(files.get('hasTextures'))
    if kind == 'rig':
        # rig report fields, when the pipeline exposes them (bone count etc.)
        for k in ('rigReport', 'rig', 'bones'):
            if a.get(k) is not None: rec[k] = a[k]
    return rec

def load():
    if os.path.exists(OUT):
        return json.load(open(OUT))
    return {'generatedAt': now(), 'target': None, 'catalog': {}, 'runs': []}

def save(data):
    with lock:
        data['updatedAt'] = now()
        tmp = OUT + '.tmp'; json.dump(data, open(tmp, 'w'), indent=1); os.replace(tmp, OUT)

def summarize(path):
    data = json.load(open(path)); runs = data['runs']
    by = {}
    for r in runs:
        e = by.setdefault(r['engine'], {'remesh': [], 'uv': [], 'rig:tripo-rig': [], 'rig:meshy-rig': []})
        for k in ('remesh', 'uv'):
            if r.get(k): e[k].append(r[k])
        for g in (r.get('rig') or []):
            e[f"rig:{g['request']['engine']}"].append(g)
    def line(recs):
        ok = [x for x in recs if x.get('status') == 'succeeded']
        walls = sorted(x['wallSeconds'] for x in ok if x.get('wallSeconds') is not None)
        med = walls[len(walls) // 2] if walls else None
        return f"{len(ok)}/{len(recs)} ok, median {med}s" if recs else '—'
    print(f"target polycount {data.get('target')}  (prices: " + ', '.join(f"{k} {v['cost']}cr" for k, v in data.get('catalog', {}).items() if k in ('pb-remesh', 'pb-uv', 'tripo-rig', 'meshy-rig')) + ')')
    print(f"{'engine':17s} {'remesh (pb-remesh)':26s} {'uv (pb-uv)':26s} {'rig tripo':22s} {'rig meshy':22s}")
    for e, v in sorted(by.items()):
        print(f"{e:17s} {line(v['remesh']):26s} {line(v['uv']):26s} {line(v['rig:tripo-rig']):22s} {line(v['rig:meshy-rig']):22s}")
    ok_remesh = [r for r in runs if (r.get('remesh') or {}).get('status') == 'succeeded' and (r['remesh'].get('stats') or {}).get('faces')]
    if ok_remesh:
        print('faces before → after (median per engine):')
        for e in sorted(by):
            rs = [r for r in ok_remesh if r['engine'] == e]
            if not rs: continue
            bef = sorted(r['sourceFaces'] for r in rs); aft = sorted(r['remesh']['stats']['faces'] for r in rs)
            print(f"  {e:17s} {bef[len(bef)//2]:>10,} → {aft[len(aft)//2]:>7,}")

if __name__ == '__main__':
    if '--summarize' in sys.argv:
        summarize(arg('--summarize', OUT) or OUT); raise SystemExit(0)
    if not KEY:
        raise SystemExit('PICOBERRY_API_KEY is not set — this mode spends credits (≈10–15 per step). Use --summarize for the free re-aggregation.')
    target = int(arg('--target', '20000')); conc = int(arg('--concurrency', '6'))
    only = set((arg('--only', 'remesh,uv,rig') or '').split(','))
    src = json.load(open(RESULTS)); src_runs = [r for r in src['runs'] if r.get('status') == 'succeeded']
    data = load(); data['target'] = target
    if not data['catalog']: data['catalog'] = catalog()
    have = {(r['engine'], r['input'], r['repeat']): r for r in data['runs']}
    for r in src_runs:
        k = (r['engine'], r['input'], r['repeat'])
        if k not in have:
            rec = {'engine': r['engine'], 'input': r['input'], 'repeat': r['repeat'], 'sourceAssetId': r['assetId'],
                   'sourceFaces': (r.get('stats') or {}).get('faces'), 'sourceVertices': (r.get('stats') or {}).get('vertices')}
            data['runs'].append(rec); have[k] = rec
    save(data)

    def work(rec):
        sid = rec['sourceAssetId']
        if 'remesh' in only and not rec.get('remesh'):
            rec['remesh'] = step('remesh', sid, {'engine': 'pb-remesh', 'polycount': target}, f'/v1/assets/{sid}/remesh'); save(data)
            print(f"remesh {rec['engine']:17s} {rec['input']:14s} #{rec['repeat']} {rec['remesh']['status']:9s} {rec['remesh'].get('wallSeconds')}s faces {rec.get('sourceFaces')} → {(rec['remesh'].get('stats') or {}).get('faces')}", flush=True)
        if 'uv' in only and not rec.get('uv') and (rec.get('remesh') or {}).get('status') == 'succeeded':
            rid = rec['remesh']['assetId']
            rec['uv'] = step('uv', rid, {'engine': 'pb-uv'}, f'/v1/assets/{rid}/uv-unwrap'); save(data)
            print(f"uv     {rec['engine']:17s} {rec['input']:14s} #{rec['repeat']} {rec['uv']['status']:9s} {rec['uv'].get('wallSeconds')}s", flush=True)
        if 'rig' in only and rec['input'] == 'character':
            done = {g['request']['engine'] for g in (rec.get('rig') or [])}
            for eng in ('tripo-rig', 'meshy-rig'):
                if eng in done: continue
                g = step('rig', sid, {'engine': eng}, f'/v1/assets/{sid}/animate')
                rec.setdefault('rig', []).append(g); save(data)
                print(f"rig    {rec['engine']:17s} {eng:10s} #{rec['repeat']} {g['status']:9s} {g.get('wallSeconds')}s {g.get('error') or ''}", flush=True)

    with ThreadPoolExecutor(max_workers=conc) as ex:
        list(ex.map(work, data['runs']))
    data['finishedAt'] = now(); save(data)
    print('wrote', OUT); summarize(OUT)