"""Offline normalization of approved waking-scene sprites; runtime is MASM."""
from pathlib import Path
import struct
import sys
from PIL import Image
from collections import deque

ROOT = Path(__file__).resolve().parents[2]
SRC = ROOT / "assets/sprites/arrival"
OUT = ROOT / "assets/animations/arrival"


def pack(name, frames):
    w, h = frames[0].size
    data = b"".join(f.tobytes("raw", "BGRA") for f in frames)
    (OUT / name).write_bytes(struct.pack("<4s4I", b"ANM1", w, h, len(frames), len(data)) + data)


def cells(name):
    sheet = Image.open(SRC / name).convert("RGBA")
    result = []
    for i in range(16):
        c, r = i % 4, i // 4
        cell = sheet.crop((c * sheet.width // 4, r * sheet.height // 4,
                           (c + 1) * sheet.width // 4, (r + 1) * sheet.height // 4))
        cell.putalpha(cell.getchannel("A").point(lambda a: 255 if a >= 160 else 0))
        result.append(cell)
    return result


def figures(name, expected=16):
    # Generator gutters are not necessarily equal grid cells. Recover whole
    # connected silhouettes before anchoring; otherwise a neighbouring head
    # fragment can become the foot anchor of a crouching pose.
    source = Image.open(SRC / name).convert("RGBA")
    source.putalpha(source.getchannel("A").point(lambda a: 255 if a >= 160 else 0))
    w, h = source.size
    mask = bytearray(source.getchannel("A").tobytes())
    blobs = []
    for start in range(len(mask)):
        if not mask[start]:
            continue
        queue = deque([start])
        mask[start] = 0
        pixels = []
        while queue:
            at = queue.popleft()
            pixels.append(at)
            x, y = at % w, at // w
            for yy in range(max(0, y - 1), min(h, y + 2)):
                for xx in range(max(0, x - 1), min(w, x + 2)):
                    neighbour = yy * w + xx
                    if mask[neighbour]:
                        mask[neighbour] = 0
                        queue.append(neighbour)
        if len(pixels) < 800:
            continue
        left = min(p % w for p in pixels)
        right = max(p % w for p in pixels) + 1
        top = min(p // w for p in pixels)
        bottom = max(p // w for p in pixels) + 1
        figure = source.crop((left, top, right, bottom))
        own_mask = bytearray(figure.width * figure.height)
        for p in pixels:
            own_mask[(p // w - top) * figure.width + p % w - left] = 255
        figure.putalpha(Image.frombytes("L", figure.size, bytes(own_mask)))
        blobs.append(((top + bottom) / 2, (left + right) / 2, figure))
    allowed = expected if isinstance(expected, tuple) else (expected,)
    assert len(blobs) in allowed, f"Unexpected figure count: {len(blobs)}"
    blobs.sort(key=lambda b: b[0])
    return [b[2] for r in range(len(blobs)//4) for b in sorted(blobs[r * 4:(r + 1) * 4], key=lambda b: b[1])]


def pack_actor():
    # Reject the atlas's shrinking departure rows. Calibrate the separate
    # in-place strip's source resolution once, then use one shared final scale.
    bodies = figures("wake-rocker-v2.png", (32,36))[:24]
    walks = figures("wake-walk-v2.png",8)
    units = max(c.height for c in bodies[19:24])/max(c.height for c in walks)
    bodies += [c.resize((round(c.width*units),round(c.height*units)),Image.Resampling.NEAREST) for c in walks]
    scale = min(176 / max(c.width for c in bodies), 304 / max(c.height for c in bodies))
    frames = []
    for body in bodies:
        body = body.resize((round(body.width * scale), round(body.height * scale)), Image.Resampling.NEAREST)
        frame = Image.new("RGBA", (192, 320))
        frame.paste(body, ((192 - body.width) // 2, 320 - body.height))
        frames.append(frame)
    pack("wake-rocker-v2.bin", frames)


def main():
    OUT.mkdir(parents=True, exist_ok=True)
    pack_actor()
    frames = []
    # Measured nozzle anchors in source-cell fractions, not a per-frame bbox fit.
    anchors = [(0.56, .25)] * 3 + [(.60, .30)] * 3 + [(.61, .29)] * 7 + [(.57, .31)] * 3
    for cell, (ax, ay) in zip(cells("wake-lighter-v1.png"), anchors):
        frame = Image.new("RGBA", (256, 256))
        sprite = cell.resize((230, 230), Image.Resampling.NEAREST)
        frame.paste(sprite, (128 - round(ax * 230), 64 - round(ay * 230)))
        frames.append(frame)
    pack("wake-lighter-v1.bin", frames)
    joints = []
    for frame in frames:
        alpha = frame.getchannel("A")
        bottom = alpha.getbbox()[3]-1
        xs = [x for x in range(256) if alpha.getpixel((x,bottom)) >= 160]
        joints.append(((min(xs)+max(xs))//2,bottom))
    joint_source = "cityLighterCuffX3D dd " + ",".join(str(x) for x,y in joints) + "\n"
    joint_source += "cityLighterCuffY3D dd " + ",".join(str(y) for x,y in joints) + "\n"
    (ROOT / "src_graphics/city_lighter_anchors.inc").write_text(joint_source,encoding="ascii")
    sleeve = Image.open(SRC / "wake-sleeve-v1.png").convert("RGBA")
    sleeve.putalpha(sleeve.getchannel("A").point(lambda a: 255 if a >= 160 else 0))
    sleeve = sleeve.crop(sleeve.getbbox()).resize((256,384),Image.Resampling.NEAREST)
    xs = [x for x in range(256) if sleeve.getchannel("A").getpixel((x,12))]
    anchor = (min(xs)+max(xs))//2
    aligned = Image.new("RGBA",(256,384))
    aligned.paste(sleeve,(128-anchor,0))
    pack("wake-sleeve-v1.bin",[aligned])
    materials = Image.open(SRC / "alley-materials-v1.png").convert("RGBA")
    textures = []
    for row in range(2):
        for col in range(2):
            tile = materials.crop((col*materials.width//2,row*materials.height//2,
                                   (col+1)*materials.width//2,(row+1)*materials.height//2))
            tile = tile.resize((128,128),Image.Resampling.NEAREST)
            tile.putalpha(255)
            textures.append(tile)
    pack("alley-materials-v1.bin",textures)
    legs = Image.open(SRC / "wake-knees-v1.png").convert("RGBA").resize((640, 240), Image.Resampling.NEAREST)
    legs.putalpha(legs.getchannel("A").point(lambda a: 255 if a >= 160 else 0))
    pack("wake-knees-v1.bin", [legs])
    print("Packed 32 grounded rocker poses, 16 cuff anchors, continuous sleeve, four alley materials and knees")


if __name__ == "__main__":
    if "--actor-only" in sys.argv:
        pack_actor()
    else:
        main()