#!/usr/bin/env python3
"""Throwaway spike: prove reikhelm (Mac) -> ComfyUI (LAN workstation) depth-render works.

Pure stdlib. No pip installs. Three subcommands:

  depth   generate a synthetic depth map PNG (no network) -- our test input
  probe   ask a ComfyUI server what models / controlnets / nodes it has
  render  upload a depth PNG, run a depth-ControlNet graph, pull the image back

This is a validation tool, not the real integration. Once the pipe is proven we
port the proven flow into Rust (reikhelm renders depth -> POST /prompt -> /view).
"""
import argparse
import json
import os
import struct
import sys
import time
import urllib.error
import urllib.request
import zlib


# --------------------------------------------------------------------------- #
# Grayscale PNG writer (stdlib only -- avoids a Pillow dependency)
# --------------------------------------------------------------------------- #
def write_gray_png(path, width, height, pixels):
    """pixels: a bytes/bytearray of length width*height, one 8-bit gray value each."""
    def chunk(tag, data):
        body = tag + data
        return struct.pack(">I", len(data)) + body + struct.pack(">I", zlib.crc32(body) & 0xFFFFFFFF)

    ihdr = struct.pack(">IIBBBBB", width, height, 8, 0, 0, 0, 0)  # 8-bit, color type 0 (gray)
    raw = bytearray()
    for y in range(height):
        raw.append(0)  # filter byte 0 (None) per scanline
        raw.extend(pixels[y * width:(y + 1) * width])
    with open(path, "wb") as f:
        f.write(b"\x89PNG\r\n\x1a\n")
        f.write(chunk(b"IHDR", ihdr))
        f.write(chunk(b"IDAT", zlib.compress(bytes(raw), 9)))
        f.write(chunk(b"IEND", b""))


# --------------------------------------------------------------------------- #
# Synthetic depth maps. ControlNet convention: NEAR = white(255), FAR = black(0).
# --------------------------------------------------------------------------- #
def corridor_depth(w, h, bands=0, falloff=2.2, floor_bias=0.0):
    """One-point-perspective stone corridor: bright near the viewer (frame edge),
    dark at the vanishing point, SMOOTH by default so no hard ring seams appear.
    `falloff` > 1 deepens the tunnel / thins the near rim. `bands` > 0 quantizes
    into concentric rings (causes visible stone ridges -- leave 0 for smooth).
    `floor_bias` (0..1) lowers the vanishing point so the floor occupies more of
    the lower frame, reading as a hall you stand in rather than a symmetric portal."""
    px = bytearray(w * h)
    cx = w / 2.0
    cy = h * (0.5 + 0.5 * floor_bias)          # push vanishing point upward
    for y in range(h):
        for x in range(w):
            # Chebyshev distance to the (possibly shifted) center -> rectangular tunnel.
            dx = abs(x - cx) / cx
            dy = abs(y - cy) / max(cy, h - cy)
            d = max(dx, dy)
            v = d ** falloff                   # steeper => deeper tunnel, thin near rim
            if bands > 0:
                v = round(v * bands) / bands   # quantize into rings (visible ridges)
            px[y * w + x] = max(0, min(255, int(v * 255)))
    return px


def room_depth(w, h):
    """Top-down 'dollhouse' room: tall walls near the camera (bright border),
    floor farther away (mid), recessed pool in the middle (dark)."""
    px = bytearray(w * h)
    bx, by = w // 7, h // 7
    px0, px1, py0, py1 = w // 3, 2 * w // 3, h // 3, 2 * h // 3
    for y in range(h):
        for x in range(w):
            if x < bx or x >= w - bx or y < by or y >= h - by:
                v = 235                         # walls: near/bright
            elif px0 <= x < px1 and py0 <= y < py1:
                v = 40                          # recessed pool: far/dark
            else:
                v = 120                         # floor: mid
            px[y * w + x] = v
    return px


def cmd_depth(args):
    w = args.width or args.size
    h = args.height or args.size
    if args.kind == "corridor":
        px = corridor_depth(w, h, bands=args.bands, falloff=args.falloff, floor_bias=args.floor_bias)
    else:
        px = room_depth(w, h)
    write_gray_png(args.out, w, h, px)
    print(f"wrote {args.out}  ({w}x{h}, kind={args.kind}, bands={args.bands}, falloff={args.falloff})")


# --------------------------------------------------------------------------- #
# ComfyUI HTTP helpers
# --------------------------------------------------------------------------- #
def _get(base, path, timeout=15):
    with urllib.request.urlopen(base + path, timeout=timeout) as r:
        return r.read()


def _get_json(base, path, timeout=15):
    return json.loads(_get(base, path, timeout))


def _node_input_options(info, class_type, input_name):
    """Pull the dropdown option list (e.g. available checkpoints) for one node input."""
    node = info.get(class_type)
    if not node:
        return None
    spec = node.get("input", {})
    for group in ("required", "optional"):
        if input_name in spec.get(group, {}):
            opt = spec[group][input_name]
            # ComfyUI encodes a dropdown as [ [list_of_values], {meta} ] or [list_of_values]
            if isinstance(opt, list) and opt and isinstance(opt[0], list):
                return opt[0]
    return None


def cmd_probe(args):
    base = args.url.rstrip("/")
    try:
        stats = _get_json(base, "/system_stats")
    except urllib.error.URLError as e:
        print(f"!! cannot reach {base} -- {e}")
        print("   ComfyUI must be started with --listen 0.0.0.0 to accept LAN connections,")
        print("   and the firewall must allow the port (default 8188).")
        sys.exit(2)

    print(f"== reachable: {base} ==")
    sysinfo = stats.get("system", {})
    print(f"   comfyui: {sysinfo.get('comfyui_version', '?')}  python: {sysinfo.get('python_version', '?')[:7]}")
    for d in stats.get("devices", []):
        free = d.get("vram_free", 0) / 1e9
        total = d.get("vram_total", 0) / 1e9
        print(f"   device: {d.get('name', '?')}  VRAM {free:.1f}/{total:.1f} GB free")

    info = _get_json(base, "/object_info", timeout=60)
    print(f"\n== installed node classes: {len(info)} ==")

    def show(label, class_type, input_name, filt=None):
        opts = _node_input_options(info, class_type, input_name)
        if opts is None:
            print(f"   [{label}] node '{class_type}' NOT present")
            return
        if filt:
            hits = [o for o in opts if any(f in o.lower() for f in filt)]
            extra = f"  (filtered {len(hits)}/{len(opts)})" if hits != opts else ""
            opts = hits or opts
            print(f"   [{label}]{extra}: {opts}")
        else:
            print(f"   [{label}] ({len(opts)}): {opts}")

    show("checkpoints", "CheckpointLoaderSimple", "ckpt_name")
    show("controlnets", "ControlNetLoader", "control_net_name")
    show("unet/diffusion (flux/z-image)", "UNETLoader", "unet_name")
    show("vae", "VAELoader", "vae_name")
    show("clip", "CLIPLoader", "clip_name")
    show("dualclip (flux)", "DualCLIPLoader", "clip_name1")

    print("\n== relevant nodes present? ==")
    for n in ["ControlNetApplyAdvanced", "ControlNetApply", "LoadImage", "KSampler",
              "UNETLoader", "DualCLIPLoader", "FluxGuidance", "VAELoader"]:
        print(f"   {'yes' if n in info else ' no'}  {n}")
    # surface any Z-Image-specific nodes by name
    zi = sorted(k for k in info if "image" in k.lower() and "z" in k.lower().split("image")[0][-2:])
    flux = sorted(k for k in info if "flux" in k.lower())
    if flux:
        print(f"   flux nodes: {flux}")
    print("\nNext: tell me the checkpoint + controlnet names above and I'll wire the render graph.")


# --------------------------------------------------------------------------- #
# Render: upload depth, run an SDXL depth-ControlNet graph, pull the result.
# (SDXL first -- simplest, most-likely-installed. Flux/Z-Image added after probe.)
# --------------------------------------------------------------------------- #
def upload_image(base, path):
    boundary = "----comfyspikeboundary"
    fn = os.path.basename(path)
    with open(path, "rb") as f:
        data = f.read()
    parts = [
        f"--{boundary}\r\n".encode(),
        f'Content-Disposition: form-data; name="image"; filename="{fn}"\r\n'.encode(),
        b"Content-Type: image/png\r\n\r\n", data, b"\r\n",
        f"--{boundary}\r\n".encode(),
        b'Content-Disposition: form-data; name="overwrite"\r\n\r\ntrue\r\n',
        f"--{boundary}--\r\n".encode(),
    ]
    body = b"".join(parts)
    req = urllib.request.Request(
        base + "/upload/image", data=body, method="POST",
        headers={"Content-Type": f"multipart/form-data; boundary={boundary}"})
    return json.loads(urllib.request.urlopen(req, timeout=30).read())


def sdxl_depth_graph(a, depth_filename):
    return {
        "ckpt": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": a.ckpt}},
        "pos": {"class_type": "CLIPTextEncode", "inputs": {"text": a.prompt, "clip": ["ckpt", 1]}},
        "neg": {"class_type": "CLIPTextEncode", "inputs": {"text": a.negative, "clip": ["ckpt", 1]}},
        "img": {"class_type": "LoadImage", "inputs": {"image": depth_filename}},
        "cnet": {"class_type": "ControlNetLoader", "inputs": {"control_net_name": a.controlnet}},
        "apply": {"class_type": "ControlNetApplyAdvanced", "inputs": {
            "positive": ["pos", 0], "negative": ["neg", 0], "control_net": ["cnet", 0],
            "image": ["img", 0], "strength": a.strength, "start_percent": 0.0, "end_percent": 1.0}},
        "latent": {"class_type": "EmptyLatentImage", "inputs": {"width": a.size, "height": a.size, "batch_size": 1}},
        "ks": {"class_type": "KSampler", "inputs": {
            "seed": a.seed, "steps": a.steps, "cfg": a.cfg, "sampler_name": a.sampler,
            "scheduler": a.scheduler, "denoise": 1.0, "model": ["ckpt", 0],
            "positive": ["apply", 0], "negative": ["apply", 1], "latent_image": ["latent", 0]}},
        "vae": {"class_type": "VAEDecode", "inputs": {"samples": ["ks", 0], "vae": ["ckpt", 2]}},
        "save": {"class_type": "SaveImage", "inputs": {"images": ["vae", 0], "filename_prefix": "comfyspike"}},
    }


def _run_graph(base, graph, out_dir, wait):
    """Queue a graph, poll history, download any output images. Returns saved paths."""
    try:
        resp = urllib.request.urlopen(urllib.request.Request(
            base + "/prompt", data=json.dumps({"prompt": graph}).encode(),
            headers={"Content-Type": "application/json"}, method="POST"), timeout=30)
    except urllib.error.HTTPError as e:
        print("!! /prompt rejected the graph (validation error):")
        print(e.read().decode()[:3000])
        sys.exit(1)
    pid = json.loads(resp.read())["prompt_id"]
    print(f"queued prompt {pid} ... waiting (up to {wait}s)")

    deadline = time.time() + wait
    while time.time() < deadline:
        hist = _get_json(base, "/history/" + pid)
        if pid in hist:
            entry = hist[pid]
            status = entry.get("status", {})
            if status.get("status_str") == "error":
                print("!! execution error:")
                print(json.dumps(status.get("messages", status), indent=2)[:3000])
            saved = []
            for node_out in entry.get("outputs", {}).values():
                for im in node_out.get("images", []):
                    q = urllib.parse.urlencode({"filename": im["filename"],
                                                "subfolder": im.get("subfolder", ""),
                                                "type": im.get("type", "output")})
                    blob = _get(base, "/view?" + q, timeout=120)
                    out = os.path.join(out_dir or ".", "out_" + im["filename"])
                    with open(out, "wb") as f:
                        f.write(blob)
                    saved.append(out)
            print("rendered:", *saved) if saved else print("finished but no images -- check the graph")
            return saved
        time.sleep(2)
    print(f"!! timed out after {wait}s")
    return []


def cmd_render(args):
    base = args.url.rstrip("/")
    up = upload_image(base, args.depth)
    depth_fn = up["name"] + (f" [{up['subfolder']}]" if up.get("subfolder") else "")
    print(f"uploaded depth -> {depth_fn}")
    _run_graph(base, sdxl_depth_graph(args, depth_fn), os.path.dirname(args.depth), args.wait)


# --------------------------------------------------------------------------- #
# Z-Image-Turbo + DiffSynth Union depth patch (the recommended path).
# Uses model_patches (NOT controlnet/) + the native ZImageFunControlnet node.
# --------------------------------------------------------------------------- #
def zimage_depth_graph(a, depth_filename):
    g = {
        "unet": {"class_type": "UNETLoader", "inputs": {"unet_name": a.unet, "weight_dtype": "default"}},
        "clip": {"class_type": "CLIPLoader", "inputs": {"clip_name": a.clip, "type": a.clip_type}},
        "vae": {"class_type": "VAELoader", "inputs": {"vae_name": a.vae}},
        "patch": {"class_type": "ModelPatchLoader", "inputs": {"name": a.patch}},
        "depth": {"class_type": "LoadImage", "inputs": {"image": depth_filename}},
        "shift": {"class_type": "ModelSamplingAuraFlow", "inputs": {"model": ["unet", 0], "shift": a.shift}},
        "cnet": {"class_type": "ZImageFunControlnet", "inputs": {
            "model": ["shift", 0], "model_patch": ["patch", 0], "vae": ["vae", 0],
            "strength": a.strength, "image": ["depth", 0]}},
        "pos": {"class_type": "CLIPTextEncode", "inputs": {"text": a.prompt, "clip": ["clip", 0]}},
        "neg": {"class_type": "CLIPTextEncode", "inputs": {"text": a.negative, "clip": ["clip", 0]}},
        "latent": {"class_type": "EmptySD3LatentImage", "inputs": {"width": a.width or a.size, "height": a.height or a.size, "batch_size": 1}},
        "ks": {"class_type": "KSampler", "inputs": {
            "seed": a.seed, "steps": a.steps, "cfg": a.cfg, "sampler_name": a.sampler,
            "scheduler": a.scheduler, "denoise": 1.0, "model": ["cnet", 0],
            "positive": ["pos", 0], "negative": ["neg", 0], "latent_image": ["latent", 0]}},
        "dec": {"class_type": "VAEDecode", "inputs": {"samples": ["ks", 0], "vae": ["vae", 0]}},
        "save": {"class_type": "SaveImage", "inputs": {"images": ["dec", 0], "filename_prefix": a.prefix}},
    }
    return g


def cmd_zrender(args):
    base = args.url.rstrip("/")
    up = upload_image(base, args.depth)
    depth_fn = up["name"] + (f" [{up['subfolder']}]" if up.get("subfolder") else "")
    print(f"uploaded depth -> {depth_fn}")
    _run_graph(base, zimage_depth_graph(args, depth_fn), os.path.dirname(args.depth), args.wait)


# --------------------------------------------------------------------------- #
import urllib.parse  # noqa: E402  (used in cmd_render)


def cmd_nodeinfo(args):
    """Dump the exact input/output signature of specific node classes from a live server."""
    base = args.url.rstrip("/")
    info = _get_json(base, "/object_info", timeout=60)
    for n in args.nodes:
        node = info.get(n)
        if not node:
            print(f"\n## {n}  -- NOT PRESENT")
            continue
        print(f"\n## {n}")
        spec = node.get("input", {})
        for group in ("required", "optional"):
            for name, val in spec.get(group, {}).items():
                meta = {}
                t = val
                if isinstance(val, list) and val:
                    t = val[0]
                    if len(val) > 1 and isinstance(val[1], dict):
                        meta = val[1]
                if isinstance(t, list):  # combo / dropdown
                    shown = t if len(t) <= 20 else t[:20] + [f"...(+{len(t) - 20})"]
                    t = f"COMBO{shown}"
                hint = {k: meta[k] for k in ("default",) if k in meta}
                print(f"  {group:8} {name}: {t}" + (f"  {hint}" if hint else ""))
        outs = node.get("output_name") or node.get("output")
        print(f"  -> outputs: {outs}")


def main():
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = p.add_subparsers(dest="cmd", required=True)

    d = sub.add_parser("depth", help="generate a synthetic depth map PNG")
    d.add_argument("--kind", choices=["corridor", "room"], default="corridor")
    d.add_argument("--size", type=int, default=1024, help="square size if --width/--height omitted")
    d.add_argument("--width", type=int, default=0)
    d.add_argument("--height", type=int, default=0)
    d.add_argument("--bands", type=int, default=0, help="corridor: >0 quantizes into rings (visible ridges); 0=smooth")
    d.add_argument("--falloff", type=float, default=2.2, help="corridor: >1 deepens the tunnel / thins the near rim")
    d.add_argument("--floor-bias", dest="floor_bias", type=float, default=0.0,
                   help="corridor: 0..1 lowers the vanishing point (more floor in lower frame)")
    d.add_argument("--out", default="sample_depth.png")
    d.set_defaults(func=cmd_depth)

    pr = sub.add_parser("probe", help="inspect a ComfyUI server's models/nodes")
    pr.add_argument("url", help="e.g. http://192.168.1.50:8188")
    pr.set_defaults(func=cmd_probe)

    r = sub.add_parser("render", help="upload depth + run an SDXL depth-ControlNet render")
    r.add_argument("url")
    r.add_argument("--depth", default="sample_depth.png")
    r.add_argument("--ckpt", required=True, help="checkpoint filename from `probe`")
    r.add_argument("--controlnet", required=True, help="depth controlnet filename from `probe`")
    r.add_argument("--prompt", default="inside an ancient stone dungeon corridor, wet flagstone, "
                                       "flickering torchlight, volumetric fog, moss, dramatic shadows, "
                                       "dark fantasy, highly detailed, cinematic")
    r.add_argument("--negative", default="blurry, text, watermark, modern, people, cartoon, flat lighting")
    r.add_argument("--strength", type=float, default=0.6)
    r.add_argument("--size", type=int, default=1024)
    r.add_argument("--steps", type=int, default=28)
    r.add_argument("--cfg", type=float, default=6.0)
    r.add_argument("--seed", type=int, default=42)
    r.add_argument("--sampler", default="dpmpp_2m")
    r.add_argument("--scheduler", default="karras")
    r.add_argument("--wait", type=int, default=300)
    r.set_defaults(func=cmd_render)

    ni = sub.add_parser("nodeinfo", help="dump exact input/output signatures of node classes")
    ni.add_argument("url")
    ni.add_argument("nodes", nargs="+")
    ni.set_defaults(func=cmd_nodeinfo)

    z = sub.add_parser("zrender", help="Z-Image-Turbo depth render via the model_patches union patch")
    z.add_argument("url")
    z.add_argument("--depth", default="sample_depth.png")
    z.add_argument("--unet", default="z_image_turbo_bf16.safetensors")
    z.add_argument("--clip", default="qwen_3_4b.safetensors")
    z.add_argument("--clip-type", dest="clip_type", default="qwen_image",
                   help="CLIPLoader type for the Qwen3-4B encoder")
    z.add_argument("--vae", default="ae.safetensors")
    z.add_argument("--patch", default="Z-Image-Turbo-Fun-Controlnet-Union-2.1-lite-2602-8steps.safetensors")
    z.add_argument("--prompt", default="inside an ancient stone dungeon corridor, wet mossy flagstone walls, "
                                       "flickering torchlight, volumetric fog, deep shadows receding into "
                                       "darkness, dark fantasy, cinematic, highly detailed")
    z.add_argument("--negative", default="")
    z.add_argument("--strength", type=float, default=0.65)
    z.add_argument("--size", type=int, default=1024, help="square size if --width/--height omitted")
    z.add_argument("--width", type=int, default=0)
    z.add_argument("--height", type=int, default=0)
    z.add_argument("--steps", type=int, default=8)
    z.add_argument("--cfg", type=float, default=1.0)
    z.add_argument("--shift", type=float, default=3.0)
    z.add_argument("--seed", type=int, default=42)
    z.add_argument("--prefix", default="zspike", help="SaveImage filename prefix (names the output)")
    z.add_argument("--sampler", default="res_multistep")
    z.add_argument("--scheduler", default="simple")
    z.add_argument("--wait", type=int, default=300)
    z.set_defaults(func=cmd_zrender)

    args = p.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()