#!/usr/bin/env python3
"""Animation sheets (4 frames x 5 views) for one creature, in the imp's hack-sheet layout.

  make_anim.py gen <id> <anim> [front|back|both] [--suffix S]   fal requests (nano-banana-pro/edit)
  make_anim.py pix <id> <anim> [--front F] [--back B]            pixelize, align, stack, preview

<anim> is walk or attack. Files (in art/generated/sprites/anim/):
  <id>-<anim>-front-raw<S>.png   request 1: the 5x4 sheet, rows 1-3 (front, front 3/4, side) used
  <id>-<anim>-back-raw<S>.png    request 2: 2x4 sheet, back three-quarter and back
  <id>-<anim>-32.png (128x160), <id>-<anim>-64.png (256x320), <id>-<anim>-<size>-preview.png

Run with the Pillow venv: /tmp/sprvenv/bin/python -I make_anim.py ...
The key is read from ~/.config/dungeoncraft/fal-key and never printed.

Pixel grid: every view row is pixelized with tools/pixelize_sprite.py (--grid 1x4, the idle
sheet's palette). The step (source pixels per output pixel) is chosen per row so that the
figures' "core" height (the body without thin weapon parts sticking out above or below, see
core()) is the median core height of the matching idle view; each frame is then shifted so that
its core bottom sits on the feet row and its core centre on the idle view's core centre.
"""
import argparse, base64, io, json, os, statistics, subprocess, sys, time, urllib.error, urllib.request

from PIL import Image

HERE = os.path.dirname(os.path.abspath(__file__))
SPR = os.path.normpath(os.path.join(HERE, ".."))
ANIM = os.path.join(SPR, "anim")
ROOT = os.path.normpath(os.path.join(SPR, "..", "..", ".."))
TOOL = os.path.join(ROOT, "tools", "pixelize_sprite.py")
sys.path.insert(0, HERE)
sys.path.insert(0, os.path.join(ROOT, "tools"))
from creatures import BY_ID as _BY_ID, GREEN, HOVER, SIEGE  # noqa: E402
import pixelize_sprite as ps  # noqa: E402

# the imp (kind Worker) has its own folder: raw = the lower row of imp-sprite/raw-1.png
IMP_DIR = os.path.normpath(os.path.join(SPR, "..", "imp-sprite"))
BY_ID = dict(_BY_ID)
BY_ID["imp"] = dict(id="imp", kind="Worker", name="Pit Imp",
                    desc="A small goblin-like worker demon: short and stocky, big head, oversized pointed ears and two small black horns, orange skin, glowing yellow eyes, a wide mischievous grin with tiny fangs, a ragged dark-red tunic with a rope belt, bare feet with claws, a small iron pickaxe with a wooden handle carried over its right shoulder. Thin tail.")
QUAD = {"spider", "cavebear", "beetle"}
PROGRESS = os.path.join(ANIM, "progress.json")


def idle_sheet(cid, size):
    if cid == "imp":
        return os.path.join(IMP_DIR, f"imp-{size}.png")
    return os.path.join(SPR, f"{cid}-{size}.png")


VIEWS = ["front", "front three-quarter", "side", "back three-quarter", "back"]
SIZES = ((32, 31), (64, 63))  # cell size, feet row

# Per creature: the weapon (as it appears in the idle sheet) and the attack frames.
# Front views: what the frames show from the front; back: the same, seen from behind.
WEAPON = {
    "knight": "the same short straight steel sword in the right hand",
    "kobold": "the same wooden spear with the steel tip in the right hand",
    "imp": "the same small iron pickaxe over the right shoulder",
    "mage": "the same wooden staff with the glowing yellow gem in the right hand",
    "archer": "the same wooden bow in the left hand",
    "dwarf": "the same steel pickaxe with the wooden handle in the right hand",
    "paladin": "the same white holy sword with the golden hilt in the right hand",
    "jailer": "the same coiled iron chain in the right hand",
    "priestess": "the same small fireball above the right hand",
    "chainbreaker": "the same steel battle axe in each hand",
    "gargoyle": "the same iron chain around one arm",
    "cogling": "the same steel hammer in the right hand",
    "revenant": "the same rusty steel sword in the right hand",
    "giant": "the same big wooden tree-trunk club in the right hand",
    "battlemage": "the same wooden staff with the burning flame on top in the right hand",
    "trader": "the same big backpack and the small glowing lantern in the right hand",
    "thief": "the same short steel dagger in the right hand and the sack of gold on the back",
    "ranger": "the same wooden bow in the left hand",
    "inquisitor": "the same glowing yellow holy book in both hands",
    "priest": "the same golden halo floating above the head",
}
ATTACK = {
    "knight": dict(
        what="a sword swing: the knight swings the sword with the right hand from high over the right shoulder diagonally down across the body at an enemy in front of it",
        front=["wind-up - the sword raised high over the right shoulder, the blade pointing up and back, body turned slightly back",
               "strike - the sword swinging down diagonally in front of the helm, body turning forward",
               "impact - the sword fully swung through, the blade low in front of the body at the left hip, body leaning forward",
               "recover - the sword pulled back up to the ready position at chest height, body upright"],
        back=["wind-up - the sword raised high over the right shoulder, the blade showing above the helm",
              "strike - the sword coming down, partly hidden by the helm and shoulders",
              "impact - the sword in front of the knight, mostly hidden by the body, only the blade tip shows low at the left side",
              "recover - the sword pulled back up, its blade visible beside the right shoulder"]),
    "kobold": dict(
        what="a spear thrust: the kobold holds the spear level with both hands and stabs it forward at an enemy in front of it",
        front=["wind-up - the spear pulled back, held level at waist height with both hands, the tip pointing forward, body leaning back",
               "strike - the spear pushed forward, body leaning forward",
               "impact - the spear fully thrust forward, arms stretched out, body lunging forward, a small white spark at the spear tip",
               "recover - the spear pulled back to the ready position, body upright"],
        back=["wind-up - the spear pulled back and held level at waist height, its butt end sticking out behind the kobold",
              "strike - the spear pushed forward away from the viewer, mostly hidden by the body",
              "impact - the spear fully thrust forward, arms stretched out, the spear hidden by the body except a short piece of shaft, a small white spark at the tip beyond the head",
              "recover - the spear pulled back to the ready position, the shaft visible beside the body"]),
}
ATTACK.update({
    "ogre": dict(
        what="a fist smash: the ogre raises both huge stone fists high and slams them down on an enemy in front of it",
        front=["wind-up - both fists raised high above the head, body leaning back",
               "strike - the fists coming down in front of the face, body bending forward",
               "impact - the fists smashed down low in front of the feet, body bent forward, a few small grey stone chips flying",
               "recover - straightening up, the arms hanging low and ready"],
        back=["wind-up - both fists raised high above the head, seen from behind",
              "strike - the fists coming down in front, the arms partly hidden by the head and shoulders",
              "impact - bent forward, the fists down in front hidden by the body, the broad back hunched",
              "recover - straightening up, the arms hanging low at the sides"]),
    "spider": dict(
        what="a pouncing bite: the spider rears up on its back legs, raises its front legs and lunges forward to bite an enemy in front of it with its fangs",
        front=["wind-up - the front of the body reared up, the two front legs raised high, fangs open",
               "strike - lunging forward and down, the front legs reaching forward",
               "impact - the body pressed low and forward, fangs biting, the front legs stabbing down",
               "recover - pulling back to the normal low stance, all eight legs on the ground"],
        back=["wind-up - the front of the body reared up, the two front legs raised high above the round back",
              "strike - lunging forward away from the viewer, the front legs reaching forward",
              "impact - the body pressed low and forward, the round back seen from behind",
              "recover - back in the normal low stance, all eight legs on the ground"]),
    "mage": dict(
        what="casting a shadow bolt: the mage raises the staff and thrusts it forward, the gem flashing and a small purple-black orb of shadow magic shooting from it",
        front=["wind-up - the staff raised high and drawn back, the gem glowing brighter, the free hand raised",
               "strike - the staff swung forward and pointed at the enemy, a small dark-purple magic orb forming at the gem",
               "impact - the staff thrust forward, the small purple orb leaving the gem with a few sparks",
               "recover - the staff pulled back upright, the gem glow fading"],
        back=["wind-up - the staff raised high above the hood, the gem glowing",
              "strike - the staff swung forward away from the viewer, mostly hidden by the hood, a purple glow beyond",
              "impact - the staff thrust forward, hidden by the body except its upper end, a small purple orb beyond",
              "recover - the staff pulled back upright beside the body"]),
    "archer": dict(
        what="shooting an arrow: the archer raises the bow, draws an arrow from the quiver, pulls the string back and shoots forward",
        front=["wind-up - drawing an arrow from the quiver on the back with the right hand, the bow lowered in the left hand",
               "strike - the bow raised and held out forward, the arrow nocked, the string pulled back to the cheek",
               "impact - the arrow just released, the string snapped forward, the bow still raised",
               "recover - the bow lowered back to the upright idle position"],
        back=["wind-up - reaching to the quiver on the back with the right hand",
              "strike - the bow raised forward away from the viewer, the right elbow pulled back drawing the string",
              "impact - the arrow released, the bow held out forward, the right arm relaxed",
              "recover - the bow lowered back to the upright idle position"]),
    "dwarf": dict(
        what="a pickaxe blow: the dwarf swings the pickaxe from over the right shoulder down in front of it onto an enemy",
        front=["wind-up - the pickaxe raised high over the right shoulder with both hands, body leaning back",
               "strike - the pickaxe swinging down over the helmet",
               "impact - the pickaxe head struck down low in front of the feet, body bent forward, a few small sparks",
               "recover - the pickaxe lifted back onto the right shoulder"],
        back=["wind-up - the pickaxe raised high over the right shoulder, its head showing above the helmet",
              "strike - the pickaxe coming down, partly hidden by the helmet and shoulders",
              "impact - bent forward, the pickaxe down in front hidden by the body, only the handle end shows",
              "recover - the pickaxe lifted back onto the right shoulder"]),
    "paladin": dict(
        what="a holy sword swing: the paladin swings the white sword with both hands from high over the right shoulder diagonally down across the body at an enemy in front of it, the blade glowing",
        front=["wind-up - the sword raised high over the right shoulder with both hands, the blade pointing up and back, glowing",
               "strike - the glowing sword swinging down diagonally in front of the helmet",
               "impact - the sword fully swung through, the blade low in front of the body at the left hip, a small golden spark",
               "recover - the sword raised back to the upright ready position"],
        back=["wind-up - the sword raised high over the right shoulder, the blade showing above the helmet",
              "strike - the sword coming down, partly hidden by the helmet and shoulders",
              "impact - the sword in front of the paladin, mostly hidden by the body, only the blade tip shows low at the left side",
              "recover - the sword raised back up, its blade visible beside the right shoulder"]),
})

WALK_FRONT = ["contact - the right leg stepping forward with the heel down, the left leg behind, arms swinging opposite to the legs",
              "passing - the legs passing each other under the body, the body slightly higher",
              "contact - the left leg stepping forward with the heel down, the right leg behind, arms swinging the other way",
              "passing - the legs passing each other under the body, the body slightly higher"]
WALK_BACK = WALK_FRONT
HOVER_FRAMES = ["upstroke - the wings raised straight up high above the head, the wing tips pointing up, the body at its highest point",
                "downstroke - the wings swept out level to the sides, fully spread",
                "down - the wings pulled down below the body, the wing tips pointing down, the body at its lowest point",
                "recovery - the wings half raised and folded a little, swinging back up"]
SIEGE_FRAMES = ["the wheels in their start position, the spokes or wheel marks at 0 degrees",
                "the wheels turned a quarter turn further (the wheel marks visibly rotated)",
                "the wheels turned another quarter turn",
                "the wheels turned another quarter turn, the machine rocking slightly"]
QUAD_FRAMES = {
    "spider": ["legs 1 and 3 on the left and 2 and 4 on the right stepping forward, the others pushing back",
               "the legs passing, the body slightly raised",
               "the other set of legs stepping forward",
               "the legs passing, the body slightly raised"],
    "beetle": ["the front and back legs on the left and the middle leg on the right stepping forward",
               "the legs passing under the shell, the body slightly raised",
               "the other three legs stepping forward",
               "the legs passing under the shell, the body slightly raised"],
    "cavebear": ["the right front paw and the left hind paw stepping forward, the head low",
                 "the paws passing under the body, the shaggy body slightly raised",
                 "the left front paw and the right hind paw stepping forward",
                 "the paws passing under the body, the shaggy body slightly raised"],
}
DIE = {
    # (what, front frames, back frames)
    "biped": ("a death: the character is struck by a fatal blow and falls over backwards (away from the direction it faces) onto its back, where it lies still - in every view it falls the same way, backwards relative to its own facing",
              ["hit - recoiling from the blow, still standing upright on both feet, head thrown back, eyes shut in pain, arms flung out",
               "falling - tipping over backwards, the knees buckling, the arms flailing",
               "hitting the ground - landing on its back, the body flat and low, limbs bouncing up a little",
               "lying still - lying flat on its back on the ground, limbs limp, eyes closed or crossed out, motionless"],
              ["hit - recoiling from the blow, still standing upright on both feet, the head thrown back towards the viewer, arms flung out",
               "falling - tipping over backwards towards the viewer, the knees buckling",
               "hitting the ground - landing on its back, the body flat and low on the ground",
               "lying still - lying flat on its back on the ground, motionless"]),
    "quad": ("a death: the creature is struck by a fatal blow, staggers, collapses and rolls over onto its side where it lies still - the same way in every view",
             ["hit - recoiling from the blow, still standing on all its legs, flinching back, eyes shut in pain",
              "falling - the legs giving way, the body collapsing and tipping over sideways",
              "hitting the ground - crashing down onto its side, the legs in the air",
              "lying still - lying on its side or back, the legs curled up and limp, motionless"],
             ["hit - recoiling from the blow, still standing on all its legs, flinching",
              "falling - the legs giving way, the body collapsing and tipping over sideways",
              "hitting the ground - crashing down onto its side, the legs in the air",
              "lying still - lying on its side or back, the legs curled up and limp, motionless"]),
    "hover": ("a death: the flying creature is struck by a fatal blow, its wings fail and it drops out of the air to the ground, where it lies still",
              ["hit - recoiling in the air from the blow, wings jerked up, eyes shut in pain",
               "falling - dropping down, the wings folding and limp, the body tumbling",
               "hitting the ground - crashing onto the ground on its back, the wings splayed",
               "lying still - lying on the ground with the wings spread flat and limp, eyes closed or crossed out, motionless"],
              ["hit - recoiling in the air from the blow, wings jerked up",
               "falling - dropping down, the wings folding and limp",
               "hitting the ground - crashing onto the ground, the wings splayed",
               "lying still - lying on the ground with the wings spread flat and limp, motionless"]),
    "siege": ("its destruction: the machine is smashed - planks crack and splinter, a wheel breaks off, and it collapses into a heap of broken wood",
              ["hit - the machine jolted by a heavy blow, cracks in the planks, a few wood splinters flying, still standing on its wheels",
               "breaking - a front wheel breaking off, the frame tilting and sagging, planks coming loose",
               "collapsing - the frame crashing down, the roof or arm broken, planks and wheels falling apart",
               "wreck - a low flat heap of broken planks, beams and a loose wheel on the ground, motionless"],
              ["hit - the machine jolted by a heavy blow, cracks in the planks, still standing on its wheels",
               "breaking - a wheel breaking off, the frame tilting and sagging",
               "collapsing - the frame crashing down, planks and wheels falling apart",
               "wreck - a low flat heap of broken planks, beams and a loose wheel on the ground, motionless"]),
}


def kind_of(cid):
    return "siege" if cid in SIEGE else "hover" if cid in HOVER else "quad" if cid in QUAD else "biped"


def bg_of(cid):
    return ("green (#00FF00)", (0, 255, 0)) if cid in GREEN else ("magenta (#FF00FF)", (255, 0, 255))


def anim_text(cid, anim, part):
    c = BY_ID[cid]
    weapon = WEAPON.get(cid, "")
    name = c['name'].lower()
    k = kind_of(cid)
    if anim == "walk":
        if k == "hover":
            what = f"a looping hovering flight cycle: the {name} hovers in place in the air, flapping its wings"
            frames = HOVER_FRAMES
            feet = "The creature stays in the middle of its cell (it does not fly across the cell); the wings flap with big, clearly different wing positions in the four frames and the body bobs up and down a little."
        elif k == "siege":
            what = f"rolling forward: the {name} rolls forward in the direction it faces, its wooden wheels turning"
            frames = SIEGE_FRAMES
            feet = "The machine rolls in place (it does not move across the cell); the wheels stand on the baseline of the row and visibly turn from frame to frame."
        else:
            what = f"a looping walk cycle: the {name} walks forward in the direction it faces, at a steady pace" + (
                f", carrying {weapon} as in the idle pose" if weapon else "")
            frames = QUAD_FRAMES.get(cid) or (WALK_FRONT if part == "front" else WALK_BACK)
            feet = "The figure walks in place (it does not move across the cell); the feet stand on the baseline of the row."
    elif anim == "die":
        what, ff, bf = DIE[k]
        frames = ff if part == "front" else bf
        feet = ("Frame 1 has the same size and position as the idle figure; in frames 2-4 the body sinks to the ground, the lowest part always on the baseline of the row, and stays within its own cell. "
                "Frame 4 lies flat and low on the ground.")
    else:
        a = ATTACK[cid]
        what = a["what"]
        frames = a[part]
        feet = "The feet stay in place in all four frames; only arms, weapon, head and upper body move."
    cols = "; ".join(f"{i + 1} {f}" for i, f in enumerate(frames))
    return what, cols, feet


def prompt_front(cid, anim):
    c = BY_ID[cid]
    bgname, _ = bg_of(cid)
    what, cols, feet = anim_text(cid, anim, "front")
    if anim == "die":
        hold = "the character keeps its clothes and equipment (a held weapon may drop beside it in frames 3-4)"
        if cid in SIEGE:
            hold = "the machine keeps its colours and materials even when broken"
    else:
        hold = f"the character keeps {WEAPON[cid]}" if cid in WEAPON else "the character keeps all its equipment"
    return f"""The attached image is a pixel art sprite sheet of the "{c['name']}" in five views (left to right: front, front three-quarter turned right, side facing right, back three-quarter turned away to the right, back). {c['desc']}

Create a new pixel art sprite sheet of exactly the same character - same proportions, colours, outline, pixel size, clothes and equipment - showing {what}.

Layout: a grid of 5 rows and 4 columns, every figure the same size, evenly spaced, each in its own equal cell, all standing on the same baseline within their row.
Rows (one view per row, top to bottom): 1 front view facing the viewer, 2 front three-quarter view turned to the right, 3 side view facing right, 4 back three-quarter view turned away to the right, 5 back view.
Columns (animation frames, left to right): {cols}.
{feet}
Every frame of a row keeps exactly that row's view direction (the side view row stays in profile facing right in all four frames{", also when falling" if anim == "die" else ""}, the back rows never show the face), and {hold} in every single frame.

Style: identical to the attached sheet - crisp retro pixel art, hard square pixels, one-pixel near-black outline, limited palette, no anti-aliasing, no blur. Flat pure {bgname} background everywhere, no ground, no shadows, no text, no labels, no frame lines, no grid lines, no motion lines."""


def prompt_back(cid, anim, retry=False):
    c = BY_ID[cid]
    bgname, _ = bg_of(cid)
    what, cols, feet = anim_text(cid, anim, "back")
    note = ("Every one of the eight figures, including the strike frames, is seen from behind: no visor, no face, no eyes and no chest emblem anywhere. "
            "The back three-quarter row shows the body turned clearly to the right (a little of its right side visible), the back row shows it straight from behind.\n") if retry else ""
    return f"""The first attached image is a pixel art sprite sheet of the "{c['name']}" in five idle views (front, front three-quarter, side, back three-quarter, back). The second attached image shows the same character's animation ({what}) in four frames (columns) for the front, front three-quarter and side views.

Create the same four-frame animation for the two views seen FROM BEHIND, where the character faces away from the viewer and we see the back of its head, its back and its clothes from behind; the face and eyes are never visible:
{note}Row 1: back three-quarter view, the character turned away from the viewer towards the upper right.
Row 2: back view, the character seen straight from behind.
Columns (frames, left to right): {cols}.
{feet}

Layout: a grid of exactly 2 rows and 4 columns, every figure the same size and the same character size as in the attached sheets, evenly spaced, each in its own equal cell, all standing on the same baseline.
Style: identical to the attached sheets - same character, proportions, colours, outline and pixel size, crisp retro pixel art, hard square pixels, no anti-aliasing, no blur. Flat pure {bgname} background everywhere, no ground, no shadows, no text, no labels, no frame lines, no grid lines."""


def idle_crop(cid):
    """The --crop of the idle sheet (scripts/extra.txt), as fractions, or None."""
    p = os.path.join(HERE, "extra.txt")
    if os.path.exists(p):
        for line in open(p):
            parts = line.split()
            if parts and parts[0] == cid and "--crop" in parts:
                return [float(v) for v in parts[parts.index("--crop") + 1].split(",")]
    return None


def clean_ref(img, key):
    """The image with every background pixel (and frame line) painted in the flat key colour."""
    img = img.convert("RGB")
    mask = ps.background_mask(img)
    w, h = img.size
    out = img.copy()
    px = out.load()
    for y in range(h):
        row = mask[y]
        for x in range(w):
            if row[x]:
                px[x, y] = key
    return out


def uri(img):
    b = io.BytesIO()
    img.save(b, "PNG")
    return "data:image/png;base64," + base64.b64encode(b.getvalue()).decode()


def fal(prompt, images, aspect):
    key = open(os.path.expanduser("~/.config/dungeoncraft/fal-key")).read().strip()
    body = {"prompt": prompt, "image_urls": [uri(i) for i in images], "aspect_ratio": aspect,
            "resolution": "2K", "num_images": 1, "output_format": "png"}
    req = urllib.request.Request("https://fal.run/fal-ai/nano-banana-pro/edit", data=json.dumps(body).encode(),
                                 headers={"Authorization": "Key " + key, "Content-Type": "application/json"})

    for attempt in range(5):
        try:
            with urllib.request.urlopen(req, timeout=900) as r:
                res = json.load(r)
            break
        except urllib.error.HTTPError as e:
            msg = e.read().decode(errors="replace")[:500]
            if e.code in (429, 500, 502, 503, 504) and attempt < 4:
                print(f"  HTTP {e.code}, retrying in {60 * (attempt + 1)} s")
                time.sleep(60 * (attempt + 1))
                continue
            raise SystemExit(f"HTTP {e.code}: {msg}")
        except (urllib.error.URLError, ConnectionError, TimeoutError) as e:
            if attempt < 4:
                print(f"  {e}, retrying in {60 * (attempt + 1)} s")
                time.sleep(60 * (attempt + 1))
                continue
            raise
    count_image()
    with urllib.request.urlopen(res["images"][0]["url"], timeout=300) as r:
        return Image.open(io.BytesIO(r.read())).convert("RGB")


def update_progress(fn):
    """Apply fn(dict) to progress.json under a file lock."""
    import fcntl
    os.makedirs(ANIM, exist_ok=True)
    with open(PROGRESS + ".lock", "w") as lk:
        fcntl.flock(lk, fcntl.LOCK_EX)
        d = json.load(open(PROGRESS)) if os.path.exists(PROGRESS) else {"images": 0, "price_per_image": 0.15, "sheets": {}}
        fn(d)
        d["cost_usd"] = round(d["images"] * d.get("price_per_image", 0.15), 2)
        tmp = PROGRESS + ".tmp"
        json.dump(d, open(tmp, "w"), indent=1, sort_keys=True)
        os.replace(tmp, PROGRESS)


CURRENT = {}


def count_image():
    def f(d):
        d["images"] += 1
        sh = d["sheets"].setdefault(CURRENT.get("sheet", "?"), {"status": "generating", "images": 0})
        sh["images"] = sh.get("images", 0) + 1
    update_progress(f)


def idle_ref(cid):
    if cid == "imp":
        raw = Image.open(os.path.join(IMP_DIR, "raw-1.png")).convert("RGB")
        f = [0, 0.27, 1, 0.72]
    else:
        raw = Image.open(os.path.join(SPR, f"{cid}-raw.png")).convert("RGB")
        f = idle_crop(cid)
    if f:
        w, h = raw.size
        raw = raw.crop((int(f[0] * w), int(f[1] * h), int(f[2] * w), int(f[3] * h)))
    raw = raw.resize((raw.width // 2, raw.height // 2), Image.LANCZOS)
    return clean_ref(raw, bg_of(cid)[1])


def row_bands(img, n):
    """The horizontal bands (y0, y1) of figure rows: split at empty rows (bands lower than 6% of
    the image are dropped); if that does not give n bands, n bands split at the emptiest rows
    near the even cuts."""
    mask = [[not v for v in r] for r in ps.background_mask(img)]
    occ = [sum(r) for r in mask]
    h = len(occ)
    raw, y0 = [], None
    for y, v in enumerate(occ + [0]):
        if v > 2 and y0 is None:
            y0 = y
        elif v <= 2 and y0 is not None:
            raw.append((y0, y))
            y0 = None
    # bands lower than 6% of the image (a halo or sparkles floating above a row, specks) belong
    # to the next band (the last one to the previous band)
    bands, pend = [], None
    for b in raw:
        if b[1] - b[0] > 0.06 * h:
            bands.append((pend if pend is not None else b[0], b[1]))
            pend = None
        elif pend is None:
            pend = b[0]
    if pend is not None and bands:
        bands[-1] = (bands[-1][0], h)
    if len(bands) == n:
        # widen to the midpoints of the gaps
        cuts = [0] + [(bands[i][1] + bands[i + 1][0]) // 2 for i in range(n - 1)] + [h]
        return list(zip(cuts[:-1], cuts[1:]))
    print(f"  note: {len(bands)} bands found by gaps, wanted {n}; cutting at the emptiest rows")
    cuts = [0]
    for k in range(1, n):
        c = k * h // n
        lo, hi = c - h // (2 * n), c + h // (2 * n)
        cuts.append(min(range(lo, hi), key=lambda y: (occ[y], abs(y - c))))
    cuts.append(h)
    return list(zip(cuts[:-1], cuts[1:]))


def cmd_gen(a):
    os.makedirs(ANIM, exist_ok=True)
    CURRENT["sheet"] = f"{a.id}-{a.anim}"
    ref = idle_ref(a.id)
    front = os.path.join(ANIM, f"{a.id}-{a.anim}-front-raw{a.suffix}.png")
    if a.part in ("front", "both"):
        img = fal(prompt_front(a.id, a.anim), [ref], "4:5")
        img.save(front)
        print("wrote", front)
    if a.part in ("back", "both"):
        src = a.front_src or front
        fr = Image.open(src).convert("RGB")
        bands = row_bands(fr, 5)
        top = fr.crop((0, 0, fr.width, bands[2][1]))
        img = fal(prompt_back(a.id, a.anim, a.strict), [ref, top], "16:9")
        p = os.path.join(ANIM, f"{a.id}-{a.anim}-back-raw{a.suffix}.png")
        img.save(p)
        print("wrote", p)


# ---- pixel grid -------------------------------------------------------------------------

def opaque(cell):
    w, h = cell.size
    px = cell.load()
    return [[px[x, y][3] > 128 for x in range(w)] for y in range(h)]


def longest_run(row):
    best, start, cur = (0, 0), 0, 0
    for x, v in enumerate(row + [False]):
        if v:
            if cur == 0:
                start = x
            cur += 1
        else:
            if cur > best[0]:
                best = (cur, start)
            cur = 0
    return best


def core(cell, size):
    """(top, bottom, centre x) of the body: top = the first row whose longest opaque run is at
    least 20% of the cell wide (a spear tip or blade above the head is narrower), bottom = the
    last row with a run of at least 3 px (feet; a thin shaft below them is not), centre = the
    median centre of the longest runs of the rows in between. None if the cell is empty."""
    m = opaque(cell)
    runs = [longest_run(r) for r in m]
    wide = max(3, round(0.2 * size))
    tops = [y for y, (n, _) in enumerate(runs) if n >= wide]
    bots = [y for y, (n, _) in enumerate(runs) if n >= 3]
    if not tops or not bots:
        return None
    top, bot = tops[0], bots[-1]
    cx = statistics.median(s + n / 2 for n, s in runs[top:bot + 1] if n >= wide)
    return top, bot, cx


def body_width(cell, size):
    """Median longest opaque run over the core rows (head and body width, for checking that the
    animation keeps the idle proportions)."""
    c = core(cell, size)
    runs = [longest_run(r)[0] for r in opaque(cell)[c[0]:c[1] + 1]]
    return statistics.median(runs)


def cells(sheet, size, row):
    return [sheet.crop((c * size, row * size, (c + 1) * size, (row + 1) * size)) for c in range(sheet.width // size)]


def run_tool(raw_crop, n, size, feet, step, palette, add):
    tmp = f"/tmp/make_anim_{os.getpid()}"
    raw_crop.save(tmp + "-in.png")
    cmd = [sys.executable, "-I", TOOL, tmp + "-in.png", tmp, "--grid", f"1x{n}", "--size", str(size), "--feet", str(feet),
           "--outline-weight", "0.6", "--palette", palette]
    if step:
        cmd += ["--step", f"{step:.4f}"]
    if add:
        cmd += ["--add-colors", add]
    r = subprocess.run(cmd, capture_output=True, text=True)
    if r.returncode:
        raise SystemExit(r.stderr)
    out = Image.open(tmp + f"-{size}.png").convert("RGBA")
    out.load()
    for f in (tmp + "-in.png", tmp + f"-{size}.png", tmp + f"-{size}-preview.png"):
        if os.path.exists(f):
            os.remove(f)
    last = r.stdout.strip().splitlines()[-1]
    used = float(last.split("step ")[1].split(",")[0])
    return out, used


def shift(cell, dx, dy):
    out = Image.new("RGBA", cell.size, (0, 0, 0, 0))
    out.alpha_composite(cell, (max(dx, 0), max(dy, 0)), (max(-dx, 0), max(-dy, 0)))
    lost = sum(1 for p in cell.get_flattened_data() if p[3] > 128) - sum(1 for p in out.get_flattened_data() if p[3] > 128)
    return out, lost


def pix_source(raw_crop, n, target, size, feet, palette, add, ref=None):
    """All n figures of one raw row, at the step that makes their median core height = target
    (ref = a figure index: only that figure's core height counts, for die rows whose later
    frames lie on the ground). Also returns each figure's horizontal offset in source pixels
    from the centre of its equal column (the tool centres every figure in its cell; die rows
    need the original positions back)."""
    step = None
    for _ in range(4):  # the core height is linear in 1/step; converge in a few passes
        sheet, used = run_tool(raw_crop, n, size, feet, step, palette, add)
        cs = [core(x, size) for x in cells(sheet, size, 0)]
        hs = [c[1] - c[0] for c in ([cs[ref]] if ref is not None else cs) if c]
        new = used * statistics.median(hs) / target
        if step and abs(new - step) / step < 0.01:
            break
        step = new
    sheet, used = run_tool(raw_crop, n, size, feet, step, palette, add)
    offs = [0.0] * n
    try:
        _, boxes = ps.figure_boxes(raw_crop.convert("RGB"), 1, n)
        if len(boxes) == n:
            col = raw_crop.width / n
            offs = [(b[0] + b[2]) / 2 - (i + 0.5) * col for i, b in enumerate(boxes)]
    except Exception as e:  # noqa: BLE001
        print("  note: no figure offsets:", e)
    return cells(sheet, size, 0), used, offs


def bbox(cell):
    b = cell.getchannel("A").point(lambda v: 255 if v > 128 else 0).getbbox()
    return b


def strip_lines(img, key):
    """Paint dark horizontal lines across nearly the whole image (ground lines, row separators
    the model drew despite the prompt; at most 12 px thick) in the key colour. Wide dark
    figure parts (a row of spider legs) are not lines: they are thicker or have gaps."""
    w, h = img.size
    px = img.load()
    corners = [px[x, y] for x, y in ((2, 2), (w - 3, 2), (2, h - 3), (w - 3, h - 3))]
    c0 = corners[0]
    d2 = lambda p, q: sum((a - b) ** 2 for a, b in zip(p, q))
    if all(d2(c, c0) < 400 for c in corners) and d2(c0, key) > 3600:
        # the model drew another flat background (white, pale blue) instead of the key colour
        img = img.copy()
        px = img.load()
        for y in range(h):
            for x in range(w):
                if d2(px[x, y], c0) < 2000:
                    px[x, y] = key
        print(f"  note: background {c0} keyed out")
    darkrow = [sum(1 for x in range(0, w, 4) if sum(px[x, y]) < 160) > 0.75 * (w / 4) for y in range(h)]
    out, y = None, 0
    while y < h:
        if darkrow[y]:
            y1 = y
            while y1 < h and darkrow[y1]:
                y1 += 1
            if y1 - y <= 12:
                if out is None:
                    out = img.copy()
                    opx = out.load()
                for yy in range(max(0, y - 1), min(h, y1 + 1)):
                    for x in range(w):
                        opx[x, yy] = key
            y = y1
        else:
            y += 1
    if out is not None:
        print("  note: removed horizontal lines")
    return out or img


def parse_view(spec, default_file):
    """'FILE:BAND/NBANDS:COUNT:PICKS[;...]' -> [(file, band, nbands, count, [picks])]; FILE may
    be empty (= default_file), PICKS like 0,2,1,3 (m = mirrored, e.g. m1)."""
    out = []
    for seg in spec.split(";"):
        parts = seg.split(":")
        f, bn, count, picks = parts[:4]
        ref = int(parts[4]) if len(parts) > 4 else None  # die: the figure that sets the scale
        band, nb = (int(v) for v in bn.split("/"))
        out.append((f or default_file, band, nb, int(count), picks.split(","), ref))
    return out


def cmd_pix(a):
    pre = os.path.join(ANIM, f"{a.id}-{a.anim}")
    front = a.front or pre + "-front-raw.png"
    back = a.back or pre + "-back-raw.png"
    specs = [f"{front}:{v}/5:4:0,1,2,3" for v in range(3)] + [f"{back}:{v}/2:4:0,1,2,3" for v in range(2)]
    for v in a.view or []:
        k, spec = v.split("=", 1)
        specs[int(k)] = spec
    views = [parse_view(sp, front) for sp in specs]
    images, bands = {}, {}
    for segs in views:
        for f, band, nb, count, picks, _ in segs:
            p = f if os.path.isabs(f) or os.path.exists(f) else os.path.join(ANIM, f)
            if p not in images:
                images[p] = strip_lines(Image.open(p).convert("RGB"), bg_of(a.id)[1])
            if (p, nb) not in bands:
                bands[p, nb] = row_bands(images[p], nb)
    for size, feet in SIZES:
        idle_path = idle_sheet(a.id, size)
        idle = Image.open(idle_path).convert("RGBA")
        sheet = Image.new("RGBA", (4 * size, 5 * size), (0, 0, 0, 0))
        log = [f"{a.id}-{a.anim}-{size}"]
        for v, segs in enumerate(views):
            icell = idle.crop((v * size, 0, (v + 1) * size, size))
            ic, iw = core(icell, size), body_width(icell, size)
            frames, notes = [], []  # frames: (cell, source offset in output pixels)
            die = a.anim == "die"
            for f, band, nb, count, picks, ref in segs:
                p = f if os.path.isabs(f) or os.path.exists(f) else os.path.join(ANIM, f)
                y0, y1 = bands[p, nb][band]
                im = images[p]
                if die and ref is None:
                    ref = int(picks[0].lstrip("m"))
                figs, used, offs = pix_source(im.crop((0, y0, im.width, y1)), count, ic[1] - ic[0], size, feet, idle_path,
                                              a.add_colors, ref if die else None)
                notes.append(f"{os.path.basename(p)} row {band}/{nb} step {used:.2f}")
                for t in picks:
                    i = int(t.lstrip("m"))
                    c, o = figs[i], offs[i] / used
                    frames.append((c.transpose(Image.FLIP_LEFT_RIGHT), -o) if t.startswith("m") else (c, o))
            if len(frames) != 4:
                raise SystemExit(f"view {v}: {len(frames)} frames, want 4")
            fn = []
            wr = (body_width(frames[0][0], size) if die else statistics.median(body_width(f, size) for f, _ in frames)) / iw
            dx0 = o0 = None
            for col, (cell, off) in enumerate(frames):
                c = core(cell, size)
                if not die or col == 0:
                    dx, dy = round(ic[2] - c[2]), feet - c[1]
                    if die:
                        dx0, o0 = dx, off
                else:
                    # keep the position relative to frame 1 as drawn; the lowest part on the feet row
                    b = bbox(cell)
                    dx = round(dx0 + off - o0)
                    bot = c[1] if c else b[3] - 1
                    dy = feet - bot
                    if b[2] - b[0] <= size:  # pull a body that fits back into the cell
                        if b[0] + dx < 0:
                            dx = -b[0]
                        elif b[2] + dx > size:
                            dx = size - b[2]
                    else:
                        dx = round(size / 2 - (b[0] + b[2]) / 2)
                moved, lost = shift(cell, dx, dy)
                sheet.alpha_composite(moved, (col * size, v * size))
                fn.append(f"h{c[1] - c[0] if c else 0} dx{dx:+d} dy{dy:+d}" + (f" CLIPPED {lost}px" if lost else ""))
            log.append(f" {VIEWS[v]} (idle core h {ic[1] - ic[0]}, width ratio {wr:.2f}"
                       + (" WIDER/NARROWER THAN IDLE" if abs(wr - 1) > 0.1 else "") + "): " + "; ".join(notes) + " | " + ", ".join(fn))
        out = pre + f"-{size}.png"
        sheet.save(out)
        preview(sheet, idle, size, pre + f"-{size}-preview.png")
        print("\n".join(log))
        print("wrote", out)


def preview(sheet, idle, size, path):
    sc = 8 if size == 32 else 4
    pad = size // 4
    w = sheet.width + pad + size
    im = Image.new("RGBA", (w, sheet.height), (96, 92, 104, 255))
    im.alpha_composite(sheet, (0, 0))
    for v in range(5):
        im.alpha_composite(idle.crop((v * size, 0, (v + 1) * size, size)), (sheet.width + pad, v * size))
    # feet row guide
    big = im.resize((im.width * sc, im.height * sc), Image.NEAREST)
    px = big.load()
    for v in range(5):
        y = (v * size + size) * sc - 1
        for x in range(0, big.width, 3):
            if px[x, y][3] and px[x, y][:3] == (96, 92, 104):
                px[x, y] = (150, 146, 160, 255)
    big.save(path)


def main():
    ap = argparse.ArgumentParser()
    sub = ap.add_subparsers(dest="cmd", required=True)
    g = sub.add_parser("gen")
    g.add_argument("id")
    g.add_argument("anim", choices=["walk", "attack", "die"])
    g.add_argument("part", nargs="?", default="both", choices=["front", "back", "both"])
    g.add_argument("--suffix", default="")
    g.add_argument("--front-src", default=None, help="front raw to show in the back request")
    g.add_argument("--strict", action="store_true", help="back request: stress that every frame is seen from behind (retry)")
    g.add_argument("--print", action="store_true", help="only print the prompts")
    p = sub.add_parser("pix")
    p.add_argument("id")
    p.add_argument("anim", choices=["walk", "attack", "die"])
    p.add_argument("--front", default=None)
    p.add_argument("--back", default=None)
    p.add_argument("--view", action="append", help="V=FILE:BAND/NBANDS:COUNT:PICKS[;...] frames of view V "
                   "(0 front .. 4 back) from other rows/files; default views 0-2 = rows 0-2 of 5 of the front raw, "
                   "3-4 = rows 0-1 of 2 of the back raw, 4 figures each")
    p.add_argument("--add-colors", default="", help="extra palette colours R,G,B;R,G,B")
    a = ap.parse_args()
    if a.cmd == "gen":
        if a.print:
            print(prompt_front(a.id, a.anim), "\n\n----\n", prompt_back(a.id, a.anim))
        else:
            cmd_gen(a)
    else:
        cmd_pix(a)


if __name__ == "__main__":
    main()