#!/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.70] [--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"]
# Closed-facing bodies: SAME theme flavor (material, contents, mood) but every
# architectural OPENING noun stripped — gatehouse/portcullis/entrance, wall
# niches, the throne's recessed dais, etc. A GPU sweep (2026-06-15) proved the
# phantom doors on closed walls were summoned by these words, not the depth map:
# `threshold`'s "gatehouse, portcullis, entrance" painted a gate onto a solid
# wall, and at cfg 1.0 the negative branch is too weak to cancel it. Working WITH
# a generic solid-wall description (no opening nouns) renders a clean dead-end;
# even "bricked-up archway" backfires because "archway" alone re-summons the arch.
THEME_BODY_CLOSED = {
    "forge":     "forge chamber, soot-blackened granite walls, iron anvils and hanging chains, dim fiery glow",
    "cistern":   "flooded cistern, dark still water, dripping wet stone walls, cold teal light",
    "crypt":     "ancient crypt, carved stone sarcophagi along blank stone walls, cobwebs, cold pale light",
    "library":   "ruined library, tall stone shelves of rotting books against the wall, drifting dust, warm dim light",
    "throne":    "throne hall, towering carved stone columns and a blank stone wall, tattered banners, cold grandeur",
    "vault":     "treasure vault, iron-bound chests and stone strongboxes against heavy locked stone walls, scattered gold",
    "hall":      "grand pillared hall, carved stone columns and a high blank stone wall, vaulted ceiling",
    "den":       "foul beast den, gnawed bones and filth, claw-scratched blank stone walls, dim red light",
    "threshold": "heavy sealed gateroom of dungeon stone, worn flagstones, blank ashlar walls",
    "stone":     "rough-hewn stone chamber, bare granite walls, plain and ancient",
}
# 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 (opening-free body above) + this sealed-wall clause.
AHEAD_WALL = ("the way ahead ends at a blank unbroken wall of solid carved stone blocks, "
              "sealed flat masonry, 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")
# Closed facings additionally negate every opening noun (cheap insurance; the
# opening-free positive body does the real work). NOT applied to open facings —
# they WANT a passage.
NEG_OPENINGS = "doorway, door, archway, arch, opening, passage, tunnel, gate, portcullis, gateway, entrance, corridor"


def prompt_for(theme, is_open):
    """Return (positive, negative) for this (theme, facing)."""
    if is_open:
        body = THEME_BODY.get(theme or "stone", DEFAULT_BODY)
        return f"{STYLE}{body}, {AHEAD_OPEN}", NEGATIVE
    body = THEME_BODY_CLOSED.get(theme or "stone", THEME_BODY_CLOSED["stone"])
    return f"{STYLE}{body}, {AHEAD_WALL}", f"{NEGATIVE}, {NEG_OPENINGS}"

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.70, help="depth control strength. 0.70 is the sweet spot now "
                    "that rev-3 maps are full-contrast: a GPU A/B showed 0.85 over-constrained the bright near-ceiling "
                    "into a blown-white band (worst on crypt frames, ~28%% of pixels), while 0.70 renders a natural "
                    "dark vaulted ceiling AND still honors geometry — closed walls stay solid (no phantom doors return).")
    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, negative = 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()