#!/usr/bin/env python3
"""Bake a first-person dungeon atlas: every room's 4 cardinal views → pixel art.

Input  (written by the browser, see explore-bake hook):
  tools/out/atlas/<seed>/_work/r<id>_<dir>_depth.png   FPV depth maps (1280x720)
  tools/out/atlas/<seed>/manifest.json                 room graph (nav + vantages + tiles)

For each depth map this runs the proven chain — comfy.py zrender (Z-Image depth
ControlNet) → pixelate.py (locked Primordyn dither) — and writes the explorer frame:
  tools/out/atlas/<seed>/r<id>_<dir>.png               640x360 indexed pixel art

The prompt is tank's "empty abandoned ruin" recipe: NO first-person/POV/crawler
trigger words (they summon player hands), a fixed style-anchor block, and a FIXED
seed across every frame so the palette + lighting stay coherent across the atlas.

  python3 tools/bake_atlas.py --seed 7 [--url http://192.168.1.26:8188]
                              [--strength 0.65] [--size 640x360] [--limit N]
"""
import argparse
import glob
import json
import os
import re
import subprocess
import sys

REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
COMFY = os.path.join(REPO, "tools", "comfy-spike", "comfy.py")
PIXELATE = os.path.join(REPO, "tools", "pixelart", "pixelate.py")
VENV_PY = os.path.join(REPO, "tools", "pixelart", ".venv", "bin", "python")
PALETTE = os.path.join(REPO, "tools", "pixelart", "primordyn_v2.json")
RUN = os.getpid()  # unique per invocation → ComfyUI never reuses a SaveImage prefix (re-bakes stay clean)

# Tank's anti-hands framing ("empty abandoned ruin, no people, architectural
# photography" — never "first person / POV / dungeon crawler"), but the room's
# CHARACTER now comes from its procgen theme and the view's actual openings, so
# the atlas stops being a mashup of one generic hall. The style anchor stays
# byte-identical across frames (consistency lever); the theme body + opening
# clause vary per (room, facing).
STYLE = ("dark fantasy, ancient torchlit stone, warm orange torchlight, deep black shadows, "
         "wet stone, volumetric haze, architectural photography of an empty ruin, "
         "empty, deserted, no people, uninhabited, atmospheric, highly detailed, ")
THEME_BODY = {
    "forge":     "molten forge chamber, channels of glowing orange lava, iron anvils and hanging chains, soot-blackened granite, fiery glow",
    "cistern":   "flooded stone cistern, still dark water with mirror reflections, dripping wet walls, cold teal light",
    "crypt":     "ancient crypt, carved stone sarcophagi and bone-filled wall niches, cobwebs, cold pale light",
    "library":   "ruined library, tall stone shelves of rotting books and scrolls, drifting dust, warm dim light",
    "throne":    "vast throne hall, a raised dais with a great stone throne, tattered banners, towering carved columns",
    "vault":     "treasure vault, iron-bound chests and stone strongboxes, scattered gold, heavy locked stone",
    "hall":      "grand pillared hall, rows of carved stone columns, high vaulted ceiling",
    "den":       "foul beast den, gnawed bones and filth, claw-scratched walls, a feral lair, dim red light",
    "threshold": "dungeon gatehouse, a great iron portcullis and worn stone steps, torchlit entrance",
    "stone":     "rough-hewn stone chamber, bare granite walls, plain and ancient",
}
DEFAULT_BODY = THEME_BODY["stone"]
# These reinforce the depth (a deep recess vs a near wall straight ahead) so the
# model paints a passage only where one actually exists — killing phantom doors.
AHEAD_OPEN = "a dark arched passage leads onward into shadow ahead"
# cfg 1.0 makes negatives inert, so the closed case is carried by emphatic
# positive description — themes like the vault say "locked", which otherwise
# summons phantom doors on blank walls.
AHEAD_WALL = ("the way ahead ends at a blank unbroken wall of carved stone blocks, "
              "sealed solid masonry with no door and no opening, a dead end")
NEGATIVE = ("person, people, human, figure, character, hands, fingers, hand, arm, arms, "
            "holding weapon, sword, feet, legs, body, silhouette, UI, HUD, text, watermark, modern")


def prompt_for(theme, is_open):
    body = THEME_BODY.get(theme or "stone", DEFAULT_BODY)
    return f"{STYLE}{body}, {AHEAD_OPEN if is_open else AHEAD_WALL}"

NAME_RE = re.compile(r"^r(\d+)_([NESW])_depth\.png$")


def run(cmd):
    p = subprocess.run(cmd, capture_output=True, text=True)
    if p.returncode != 0:
        sys.stderr.write(p.stdout + p.stderr + "\n")
    return p.returncode == 0


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--seed", type=int, required=True, help="dungeon seed (selects the atlas dir)")
    ap.add_argument("--url", default="http://192.168.1.26:8188")
    ap.add_argument("--strength", type=float, default=0.85, help="depth control strength (0.85 honors the room footprint)")
    ap.add_argument("--render-seed", type=int, default=7, help="FIXED diffusion seed for every frame (style lock)")
    ap.add_argument("--size", default="640x360", help="pixel-art frame size (16:9; 320x180 for chunky)")
    ap.add_argument("--width", type=int, default=1280)
    ap.add_argument("--height", type=int, default=720)
    ap.add_argument("--rooms", default="", help="comma-separated room ids to bake (default all)")
    ap.add_argument("--limit", type=int, default=0, help="bake only the first N frames (smoke test)")
    a = ap.parse_args()

    atlas = os.path.join(REPO, "tools", "out", "atlas", str(a.seed))
    work = os.path.join(atlas, "_work")

    # Per-room theme + per-facing openness come from the manifest the browser wrote.
    rooms = {}
    mpath = os.path.join(atlas, "manifest.json")
    if os.path.exists(mpath):
        with open(mpath) as f:
            for r in json.load(f).get("rooms", []):
                rooms[int(r["id"])] = r

    only = {int(x) for x in a.rooms.split(",") if x.strip()} if a.rooms else None

    def order(p):
        m = NAME_RE.match(os.path.basename(p))
        return (int(m.group(1)), "NESW".index(m.group(2))) if m else (1 << 30, 0)
    depths = sorted(glob.glob(os.path.join(work, "r*_*_depth.png")), key=order)
    if only is not None:
        depths = [p for p in depths if NAME_RE.match(os.path.basename(p)) and int(NAME_RE.match(os.path.basename(p)).group(1)) in only]
    if not depths:
        sys.exit(f"!! no depth maps to bake in {work} — run the browser export first")
    if a.limit:
        depths = depths[:a.limit]

    print(f"baking {len(depths)} frames → {os.path.relpath(atlas, REPO)}  "
          f"(strength {a.strength}, render-seed {a.render_seed}, {a.size}, per-theme prompts)")
    ok = 0
    for i, depth in enumerate(depths, 1):
        m = NAME_RE.match(os.path.basename(depth))
        if not m:
            continue
        rid, d = m.group(1), m.group(2)
        room = rooms.get(int(rid), {})
        theme = room.get("theme")
        is_open = bool(room.get("open", {}).get(d, False))
        prompt = prompt_for(theme, is_open)
        strength = a.strength
        # Walls render under their own FIXED seed: the open-frame seed's
        # favorite composition is a centered doorway — exactly what a closed
        # wall must not grow. One class, one seed → style stays coherent.
        render_seed = a.render_seed if is_open else a.render_seed + 13
        prefix = f"r{rid}_{d}_{RUN}"  # per-run unique; avoids ComfyUI prefix-counter collisions on re-bake
        frame = os.path.join(atlas, f"r{rid}_{d}.png")
        tag = f"[{i}/{len(depths)}] r{rid} {d} {theme or '?'}{'·open' if is_open else '·wall'}"

        if not run(["python3", COMFY, "zrender", a.url, "--depth", depth,
                    "--width", str(a.width), "--height", str(a.height),
                    "--strength", str(strength), "--seed", str(render_seed),
                    "--prompt", prompt, "--negative", NEGATIVE,
                    "--prefix", prefix, "--wait", "300"]):
            print(f"  {tag}  RENDER FAILED"); continue

        renders = glob.glob(os.path.join(work, f"out_{prefix}_*.png"))
        if not renders:
            print(f"  {tag}  no render output"); continue
        render = renders[0]

        if not run([VENV_PY, PIXELATE, render, "--palette", PALETTE,
                    "--size", a.size, "--preview-scale", "0", "--overshoot", "6",
                    "--out", frame]):
            print(f"  {tag}  DITHER FAILED"); continue

        os.remove(render)  # keep the depth, drop the big intermediate render
        ok += 1
        print(f"  {tag}  → {os.path.relpath(frame, REPO)}")

    print(f"done: {ok}/{len(depths)} frames baked into {os.path.relpath(atlas, REPO)}")


if __name__ == "__main__":
    main()