#!/usr/bin/env python3
"""
Build the PixelLab part library (64x64 sprites, registry: sprite_parts_pl.py).

Each base creature has a spec, pixellab_library/<body>.json. The base is generated once with
PixelLab's pixen model; every part is an edit of that exact image ("add horns ...") made with
edit-image-pixen, and the part is the pixel difference between the edit and the base. All art
is drawn in a stand-in palette (teal body = main, crimson = accent, gold = type) that maps to
the palette keys, so the generator can recolor it with brand colors.

Workflow
  python pixellab_library.py generate pl_fox --base-only  # base only (pixellab_src/pl_fox/base.png): review it first
  python pixellab_library.py generate pl_fox            # every missing edit
  python pixellab_library.py extract pl_fox             # grids parts_pl_src/, PNGs parts_pl/, diffs preview/pixellab/diff/
  python pixellab_library.py preview pl_fox             # preview/pixellab/body_pl_fox.png
  python pixellab_library.py generate pl_fox --parts pl_tail_flame --force   # redo one edit (new seed!)
  python pixellab_library.py preview --cross            # every part on every body: preview/pixellab/cross_<slot>.png
  python pixellab_library.py preview --random 60        # real random picks: preview/pixellab/random.png
  python pixellab_library.py cost                       # spend so far (ledger) and live balance
  python validate_parts.py --library pixellab           # colors, markers, canvas fit

Spec (pixellab_library/<body>.json)
  {
    "name": "pl_fox",                    # unique, "pl_" prefix; also the body part name
    "plan": "quadruped",                 # quadruped biped serpentine avian aquatic blob insectoid
    "head": ["mammal_head"],             # head-family tags; head_top parts gate on them via "requires"
    "types": {"fire": 2}, "other": 1,    # type boosts; every body is open to every type (other > 0)
    "tags": [], "excludes": [], "deny": [],   # as in sprite_parts.py (excludes: tags of parts it can't take)
    "slots": ["head_top", "back", "tail"],    # feature slots this body accepts (default all three)
    "desc": "Chibi fox, 3/4 view",
    "prompt": "chibi fox creature with ...",  # body only; BASE_STYLE (colors, framing) is appended
    "size": 48,                          # pixen canvas, multiple of 4; the creature is re-placed on 64x64
    "seed": 11, "direction": "south", "detail": "medium detail",
    "anchors": {"tail": [40, 50]},       # optional, pixel coords on the 64x64 base
    "parts": [{
      "slot": "tail", "name": "pl_tail_flame",
      "edit": "add a long tail ending in a big orange and yellow flame, curling up on the right side",
      "seed": 3,                         # change it to get a different edit
      "types": {"fire": 6}, "other": 0,  # picking weight per type / for other types
      "plans": "all",                    # or a list of plans it may go on
      "requires": [], "tags": ["tailed", "fiery"], "excludes": [], "motif": "fire",
      "palette": "type",                 # main | accent | type: the ramp it mostly uses (for previews)
      "fixed": "fire",                   # optional fixed-color ramp (generate_sprite.FIXED); needs types + other 0
      "fixed_from": ["accent", "type"],  # ramps folded into the fixed ramp (default accent, type)
      "z": null,                         # draw order override (back 10, tail 20, body 30, head_top 50)
      "deny_bodies": [],                 # bodies it must never go on
      "box": [x0, y0, x1, y1], "chars": "xyz", "min_area": 8,   # optional diff filters
      "skip": false                      # keep the spec entry but leave the part out of the library
    }]
  }

Conventions every base and edit must follow so parts mix across bodies
  - Front or 3/4 view facing the viewer (direction south / south-west), feet at the bottom.
  - Tails come out on the image's RIGHT side. Wings and horns are symmetric pairs.
  - Body in teal, details crimson; generic elemental features gold (type ramp, recolored per
    type); fixed-color features in their natural colors (orange/yellow fire, pale ice, ...).
  - The base creature fills its pixen canvas; at 64x64 it should leave room for features
    (height <= ~46px, width <= ~48px).
"""
import argparse
import importlib
import json
import random
import sys
import zlib
from collections import deque
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

import numpy as np
from PIL import Image, ImageColor, ImageDraw

import build_parts as bp
import generate_sprite as gs
import sprite_parts_pl as reg
from sprite_parts import TYPES

SPEC_DIR = reg.SPEC_DIR
SRC_DIR = Path("pixellab_src")
GRID_DIR = Path("parts_pl_src")
PARTS_DIR = Path(reg.PARTS_DIR)
PREVIEW_DIR = Path("preview/pixellab")
CANVAS, GROUND_MARGIN = reg.CANVAS, reg.GROUND_MARGIN
SLOTS = reg.FEATURE_SLOTS

# No franchise names here: "pokemon-style" made pixen copy existing designs (a Bulbasaur fox).
BASE_STYLE = (", original cute monster creature design, chibi proportions, teal body with darker teal shading, "
              "small crimson details, big expressive eyes, full body, centered, retro handheld game sprite")
EDIT_SUFFIX = ", keep everything else exactly the same"

# Stand-in palette: what PixelLab draws in, mapped to the legend characters of build_parts.py
STANDIN = {
    "main": ("123", ("#9FE0D0", "#3E9C8F", "#1F5560")),        # teal
    "accent": ("abc", ("#F07A8A", "#C42E4A", "#7A1830")),      # crimson
    "type": ("xyz", ("#FFF0A0", "#F2C230", "#A87A10")),        # gold
}
OUTLINE_RGB, WHITE_RGB = ImageColor.getrgb("#10131F"), (255, 255, 255)
OUTLINE_MAX_LIGHT = 0.17       # darker than every ramp's dark shade -> outline
# On a base, "gold" this pale or this washed-out is a cream/tan belly or highlight, not a gold
# detail (beaks, irises measure saturation >= 0.6): it becomes the main ramp, same shade.
BODY_PALE_GOLD_LIGHT = 0.7
BODY_TAN_MAX_SAT = 0.45
WHITE_MIN_LIGHT, WHITE_MAX_SAT = 0.9, 0.35
GRAY_MAX_SAT = 0.12            # too gray to tell the hue -> nearest RGB
# Separating the feature from the body redrawn by the edit, on and right around the body:
NOISE_RGB = 48                 # any edit pixel this close (raw RGB) to the base pixel is the body
SAME_RAMP_RGB = 120            # same ramp as the base pixel and closer than this: shading shift, body
JITTER_RAMPS = ("main", "outline", "white")   # just outside the old silhouette: edge jitter, body
BODY_CHANGE_WARN = 0.12
MARKER_CHAR = {"pivot": "P", "head_top": "^", "back": "B", "tail": "T", "aura": "A"}
PREVIEW_PALETTES = [("blue/amber", ["#1565C0", "#FFB300"]), ("green/red", ["#2E7D32", "#E53935"])]
BG = (226, 226, 220, 255)


class SpecError(Exception):
    pass


# ---------------------------------------------------------------- specs

def load_specs(names=()):
    specs = reg.load_specs(SPEC_DIR)
    seen = {}
    for s in specs:
        for key in ("name", "plan", "prompt"):
            if key not in s:
                raise SpecError(f"{s['_path']}: missing '{key}'")
        if s["plan"] not in reg.PLANS:
            raise SpecError(f"{s['_path']}: plan must be one of {reg.PLANS}")
        for p in [s] + s.get("parts", []):
            if p["name"] in seen:
                raise SpecError(f"{s['_path']}: name {p['name']} already used in {seen[p['name']]}")
            seen[p["name"]] = s["_path"]
        for p in s.get("parts", []):
            if p.get("slot") not in SLOTS:
                raise SpecError(f"{s['_path']}: {p.get('name')}: slot must be one of {SLOTS}")
            if p.get("fixed") and p["fixed"] not in gs.FIXED:
                raise SpecError(f"{s['_path']}: {p['name']}: fixed must be one of {list(gs.FIXED)}")
            if "edit" not in p:
                raise SpecError(f"{s['_path']}: {p['name']}: missing 'edit'")
    if names:
        unknown = set(names) - {s["name"] for s in specs}
        if unknown:
            raise SpecError(f"no spec for {sorted(unknown)}")
        specs = [s for s in specs if s["name"] in names]
    return specs


def stable_seed(text):
    return zlib.crc32(text.encode()) % 100000


# ---------------------------------------------------------------- generation

def normalize_base(raw):
    """Binarize alpha, crop to the creature, put it bottom-centre on the 64x64 canvas."""
    arr = np.array(raw.convert("RGBA"))
    arr[..., 3] = np.where(arr[..., 3] >= 128, 255, 0)
    img = Image.fromarray(arr)
    bbox = img.getbbox()
    if bbox is None:
        raise SpecError("base image is empty")
    img = img.crop(bbox)
    if img.width > CANVAS - 2 or img.height > CANVAS - GROUND_MARGIN - 1:
        raise SpecError(f"base creature is {img.width}x{img.height}; must fit {CANVAS - 2}x{CANVAS - GROUND_MARGIN - 1} "
                        f"(use a smaller 'size')")
    out = Image.new("RGBA", (CANVAS, CANVAS), (0, 0, 0, 0))
    out.paste(img, ((CANVAS - img.width) // 2, CANVAS - GROUND_MARGIN - img.height))
    return out


def base_payload(spec):
    size = spec.get("size", 48)
    return {
        "description": spec["prompt"] + BASE_STYLE,
        "image_size": {"width": size, "height": size},
        "no_background": True,
        "outline": "single color black outline",
        "detail": spec.get("detail", "medium detail"),
        "view": "side",
        "direction": spec.get("direction", "south"),
        "seed": spec.get("seed", stable_seed(spec["name"])),
    }


def edit_payload(spec, part, base):
    return {
        "image": gs_b64(base),
        "description": part["edit"] + EDIT_SUFFIX,
        "width": CANVAS, "height": CANVAS,
        "seed": part.get("seed", stable_seed(part["name"])),
        "no_background": True,
    }


def gs_b64(img):
    from pixellab_api import b64
    return b64(img)


def save_meta(path, endpoint, payload, usage):
    from pixellab_api import redact
    path.write_text(json.dumps({"endpoint": endpoint, "request": redact(payload), "usage": usage}, indent=2) + "\n")


MIN_USD = 0.30   # generate stops starting new bodies below this balance once trial generations are gone


def budget_ok(client):
    bal = client.balance()
    trial = (bal.get("subscription") or {}).get("generations") or 0
    usd = (bal.get("credits") or {}).get("usd") or 0
    if trial <= 0 and usd < MIN_USD:
        print(f"budget stop: ${usd:.2f} left (floor ${MIN_USD:.2f}), no trial generations")
        return False
    return True


def generate(specs, only=(), force=False, dry=False, workers=1, base_only=False):
    client = None
    if not dry:
        from pixellab_api import Client
        client = Client()
    for spec in specs:
        if client and not budget_ok(client):
            return
        d = SRC_DIR / spec["name"]
        d.mkdir(parents=True, exist_ok=True)
        base_path = d / "base.png"
        if force and not only or not base_path.exists():
            payload = base_payload(spec)
            if dry:
                print(f"{spec['name']}: create-image-pixen {json.dumps(payload)}")
            else:
                print(f"{spec['name']}: generating base ...")
                try:
                    raw, usage = client.run("create-image-pixen", payload, f"{spec['name']}/base")
                    raw.save(d / "base_raw.png")
                    normalize_base(raw).save(base_path)
                except Exception as e:      # API error after retries, or a creature too big for 64x64
                    print(f"  base FAILED: {e}")
                    continue
                save_meta(d / "base.json", "create-image-pixen", payload, usage)
                stale = [f for f in d.glob("*.png") if not f.name.startswith("base")]
                for f in stale:             # edits of the old base no longer match it
                    f.unlink()
                    f.with_suffix(".json").unlink(missing_ok=True)
                print(f"  base saved ({usage})" + (f", removed {len(stale)} old edits" if stale else ""))
        if base_only:
            continue
        if not base_path.exists():
            if dry:
                print(f"  (edits need {base_path})")
            continue
        base = Image.open(base_path).convert("RGBA")
        todo = [p for p in spec.get("parts", [])
                if (not only or p["name"] in only) and (force and (not only or p["name"] in only)
                                                        or not (d / f"{p['name']}.png").exists())]
        todo = [p for p in todo if not p.get("skip")]

        def one(p):
            payload = edit_payload(spec, p, base)
            if dry:
                return f"  {p['name']}: edit-image-pixen \"{payload['description']}\" seed {payload['seed']}"
            try:
                img, usage = client.run("edit-image-pixen", payload, f"{spec['name']}/{p['name']}")
            except Exception as e:   # one failed edit shouldn't stop the others
                return f"  {p['name']}: FAILED {e}"
            img.save(d / f"{p['name']}.png")
            save_meta(d / f"{p['name']}.json", "edit-image-pixen", payload, usage)
            return f"  {p['name']}: saved ({usage})"

        with ThreadPoolExecutor(max_workers=max(1, workers)) as ex:
            for line in ex.map(one, todo):
                print(line)


# ---------------------------------------------------------------- color conversion

def hls(rgb):
    """(hue 0-1, lightness, saturation) of an (..., 3) array of 0-255 colors."""
    c = np.asarray(rgb, float) / 255.0
    mx, mn = c.max(-1), c.min(-1)
    light, d = (mx + mn) / 2, mx - mn
    sat = np.where(d == 0, 0.0, d / np.maximum(1 - np.abs(2 * light - 1), 1e-9))
    r, g, b = c[..., 0], c[..., 1], c[..., 2]
    dd = np.maximum(d, 1e-9)
    hue = np.where(mx == r, ((g - b) / dd) % 6, np.where(mx == g, (b - r) / dd + 2, (r - g) / dd + 4)) / 6
    return hue, light, sat


def ramp_table(fixed=None):
    """[(ramp name, chars, rgb array 3x3)] for classification: stand-in ramps (+ one fixed ramp)."""
    table = [(name, chars, np.array([ImageColor.getrgb(h) for h in hexes])) for name, (chars, hexes) in STANDIN.items()]
    if fixed:
        table.append((fixed, bp.FIXED_CHARS[fixed], np.array(gs.FIXED[fixed])))
    return table


def shade_centers(values, default):
    """Light/base/dark lightness centres of one ramp's pixels (1D k-means, k=3), so an image's
    own three tones map to the three shades whatever their absolute lightness."""
    if len(np.unique(np.round(values, 3))) < 3:
        return default
    centers = np.percentile(values, [85, 50, 15])
    for _ in range(20):
        idx = np.abs(values[:, None] - centers).argmin(-1)
        new = np.array([values[idx == k].mean() if (idx == k).any() else centers[k] for k in range(3)])
        if np.allclose(new, centers):
            break
        centers = new
    return centers


def to_grid(arr, fixed=None, body=False):
    """RGBA array -> (char grid, ramp-name grid).

    Exact stand-in / fixed colors map directly. Otherwise: very dark -> outline, bright and
    unsaturated -> eye white, gray -> nearest color, anything else picks a ramp by hue and a
    shade by clustering that ramp's lightness in this image into three tones.
    body=True (base creatures): cream/tan gold (bellies, highlights) becomes the main ramp at the
    same shade instead of type color; saturated gold details (beaks, irises) stay type."""
    table = ramp_table(fixed)
    colors = np.concatenate([t[2] for t in table] + [np.array([OUTLINE_RGB, WHITE_RGB])])
    chars = np.array([c for t in table for c in t[1]] + ["o", "w"])
    names = np.array([t[0] for t in table for _ in range(3)] + ["outline", "white"])
    rgb = arr[..., :3].astype(int)
    opaque = arr[..., 3] >= 128
    dist = ((rgb[:, :, None, :] - colors[None, None]) ** 2).sum(-1)
    nearest = dist.argmin(-1)
    hue, light, sat = hls(rgb)
    exact = dist.min(-1) == 0
    outline = light <= OUTLINE_MAX_LIGHT
    white = (light >= WHITE_MIN_LIGHT) & (sat <= WHITE_MAX_SAT)
    gray = sat <= GRAY_MAX_SAT
    ramp_hue = np.array([hls(t[2][1])[0] for t in table])
    gap = np.abs(hue[..., None] - ramp_hue)
    ramp = np.minimum(gap, 1 - gap).argmin(-1)
    free = opaque & ~(exact | outline | white | gray)
    shade = np.zeros(ramp.shape, int)
    for r, (_, _, rgbs) in enumerate(table):
        sel = ramp == r
        centers = shade_centers(light[free & sel], hls(rgbs)[1])
        shade[sel] = np.abs(light[sel][:, None] - centers).argmin(-1)
    by_hue = ramp * 3 + shade
    idx = np.select([exact | gray, outline, white], [nearest, len(chars) - 2, len(chars) - 1], by_hue)
    if body:
        tan = (names[idx] == "type") & ((light >= BODY_PALE_GOLD_LIGHT) | (sat < BODY_TAN_MAX_SAT))
        idx = np.where(tan, np.where(light >= BODY_PALE_GOLD_LIGHT, 0, idx % 3), idx)   # main ramp is 0-2
    grid = np.where(opaque, chars[idx], ".")
    ramps = np.where(opaque, names[idx], "empty")
    return grid, ramps


def load_rgba(path):
    return np.array(Image.open(path).convert("RGBA"))


# ---------------------------------------------------------------- extraction

def grow(m, diag=True):
    out = m.copy()
    out[1:] |= m[:-1]
    out[:-1] |= m[1:]
    out[:, 1:] |= m[:, :-1]
    out[:, :-1] |= m[:, 1:]
    if diag:
        out[1:, 1:] |= m[:-1, :-1]
        out[1:, :-1] |= m[:-1, 1:]
        out[:-1, 1:] |= m[1:, :-1]
        out[:-1, :-1] |= m[1:, 1:]
    return out


def components(mask):
    """8-connected components as lists of (y, x)."""
    seen = np.zeros_like(mask)
    h, w = mask.shape
    out = []
    for y, x in zip(*np.nonzero(mask)):
        if seen[y, x]:
            continue
        seen[y, x] = True
        comp, q = [], deque([(y, x)])
        while q:
            cy, cx = q.popleft()
            comp.append((cy, cx))
            for dy in (-1, 0, 1):
                for dx in (-1, 0, 1):
                    ny, nx = cy + dy, cx + dx
                    if 0 <= ny < h and 0 <= nx < w and mask[ny, nx] and not seen[ny, nx]:
                        seen[ny, nx] = True
                        q.append((ny, nx))
        out.append(comp)
    return out


def nearest(mask, x, y):
    ys, xs = np.nonzero(mask)
    i = int(np.argmin((xs - x) ** 2 + (ys - y) ** 2))
    return int(xs[i]), int(ys[i])


def extract(base, part):
    """Pixels of the edit that belong to the new feature.

    base/var: dicts with grid (chars), ramps (ramp names), rgb (raw array). Returns
    (part char grid, fill mask, body-change fraction)."""
    bgrid, vgrid = base["grid"], part["var"]["grid"]
    bramp, vramp = base["ramps"], part["var"]["ramps"]
    spec = part["spec"]
    raw_dist = np.sqrt(((base["rgb"][..., :3].astype(int) - part["var"]["rgb"][..., :3].astype(int)) ** 2).sum(-1))
    near_body = grow(bgrid != ".")
    over_body = bgrid != "."
    noise = near_body & ((vramp == "outline")
                         | (over_body & (raw_dist < NOISE_RGB))
                         | (over_body & (vramp == bramp) & (raw_dist < SAME_RAMP_RGB))
                         | (~over_body & np.isin(vramp, JITTER_RAMPS)))
    changed = (vgrid != ".") & (vgrid != bgrid)
    allowed = np.ones_like(changed)
    if "box" in spec:
        x0, y0, x1, y1 = spec["box"]
        allowed[:] = False
        allowed[y0:y1 + 1, x0:x1 + 1] = True
    if "chars" in spec:
        allowed &= np.isin(vgrid, list(spec["chars"] + "o"))
    keep = np.zeros_like(changed)
    for comp in components(changed & allowed & ~noise):
        if len(comp) >= spec.get("min_area", 8):
            keep[tuple(np.array(comp).T)] = True
    # The noise rule also drops the feature's root where it meets the body (a teal tail base,
    # say). Take back changed pixels outside the original silhouette that touch the feature.
    keep |= changed & allowed & (bgrid == ".") & grow(keep & (vgrid != "o"))
    fill = keep & (vgrid != "o")
    if not fill.any():
        raise SpecError(f"{spec['name']}: the edit added nothing (after box/chars/min_area)")
    mask = keep | (grow(fill) & (vgrid == "o"))         # plus the outline hugging the new fill
    grid = np.where(mask, vgrid, ".")
    if spec.get("fixed"):
        grid = fold_fixed(grid, part["var"], fill, spec)
    # Close the silhouette where the part ends in empty space; where it meets the creature it
    # stays open, so a head_top drawn over the head joins without a seam.
    grid[grow(fill, diag=False) & (vgrid == ".") & ~mask] = "o"
    body_px = bgrid != "."
    redrawn = body_px & ~grow(mask, diag=True) & ((vgrid == ".") | (raw_dist >= SAME_RAMP_RGB))
    return grid, fill, float(redrawn.sum() / max(1, body_px.sum()))


def fold_fixed(grid, var, fill, spec):
    """Recolor the feature's pixels from the listed ramps into the fixed ramp, by lightness."""
    fixed = spec["fixed"]
    sources = set(spec.get("fixed_from", ["accent", "type"])) | {fixed}
    sel = fill & np.isin(var["ramps"], list(sources))
    if not sel.any():
        return grid
    _, light, _ = hls(var["rgb"][..., :3])
    centers = shade_centers(light[sel], hls(np.array(gs.FIXED[fixed]))[1])
    shade = np.abs(light[sel][:, None] - centers).argmin(-1)
    out = grid.copy()
    out[sel] = np.array(list(bp.FIXED_CHARS[fixed]))[shade]
    return out


def contact_anchor(bgrid, vgrid, fill):
    """Body pixel nearest the middle of where the part touches the body (for a pair of horns or
    wings, the middle of the head or back); None if the part doesn't touch the body."""
    touch = grow(fill, diag=False) & ~fill & (bgrid != ".") & (vgrid != ".")
    if not touch.any():
        return None
    ys, xs = np.nonzero(touch)
    return nearest(bgrid != ".", xs.mean(), ys.mean())


def default_anchors(bgrid):
    opaque = bgrid != "."
    ys, xs = np.nonzero(opaque)
    x0, x1, y0, y1 = xs.min(), xs.max(), ys.min(), ys.max()
    cx, cy, h = xs.mean(), ys.mean(), y1 - y0 + 1
    bottom = opaque & (np.arange(bgrid.shape[0])[:, None] == y1)
    return {
        "pivot": nearest(bottom, (x0 + x1) / 2, y1),
        "aura": nearest(opaque, cx, cy),
        "head_top": nearest(opaque, cx, y0),
        "back": nearest(opaque, cx, cy - 0.2 * h),
        "tail": nearest(opaque, x1, y0 + 0.75 * h),
    }


def crop_box(mask, points):
    ys, xs = np.nonzero(mask)
    xs = list(xs) + [p[0] for p in points]
    ys = list(ys) + [p[1] for p in points]
    return min(xs), min(ys), max(xs), max(ys)


def write_part(slot, name, grid, header):
    src = GRID_DIR / slot / f"{name}.txt"
    src.parent.mkdir(parents=True, exist_ok=True)
    src.write_text("".join(f"# {line}\n" for line in header) + "\n".join("".join(r) for r in grid) + "\n")
    rows = bp.read_grid(src)                      # same validation build_parts applies
    png = PARTS_DIR / slot / f"{name}.png"
    png.parent.mkdir(parents=True, exist_ok=True)
    bp.to_image(rows).save(png)
    return png


def standin_image(grid):
    """Char grid -> RGBA in the stand-in palette (readable, unlike the key colors)."""
    lut = {"o": OUTLINE_RGB, "w": WHITE_RGB}
    for chars, hexes in STANDIN.values():
        lut.update(zip(chars, (ImageColor.getrgb(h) for h in hexes)))
    for ramp, chars in bp.FIXED_CHARS.items():
        lut.update(zip(chars, gs.FIXED[ramp]))
    lut.update({m: (255, 0, 255) for m in MARKER_CHAR.values()})
    img = Image.new("RGBA", (grid.shape[1], grid.shape[0]), (0, 0, 0, 0))
    img.putdata([lut[c] + (255,) if c in lut else (0, 0, 0, 0) for c in grid.ravel()])
    return img


def zoom(img, k, bg=(27, 31, 42, 255)):
    big = img.resize((img.width * k, img.height * k), Image.NEAREST)
    return Image.alpha_composite(Image.new("RGBA", big.size, bg), big)


def hstack(tiles, gap=8, bg=(0, 0, 0, 0)):
    sheet = Image.new("RGBA", (sum(t.width for t in tiles) + gap * (len(tiles) - 1), max(t.height for t in tiles)), bg)
    x = 0
    for t in tiles:
        sheet.paste(t, (x, 0))
        x += t.width + gap
    return sheet


def extract_body(spec):
    d = SRC_DIR / spec["name"]
    if not (d / "base.png").exists():
        print(f"{spec['name']}: no base yet, run generate")
        return
    rgb = load_rgba(d / "base.png")
    bgrid, bramps = to_grid(rgb, body=True)
    base = {"grid": bgrid, "ramps": bramps, "rgb": rgb}
    anchors = default_anchors(bgrid)
    from_part, extracted = {}, []
    for p in spec.get("parts", []):
        if p.get("skip"):
            continue
        src = d / f"{p['name']}.png"
        if not src.exists():
            print(f"  {p['name']}: no edit yet, run generate")
            continue
        vrgb = load_rgba(src)
        if vrgb.shape != rgb.shape:
            print(f"  {p['name']}: edit is {vrgb.shape[1]}x{vrgb.shape[0]}, expected {CANVAS}x{CANVAS}; skipped")
            continue
        vgrid, vramps = to_grid(vrgb, p.get("fixed"))
        var = {"grid": vgrid, "ramps": vramps, "rgb": vrgb}
        try:
            grid, fill, change = extract(base, {"spec": p, "var": var})
        except SpecError as e:
            print(f"  {p['name']}: {e}")
            continue
        if p["slot"] not in from_part:
            a = contact_anchor(bgrid, vgrid, fill)
            if a is None:
                print(f"  ! {p['name']}: doesn't touch the body; {p['slot']} anchor uses the default")
            else:
                from_part[p["slot"]] = a
        extracted.append((p, var, grid, fill, change))
    anchors.update(from_part)
    anchors.update({k: tuple(v) for k, v in spec.get("anchors", {}).items()})
    taken = {}
    for k, xy in anchors.items():
        if xy in taken:
            raise SpecError(f"{spec['name']}: {k} and {taken[xy]} anchors share pixel {xy}; set one in 'anchors'")
        taken[xy] = k

    x0, y0, x1, y1 = crop_box(bgrid != ".", anchors.values())
    body = bgrid[y0:y1 + 1, x0:x1 + 1].copy()
    for k, (x, y) in anchors.items():
        body[y - y0, x - x0] = MARKER_CHAR[k]
    write_part("body", spec["name"], body,
               [f"{spec['name']}: from pixellab_src/{spec['name']}/base.png (pixellab_library.py extract)"]
               + [f"{k} at {xy}" + (" (from part)" if k in from_part else "") for k, xy in anchors.items()])
    print(f"{spec['name']}: body {body.shape[1]}x{body.shape[0]}  "
          + " ".join(f"{k}={anchors[k]}" for k in MARKER_CHAR))

    diff_dir = PREVIEW_DIR / "diff"
    diff_dir.mkdir(parents=True, exist_ok=True)
    for p, var, grid, fill, change in extracted:
        pivot = anchors[p["slot"]]
        px0, py0, px1, py1 = crop_box(grid != ".", [pivot])
        pgrid = grid[py0:py1 + 1, px0:px1 + 1].copy()
        pgrid[pivot[1] - py0, pivot[0] - px0] = "P"
        write_part(p["slot"], p["name"], pgrid,
                   [f"{p['name']}: pixellab_src/{spec['name']}/{p['name']}.png minus base.png", f"pivot at {pivot}"])
        tiles = [Image.fromarray(rgb), Image.fromarray(var["rgb"]), standin_image(var["grid"]), standin_image(grid)]
        hstack([zoom(t, 6) for t in tiles]).save(diff_dir / f"{p['name']}.png")
        flag = "  ! body redrawn, check the diff" if change > BODY_CHANGE_WARN else ""
        print(f"  {p['slot']:8} {p['name']:28} {pgrid.shape[1]}x{pgrid.shape[0]}  {int(fill.sum())} px  "
              f"body change {change:.0%}{flag}")


def sweep_orphans(specs):
    """Remove grids/PNGs for parts and bodies no longer in any spec (or marked skip)."""
    live = {("body", s["name"]) for s in specs}
    live |= {(p["slot"], p["name"]) for s in specs for p in s.get("parts", []) if not p.get("skip")}
    for root, ext in ((GRID_DIR, "txt"), (PARTS_DIR, "png")):
        for f in root.glob(f"*/*.{ext}"):
            if (f.parent.name, f.stem) not in live:
                f.unlink()
                print(f"removed stale {f}")


# ---------------------------------------------------------------- previews

def registry():
    gs._part_cache.clear()
    reg._extents.clear()
    return importlib.reload(reg)


def top_type(p, fallback="normal"):
    if not p["types"]:
        return fallback
    return max(p["types"], key=p["types"].get)


def render(lib, body, picks, ptype, colors):
    chosen = {"body": body, **{s: p for s, p in picks.items()}}
    spr = gs.compose(lib.PARTS_DIR, chosen, "none", random.Random(0))
    spr = gs.place_on_canvas(spr, lib)
    main, accent = gs.pick_main_accent([(gs.hex_to_rgb(c), 1.0) for c in colors], ptype)
    return gs.recolor(spr, gs.build_palette(main, accent, ptype))


def label_grid(cells, cols, path, scale, title=None):
    """cells: [(label, image or None)] in a grid; None draws an empty cell."""
    cw, lh = CANVAS * scale, 12
    rows = -(-len(cells) // cols)
    top = 16 if title else 0
    sheet = Image.new("RGBA", (cols * (cw + 4) + 4, top + rows * (cw + lh + 4) + 4), BG)
    d = ImageDraw.Draw(sheet)
    if title:
        d.text((4, 2), title, fill=(30, 30, 30, 255))
    for i, (label, img) in enumerate(cells):
        x, y = 4 + (i % cols) * (cw + 4), top + 4 + (i // cols) * (cw + lh + 4)
        d.text((x, y), label[:28], fill=(40, 40, 40, 255))
        if img is not None:
            sheet.alpha_composite(img.resize((cw, cw), Image.NEAREST), (x, y + lh))
    path.parent.mkdir(parents=True, exist_ok=True)
    sheet.save(path)
    print(path)


def preview_body(lib, spec):
    by = {p["name"]: p for s in lib.SLOTS.values() for p in s}
    if spec["name"] not in by:
        print(f"{spec['name']}: not extracted")
        return
    body = by[spec["name"]]
    btype = top_type(body)
    own = [(p["slot"], by[p["name"]]) for p in spec.get("parts", []) if p["name"] in by]
    cells = []
    for _, colors in PREVIEW_PALETTES:
        base = Image.open(SRC_DIR / spec["name"] / "base.png").convert("RGBA")
        cells.append(("pixellab base", Image.alpha_composite(Image.new("RGBA", base.size, (27, 31, 42, 255)), base)))
        cells.append((spec["name"], render(lib, body, {}, btype, colors)))
        for slot, p in own:
            cells.append((p["name"], render(lib, body, {slot: p}, top_type(p, btype), colors)))
    label_grid(cells, len(own) + 2, PREVIEW_DIR / f"body_{spec['name']}.png", 3)


def allowed_on(lib, body, slot, p):
    tags = set(body["tags"]) | {f"plan:{body['plans'][0]}"}
    if body["plans"][0] not in p["plans"] or p["name"] in body["deny"]:
        return False
    if p["requires"] and not set(p["requires"]) & tags:
        return False
    return not (set(p["tags"]) & set(body["excludes"]) or set(p["excludes"]) & tags)


def preview_cross(lib):
    """Per slot: rows = bodies, columns = parts; every pairing the rules allow is rendered.
    Gray cell: not allowed by plans/tags/deny. Red frame: rendered for review, but the generator
    never picks it (detached from the body, or too big for the canvas)."""
    colors = PREVIEW_PALETTES[0][1]
    for slot in SLOTS:
        parts = [p for p in lib.SLOTS[slot] if p["name"] != "none"]
        if not parts:
            continue
        cells = []
        for body in lib.BODY:
            for p in parts:
                label = f"{body['name'][3:]}+{p['name'][3:]}"
                if not allowed_on(lib, body, slot, p):
                    cells.append((label, Image.new("RGBA", (CANVAS, CANVAS), (200, 200, 196, 255))))
                    continue
                img = render(lib, body, {slot: p}, top_type(p, top_type(body)), colors)
                if not lib.fits(body, {slot: p}):
                    if img.width > CANVAS:      # expanded canvas: show the bottom-centre 64x64
                        img = img.crop(((img.width - CANVAS) // 2, img.height - CANVAS,
                                        (img.width + CANVAS) // 2, img.height))
                    ImageDraw.Draw(img).rectangle((0, 0, CANVAS - 1, CANVAS - 1), outline=(220, 0, 0, 255), width=2)
                    label += " (detached)" if not lib.attached(body, slot, p) else " (too big)"
                cells.append((label, img))
        label_grid(cells, len(parts), PREVIEW_DIR / f"cross_{slot}.png", 2,
                   title=f"{slot}: rows = bodies, columns = parts ({', '.join(p['name'] for p in parts)})")


def preview_random(lib, n):
    palettes = [["#1565C0", "#FFB300"], ["#2E7D32", "#E53935"], ["#6A1B9A", "#FDD835"], ["#D84315", "#263238"],
                ["#00897B", "#F06292"], ["#37474F", "#FFCA28"], ["#C62828", "#ECEFF1"], ["#F9A825", "#1A237E"]]
    cells = []
    for i in range(n):
        ptype = TYPES[i % len(TYPES)]
        rng = random.Random(i)
        plan, chosen, markings, _ = gs.pick_parts(lib, ptype, rng)
        spr = gs.place_on_canvas(gs.compose(lib.PARTS_DIR, chosen, markings, rng), lib)
        main, accent = gs.pick_main_accent([(gs.hex_to_rgb(c), 1.0) for c in palettes[i % len(palettes)]], ptype)
        spr = gs.recolor(spr, gs.build_palette(main, accent, ptype))
        extras = [p["name"][3:] for s, p in chosen.items() if s != "body" and p["name"] != "none"]
        cells.append((f"{ptype} {chosen['body']['name'][3:]} {' '.join(extras)}", spr))
    label_grid(cells, 10, PREVIEW_DIR / "random.png", 2)


# ---------------------------------------------------------------- cost

def cost():
    ledger = SRC_DIR / "ledger.jsonl"
    usd = gens = n = 0
    if ledger.exists():
        for line in ledger.read_text().splitlines():
            u = json.loads(line).get("usage") or {}
            n += 1
            usd += u.get("usd", 0) or 0
            gens += u.get("generations", 0) or 0
    print(f"ledger: {n} calls, ${usd:.3f}, {gens:g} trial generations")
    from pixellab_api import Client
    print("balance:", Client().balance())


def main(argv=None):
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = ap.add_subparsers(dest="cmd", required=True)
    g = sub.add_parser("generate", help="create missing bases and edits")
    g.add_argument("bodies", nargs="*")
    g.add_argument("--parts", nargs="*", default=(), help="only these parts (with --force: redo them)")
    g.add_argument("--force", action="store_true", help="regenerate even if the image exists")
    g.add_argument("--dry-run", action="store_true")
    g.add_argument("--workers", type=int, default=1,
                   help="parallel edit requests (the account runs one job at a time, so extra workers just queue)")
    g.add_argument("--base-only", action="store_true", help="only the base image (review it before paying for edits)")
    e = sub.add_parser("extract", help="convert + diff + anchors -> parts_pl_src/ and parts_pl/")
    e.add_argument("bodies", nargs="*")
    pv = sub.add_parser("preview", help="review sheets in preview/pixellab/")
    pv.add_argument("bodies", nargs="*")
    pv.add_argument("--cross", action="store_true", help="every part on every body, one sheet per slot")
    pv.add_argument("--random", type=int, metavar="N", help="N real random picks")
    sub.add_parser("cost", help="spend so far and live balance")
    a = ap.parse_args(argv)
    try:
        if a.cmd == "cost":
            cost()
            return 0
        specs = load_specs(getattr(a, "bodies", ()))
        if a.cmd == "generate":
            generate(specs, set(a.parts), a.force, a.dry_run, a.workers, a.base_only)
        elif a.cmd == "extract":
            for spec in specs:
                extract_body(spec)
            if not a.bodies:
                sweep_orphans(specs)
        elif a.cmd == "preview":
            lib = registry()
            if a.cross:
                preview_cross(lib)
            if a.random:
                preview_random(lib, a.random)
            if a.bodies or not (a.cross or a.random):
                for spec in specs:
                    preview_body(lib, spec)
    except (SpecError, reg.SpecError, bp.SourceError) as err:
        print(f"error: {err}", file=sys.stderr)
        return 1
    return 0


if __name__ == "__main__":
    sys.exit(main())