#!/usr/bin/env python3
"""mb.py — MODELBEAST client for the not-tonight asset pipeline.

    python3 tools/gen/mb.py image  <out.png> "<prompt>" [--seed 4300]
    python3 tools/gen/mb.py mesh   <in.png> <out.glb> [--operator trellis2_mlx]
    python3 tools/gen/mb.py ops

The farm is free and unlimited on our own silicon, which makes it the fallback
when Cloudflare's daily neurons run out — and the ONLY route for meshes.

Contract notes that cost previous sessions time, encoded here so they cannot be
got wrong again:

- `asset_id` is a **TOP-LEVEL** job field, not a member of `params`. Nested, the
  mesh operators fail with a flat "no input image" that reads like a bad upload.
- Job logs carry raw control characters, so every response is scrubbed before
  `json.loads` — otherwise polling dies mid-batch on a malformed payload.
- Jobs do not back-link their output asset id. Retrieval is "newest asset whose
  `parent_job` is my job id"; the newest-of-type shortcut RACES when two jobs
  run at once, and on 2026-07-16 that handed one project's mesh to another.
- The asset list field is `name`, not `filename`.
- The guest token allows 4 active jobs, so the batch driver keeps its own gate.
"""
from __future__ import annotations

import json
import mimetypes
import sys
import time
import urllib.error
import urllib.request
import uuid
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
import creds  # noqa: E402

HOST = "http://100.89.131.57:8777"
MAX_ACTIVE = 4          # guest-token ceiling
POLL_SECONDS = 6


def _scrub(raw: bytes) -> object:
    text = raw.decode("utf-8", "replace")
    return json.loads("".join(c if c >= " " or c == "\t" else " " for c in text))


def _req(path: str, data: bytes | None = None, headers: dict[str, str] | None = None,
         raw: bool = False, timeout: int = 180) -> object:
    head = {"Authorization": f"Bearer {creds.get('MB_TOKEN')}"}
    head.update(headers or {})
    req = urllib.request.Request(HOST + path, data=data, headers=head)
    try:
        body = urllib.request.urlopen(req, timeout=timeout).read()
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode(errors="ignore")[:400]
        raise SystemExit(f"MODELBEAST HTTP {exc.code} on {path}: {detail}") from None
    except urllib.error.URLError as exc:
        raise SystemExit(f"MODELBEAST unreachable at {HOST}: {exc.reason}") from None
    return body if raw else _scrub(body)


def ops() -> list[str]:
    data = _req("/api/operators")
    items = data if isinstance(data, list) else data.get("operators", data.get("items", []))
    return sorted(o.get("name", str(o)) if isinstance(o, dict) else str(o) for o in items)


def upload(path: Path) -> str:
    boundary = uuid.uuid4().hex
    ctype = mimetypes.guess_type(str(path))[0] or "application/octet-stream"
    body = (
        f'--{boundary}\r\nContent-Disposition: form-data; name="file"; filename="{path.name}"\r\n'
        f"Content-Type: {ctype}\r\n\r\n"
    ).encode() + path.read_bytes() + f"\r\n--{boundary}--\r\n".encode()
    res = _req("/api/assets", data=body,
               headers={"Content-Type": f"multipart/form-data; boundary={boundary}"})
    if isinstance(res, dict):
        return res.get("id") or (res.get("items") or [{}])[0].get("id", "")
    return res[0]["id"] if res else ""


def submit(operator: str, params: dict | None = None, asset_id: str | None = None) -> str:
    """Submit a job. `asset_id` goes TOP-LEVEL — see the module docstring."""
    payload: dict[str, object] = {"operator": operator, "params": params or {}}
    if asset_id:
        payload["asset_id"] = asset_id
    res = _req("/api/jobs", data=json.dumps(payload).encode(),
               headers={"Content-Type": "application/json"})
    return res["id"]


def status(job_id: str) -> str:
    return _req(f"/api/jobs/{job_id}").get("status", "unknown")


def wait(job_id: str, *, timeout: int = 2400, label: str = "") -> str:
    t0 = time.time()
    last = ""
    while True:
        state = status(job_id)
        if state != last:
            print(f"[mb] {label or job_id}: {state} ({int(time.time() - t0)}s)", flush=True)
            last = state
        if state in ("done", "error", "cancelled", "failed"):
            return state
        if time.time() - t0 > timeout:
            return "timeout"
        time.sleep(POLL_SECONDS)


def output_asset(job_id: str, suffix: str) -> str | None:
    """The asset THIS job produced. parent_job only — never newest-of-type."""
    data = _req("/api/assets?limit=200")
    items = data if isinstance(data, list) else data.get("items", [])
    mine = [
        a for a in items
        if a.get("parent_job") == job_id
        and str(a.get("name", a.get("filename", ""))).lower().endswith(suffix)
    ]
    return mine[0]["id"] if mine else None


def fetch(asset_id: str, out: Path) -> Path:
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_bytes(_req(f"/api/assets/{asset_id}/file", raw=True, timeout=600))
    return out


# ---- the two things the asset batch actually asks for ------------------------

def image(out: Path, prompt: str, *, seed: int = 0, steps: int = 4,
          width: int = 1024, height: int = 1024, label: str = "") -> Path:
    job = submit("flux_local", {
        "prompt": prompt, "seed": seed, "steps": steps,
        "width": width, "height": height, "model": "flux2-klein-4b",
    })
    state = wait(job, label=label or out.stem)
    if state != "done":
        raise RuntimeError(f"flux_local {state} for {out.stem}")
    asset = output_asset(job, ".png")
    if not asset:
        raise RuntimeError(f"flux_local produced no png for {out.stem}")
    return fetch(asset, out)


# Props ship at 16-48px. TRELLIS's defaults (1024_cascade + a 2048px metal PBR
# bake) are tuned for hero meshes you look at up close, and they cost 5-12
# minutes each on a queue that runs ONE gpu job at a time — call it four hours
# for a 37-prop tileset. The 512 tier with a vertex bake lands in ~90s, and at
# the size these are actually seen (docs/SPACES.md: "judge at ship size, not
# render size") the difference does not survive the downscale to 32x32, let
# alone the venue's own darkness sheet.
FAST_MESH = {
    "pipeline_type": "512",
    "baker": "vertex",        # 1s bake, cleanest colours, no MR maps we'd use
    "max_bake_faces": 200000,
    "alpha_mode": "opaque",   # props_import keys off Blender's alpha, not the mesh's
}


def mesh(src: Path, out: Path, *, operator: str = "trellis2_mlx", label: str = "",
         params: dict | None = None) -> Path:
    asset_id = upload(src)
    job = submit(operator, {**FAST_MESH, **(params or {})}, asset_id=asset_id)
    state = wait(job, timeout=3600, label=label or out.stem)
    if state != "done":
        raise RuntimeError(f"{operator} {state} for {out.stem}")
    glb = output_asset(job, ".glb")
    if not glb:
        raise RuntimeError(f"{operator} produced no glb for {out.stem}")
    return fetch(glb, out)


def main() -> None:
    if len(sys.argv) < 2:
        sys.exit(__doc__)
    cmd = sys.argv[1]
    if cmd == "ops":
        print("\n".join(ops()))
    elif cmd == "image":
        seed = int(sys.argv[sys.argv.index("--seed") + 1]) if "--seed" in sys.argv else 0
        print(image(Path(sys.argv[2]), sys.argv[3], seed=seed))
    elif cmd == "mesh":
        op = sys.argv[sys.argv.index("--operator") + 1] if "--operator" in sys.argv else "trellis2_mlx"
        print(mesh(Path(sys.argv[2]), Path(sys.argv[3]), operator=op))
    else:
        sys.exit(__doc__)


if __name__ == "__main__":
    main()