#!/usr/bin/env python3
"""Cashies photo -> game prop: bg_remove_local then trellis_mac on MODELBEAST.

Usage: cc_to_glb.py <name=image> [name=image ...] [--out DIR] [--workers N]

For each input image: upload -> bg_remove_local (clean cutout; the op's own
docs call this the big quality jump before image->3D) -> trellis_mac
(1024 tier, 1024 texture: prop-grade, keeps GLBs a few MB) -> download the
GLB to <out>/<name>.glb. Progress lines are grep-friendly: "DONE name" /
"FAIL name reason". Queue + token per ~/.claude/skills/fleet/SKILL.md; the
guest token allows 4 active jobs, so keep --workers at 2 (each worker holds
one job at a time, bg+trellis are sequential per item).

Job log JSON carries raw control chars -- scrub before json.loads.
"""
import json
import mimetypes
import os
import re
import sys
import threading
import time
import urllib.request
import uuid

HOST = os.environ.get("MB_HOST", "http://100.89.131.57:8777")
ENVS = ["~/Documents/fluxgod-work/.env", "~/Documents/backnforth/.env"]


def token():
    for p in ENVS:
        p = os.path.expanduser(p)
        if not os.path.exists(p):
            continue
        for line in open(p):
            if line.startswith("MB_TOKEN"):
                return line.split("=", 1)[1].strip().strip("'\"")
    sys.exit("no MB_TOKEN found on disk")


TOK = token()


def call(path, data=None, headers=None, raw=False, tries=4):
    # the queue occasionally throws a transient 401/5xx mid-batch (seen live:
    # one 401 between a working bg job and a working trellis submit) -- retry
    # with backoff before declaring an item dead
    h = {"Authorization": "Bearer " + TOK}
    h.update(headers or {})
    last = None
    for attempt in range(tries):
        try:
            req = urllib.request.Request(HOST + path, data=data, headers=h)
            blob = urllib.request.urlopen(req, timeout=180).read()
            if raw:
                return blob
            return json.loads(re.sub(r"[\x00-\x1f]", "", blob.decode("utf-8", "replace")))
        except Exception as e:
            last = e
            time.sleep(4 * (attempt + 1))
    raise last


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


def run_job(operator, asset_id, params, timeout_s):
    j = call("/api/jobs", data=json.dumps(
        {"operator": operator, "asset_id": asset_id, "params": params}).encode(),
        headers={"Content-Type": "application/json"})
    jid = j["id"]
    t0 = time.time()
    while time.time() - t0 < timeout_s:
        time.sleep(10)
        st = call("/api/jobs/%s" % jid)
        if st.get("status") in ("done", "error", "cancelled"):
            if st["status"] != "done":
                raise RuntimeError("%s job %s: %s (%s)" % (
                    operator, jid, st["status"], str(st.get("error", ""))[:160]))
            return jid
    raise RuntimeError("%s job %s timed out after %ds" % (operator, jid, timeout_s))


def job_asset(jid, suffix=None):
    # jobs don't back-link outputs: list assets, match parent_job (mb_recon contract)
    a = call("/api/assets?limit=100")
    items = a if isinstance(a, list) else a.get("items", [])
    mine = [x for x in items if x.get("parent_job") == jid]
    if suffix:
        mine = [x for x in mine
                if str(x.get("filename", x.get("name", ""))).endswith(suffix)]
    if not mine:
        raise RuntimeError("job %s produced no %s asset" % (jid, suffix or ""))
    return mine[0]["id"]


MESH_OP = "trellis_mac"
MESH_PARAMS = {
    # jerry cans made trellis_mac exit 2 on two different photos; hunyuan3d_mlx
    # chewed the same cutout happily -- keep both dialled in
    "trellis_mac": {"pipeline_type": "1024", "texture_size": 1024},
    "hunyuan3d_mlx": {"texture_size": 1024},
}


def one(name, img, outdir):
    try:
        aid = upload(img)
        print("UPLOADED %s -> %s" % (name, aid), flush=True)
        bj = run_job("bg_remove_local", aid, {"resolution": 1024}, 600)
        cut = job_asset(bj)
        print("CUTOUT %s" % name, flush=True)
        tj = run_job(MESH_OP, cut, MESH_PARAMS.get(MESH_OP, {}), 3600)
        gid = job_asset(tj, ".glb")
        out = os.path.join(outdir, name + ".glb")
        open(out, "wb").write(call("/api/assets/%s/file" % gid, raw=True))
        print("DONE %s (%dKB)" % (name, os.path.getsize(out) // 1024), flush=True)
    except Exception as e:
        print("FAIL %s %s" % (name, str(e)[:200]), flush=True)


def main():
    args = sys.argv[1:]
    outdir = "assets/tools"
    workers = 2
    if "--out" in args:
        i = args.index("--out")
        outdir = args[i + 1]
        del args[i:i + 2]
    if "--workers" in args:
        i = args.index("--workers")
        workers = int(args[i + 1])
        del args[i:i + 2]
    if "--operator" in args:
        i = args.index("--operator")
        global MESH_OP
        MESH_OP = args[i + 1]
        del args[i:i + 2]
    os.makedirs(outdir, exist_ok=True)
    queue = [a.split("=", 1) for a in args if "=" in a]
    lock = threading.Lock()

    def worker():
        while True:
            with lock:
                if not queue:
                    return
                name, img = queue.pop(0)
            one(name, img, outdir)

    threads = [threading.Thread(target=worker) for _ in range(workers)]
    for t in threads:
        t.start()
    for t in threads:
        t.join()
    print("BATCH_COMPLETE", flush=True)


if __name__ == "__main__":
    main()