#!/usr/bin/env python3
"""mesh_pipeline.py — flux image -> bg_remove -> image->GLB, per spec entry.

Usage: mesh_pipeline.py tools/mesh-wave-1.json assets/meshes/
Retries each farm job up to 3x (the m4 node's broken mflux venv 126s any flux
job routed to it — resubmission re-rolls the load-balancer). Skips any entry
whose .glb already exists, so reruns are cheap.
"""
import json, os, sys, time
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import mb

RETRIES = 3


def run_job(operator, params=None, asset_id=None, label=''):
    for attempt in range(1, RETRIES + 1):
        jid = None
        while jid is None:
            jid = mb.submit(operator, params, asset_id=asset_id)
            if jid is None:
                print(f'[pipe] {label}: throttled, 15s', flush=True)
                time.sleep(15)
        st = mb.wait(jid, f'{label} ({operator} try {attempt})', interval=10)
        if st == 'done':
            return jid
        time.sleep(5)
    raise RuntimeError(f'{label}: {operator} failed {RETRIES}x')


PRIMARY = 'm3ultra@100.89.131.57'


def scp_from_outdir(jid, out, suffix='.glb'):
    """Fallback for a farm bug (2026-07-17): hunyuan jobs run on remote workers
    ship outputs back to the primary's job outdir but never register them as
    assets. Pull the textured mesh (not *_shape.glb) straight from outdir."""
    import subprocess
    j = mb.req(f'/api/jobs/{jid}')
    outdir = j.get('outdir')
    if not outdir:
        return False
    ls = subprocess.run(['ssh', '-o', 'ConnectTimeout=8', '-o', 'BatchMode=yes',
                         PRIMARY, f'ls {outdir}'], capture_output=True, text=True)
    names = [n for n in ls.stdout.split()
             if n.endswith(suffix) and not n.endswith('_shape.glb')]
    if not names:
        return False
    r = subprocess.run(['scp', '-o', 'ConnectTimeout=8', '-o', 'BatchMode=yes',
                        f'{PRIMARY}:{outdir}/{names[0]}', out], capture_output=True)
    return r.returncode == 0 and os.path.exists(out)


def fetch_output(jid, out, suffix=None):
    outs = mb.outputs_for(jid, suffix)
    if outs:
        n = mb.fetch(outs[0]['id'], out)
        print(f'[pipe] wrote {out} ({n//1024}KB)', flush=True)
        return
    if suffix and scp_from_outdir(jid, out, suffix):
        print(f'[pipe] wrote {out} ({os.path.getsize(out)//1024}KB, outdir fallback)', flush=True)
        return
    raise RuntimeError(f'no output asset for {jid} (api + outdir fallback both empty)')


def main(spec_path, outdir):
    specs = json.load(open(spec_path))
    srcdir = os.path.join(outdir, 'src')
    os.makedirs(srcdir, exist_ok=True)
    failed = []
    for s in specs:
        name = s['name']
        glb = os.path.join(outdir, f'{name}.glb')
        if os.path.exists(glb):
            print(f'[pipe] {name}: glb exists, skip', flush=True)
            continue
        try:
            img = os.path.join(srcdir, f'{name}.png')
            if not os.path.exists(img):
                jid = run_job('flux_local', {'prompt': s['prompt']}, label=name)
                fetch_output(jid, img)
            cut = os.path.join(srcdir, f'{name}-cut.png')
            if not os.path.exists(cut):
                aid = mb.upload(img)
                jid = run_job('bg_remove_local', {}, asset_id=aid, label=name)
                fetch_output(jid, cut)
            aid = mb.upload(cut)
            jid = run_job('hunyuan3d_mlx', {}, asset_id=aid, label=name)
            fetch_output(jid, glb, suffix='.glb')
        except Exception as e:
            print(f'[pipe] {name}: FAILED — {e}', flush=True)
            failed.append(name)
    ok = len(specs) - len(failed)
    print(f'[pipe] mesh wave complete: {ok}/{len(specs)} ok'
          + (f' — failed: {failed}' if failed else ''), flush=True)


if __name__ == '__main__':
    main(sys.argv[1], sys.argv[2])