#!/usr/bin/env python3
"""Rearrange a documented Mini Medieval sheet into the entities.png slot layout.

The packs' Animals/Monsters sheets share our state order (idle, sit, walk, action 1,
action 2, hit, dead) but pad groups with blank rows and show every frame twice (on
light and dark preview boxes). This writes assets/work/<Name>-aligned.png: the label
column, then per row the light frames of each state placed in the manifest's slot
columns, blank rows removed. That block pastes into the template in one go.

    python3 tools/realign_sheet.py assets/packs/Mini-Medieval-Documented-8x8/Animals.png
"""

import os
import sys
import tomllib
from PIL import Image

MANIFEST = "assets/tilesets/entities.toml"
HEADER = 16          # y of the first sprite row in the source
CELL = 8
LABEL_W = 91         # source label column (up to the first separator)
BLOCK_X = 96         # where the realigned frames start in the output


def is_dark(p):
    return p[3] > 0 and p[0] + p[1] + p[2] < 120


def separators(px, w, h):
    """x of the vertical rule lines (dark for most of the sheet's height)."""
    xs = [x for x in range(LABEL_W, w) if sum(is_dark(px[x, y]) for y in range(h)) > h * 0.9]
    # Rules are 2 px wide; keep one x per rule.
    return [x for i, x in enumerate(xs) if i == 0 or x != xs[i - 1] + 1]


def region_frames(px, x0, x1, h):
    """Cell columns of the light frames in a region: from the first content column
    (past the rule line, snapped to the 8 px grid) up to the first cell column where
    a preview box (a fully opaque cell in some row) starts."""
    def is_rule(x):
        return sum(is_dark(px[x, y]) for y in range(h)) > h * 0.9
    ys = range(HEADER, h)
    start = next((x for x in range(x0, x1)
                  if not is_rule(x) and any(px[x, y][3] > 0 for y in ys)), None)
    if start is None:
        return None, 0
    start -= start % CELL
    n = 0
    x = start
    while x + CELL <= x1:
        boxed = any(all(px[x + dx, y + dy][3] > 0 for dx in range(CELL) for dy in range(CELL))
                    for y in range(HEADER, h - CELL + 1, CELL))
        if boxed:
            break
        n += 1
        x += CELL
    return start, n


def main(path):
    m = tomllib.load(open(MANIFEST, "rb"))
    slots = list(m["layout"]["states"].values())      # [(col, max)] in state order
    total_cols = max(c + n for c, n in slots)

    img = Image.open(path).convert("RGBA")
    px = img.load()
    w, h = img.size

    # Non-empty source rows (labels or sprites, ignoring rules and boxes).
    def band_empty(y0):
        for y in range(y0, y0 + CELL):
            for x in range(w):
                p = px[x, y]
                if p[3] == 0 or (p[:3] == (18, 14, 35)) or (is_dark(p) and x > LABEL_W - 4):
                    continue
                return False
        return True
    src_rows = [y for y in range(HEADER, h - CELL + 1, CELL) if not band_empty(y)]

    seps = separators(px, w, h)
    regions = list(zip(seps, seps[1:] + [w]))
    if len(regions) < len(slots):
        sys.exit(f"{path}: found {len(regions)} state columns, manifest has {len(slots)}")

    out = Image.new("RGBA", (BLOCK_X + total_cols * CELL, HEADER + len(src_rows) * CELL), (0, 0, 0, 0))
    out.paste(img.crop((0, 0, LABEL_W, HEADER)), (0, 0))
    for (x0, x1), (slot_col, slot_max) in zip(regions, slots):
        start, n = region_frames(px, x0, x1, h)
        n = min(n, slot_max)
        for i, sy in enumerate(src_rows):
            dy = HEADER + i * CELL
            out.paste(img.crop((0, sy, LABEL_W, sy + CELL)), (0, dy))
            if n:
                frames = img.crop((start, sy, start + n * CELL, sy + CELL))
                out.paste(frames, (BLOCK_X + slot_col * CELL, dy), frames)

    name = os.path.splitext(os.path.basename(path))[0]
    dst = f"assets/work/{name}-aligned.png"
    out.save(dst)
    print(f"wrote {dst}: {len(src_rows)} rows; paste the block ({BLOCK_X}, {HEADER})–"
          f"({BLOCK_X + total_cols * CELL}, {out.size[1]}) at (64, 16 + 8·row)")


if __name__ == "__main__":
    for p in sys.argv[1:]:
        main(p)