#!/usr/bin/env python3
from __future__ import annotations

import argparse
from collections import deque
from pathlib import Path

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


ROWS = ["down", "up", "right", "left"]


def resize_rgba_alpha_aware(image: Image.Image, size: tuple[int, int]) -> Image.Image:
    rgba = image.convert("RGBA")
    arr = np.asarray(rgba).astype(np.float32)
    alpha = arr[:, :, 3:4] / 255.0
    premultiplied = arr[:, :, :3] * alpha
    premul_image = Image.fromarray(np.clip(premultiplied, 0, 255).astype(np.uint8), "RGB")
    alpha_image = Image.fromarray(arr[:, :, 3].astype(np.uint8), "L")
    resized_premul = np.asarray(premul_image.resize(size, Image.Resampling.LANCZOS)).astype(np.float32)
    resized_alpha = np.asarray(alpha_image.resize(size, Image.Resampling.LANCZOS)).astype(np.float32)
    alpha_fraction = resized_alpha[:, :, None] / 255.0
    rgb = np.zeros_like(resized_premul)
    np.divide(resized_premul, alpha_fraction, out=rgb, where=alpha_fraction > 0.001)
    output = np.dstack([np.clip(rgb, 0, 255), resized_alpha])
    return Image.fromarray(np.clip(output, 0, 255).astype(np.uint8), "RGBA")


def largest_component_mask(alpha: np.ndarray, threshold: int = 8) -> np.ndarray:
    foreground = alpha > threshold
    visited = np.zeros(foreground.shape, dtype=bool)
    best: list[tuple[int, int]] = []
    height, width = foreground.shape
    for y in range(height):
        for x in range(width):
            if not foreground[y, x] or visited[y, x]:
                continue
            component: list[tuple[int, int]] = []
            queue: deque[tuple[int, int]] = deque([(y, x)])
            visited[y, x] = True
            while queue:
                cy, cx = queue.popleft()
                component.append((cy, cx))
                for ny, nx in ((cy - 1, cx), (cy + 1, cx), (cy, cx - 1), (cy, cx + 1)):
                    if 0 <= ny < height and 0 <= nx < width and foreground[ny, nx] and not visited[ny, nx]:
                        visited[ny, nx] = True
                        queue.append((ny, nx))
            if len(component) > len(best):
                best = component
    mask = np.zeros(foreground.shape, dtype=bool)
    for y, x in best:
        mask[y, x] = True
    return mask


def keep_largest_component(cell: Image.Image, dilate: int) -> Image.Image:
    arr = np.asarray(cell.convert("RGBA")).copy()
    keep = largest_component_mask(arr[:, :, 3], threshold=10)
    for _ in range(dilate):
        expanded = keep.copy()
        expanded[:-1, :] |= keep[1:, :]
        expanded[1:, :] |= keep[:-1, :]
        expanded[:, :-1] |= keep[:, 1:]
        expanded[:, 1:] |= keep[:, :-1]
        keep = expanded
    arr[:, :, 3] = np.where(keep, arr[:, :, 3], 0).astype(np.uint8)
    return Image.fromarray(arr, "RGBA")


def remove_magenta_contamination(cell: Image.Image) -> Image.Image:
    arr = np.asarray(cell.convert("RGBA")).copy()
    rgb = arr[:, :, :3].astype(np.int16)
    alpha = arr[:, :, 3]
    magenta = (alpha > 0) & (rgb[:, :, 0] > 120) & (rgb[:, :, 1] < 80) & (rgb[:, :, 2] > 120)
    if magenta.any():
        arr[:, :, 3] = np.where(magenta, 0, alpha).astype(np.uint8)
    return Image.fromarray(arr, "RGBA")


def fit_cell(
    cell: Image.Image,
    *,
    frame_size: int,
    target_body_height: int,
    target_foot_y: int,
    component_dilate: int,
) -> Image.Image:
    cell = remove_magenta_contamination(keep_largest_component(cell, component_dilate))
    bbox = cell.getchannel("A").getbbox()
    if bbox is None:
        return Image.new("RGBA", (frame_size, frame_size), (0, 0, 0, 0))
    trimmed = cell.crop(bbox)
    scale = min(target_body_height / trimmed.height, (frame_size - 8) / trimmed.width, 1.0)
    next_size = (max(1, round(trimmed.width * scale)), max(1, round(trimmed.height * scale)))
    resized = resize_rgba_alpha_aware(trimmed, next_size)
    output = Image.new("RGBA", (frame_size, frame_size), (0, 0, 0, 0))
    x = (frame_size - resized.width) // 2
    y = target_foot_y - resized.height
    y = max(0, min(frame_size - resized.height, y))
    output.alpha_composite(resized, (x, y))
    return output


def extract_row(path: Path, args: argparse.Namespace) -> list[Image.Image]:
    image = Image.open(path).convert("RGBA")
    width, height = image.size
    x_edges = [round(i * width / args.columns) for i in range(args.columns + 1)]
    frames: list[Image.Image] = []
    for column in range(args.columns):
        cell = image.crop((x_edges[column], 0, x_edges[column + 1], height))
        frames.append(
            fit_cell(
                cell,
                frame_size=args.frame_size,
                target_body_height=args.target_body_height,
                target_foot_y=args.target_foot_y,
                component_dilate=args.component_dilate,
            )
        )
    return align_frames_to_neutral_upper_body(frames)


def shift_frame_horizontally(frame: Image.Image, offset: int) -> Image.Image:
    if offset == 0:
        return frame
    shifted = Image.new("RGBA", frame.size, (0, 0, 0, 0))
    shifted.alpha_composite(frame, (offset, 0))
    return shifted


def align_frames_to_neutral_upper_body(frames: list[Image.Image], max_shift: int = 8) -> list[Image.Image]:
    """Undo per-frame recentering caused by wider walking-leg silhouettes."""
    if not frames:
        return frames
    upper_end = round(frames[0].height * 0.60)
    neutral = np.asarray(frames[0].convert("RGBA"))[:upper_end, :, 3] > 20
    aligned = [frames[0]]
    for frame in frames[1:]:
        candidate = np.asarray(frame.convert("RGBA"))[:upper_end, :, 3] > 20
        best_score = float("inf")
        best_shift = 0
        for offset in range(-max_shift, max_shift + 1):
            shifted = np.zeros_like(candidate)
            if offset < 0:
                shifted[:, :offset] = candidate[:, -offset:]
            elif offset > 0:
                shifted[:, offset:] = candidate[:, :-offset]
            else:
                shifted = candidate
            union = np.logical_or(neutral, shifted)
            score = float(np.logical_xor(neutral, shifted).sum() / union.sum()) if union.any() else 0.0
            if score < best_score:
                best_score = score
                best_shift = offset
        aligned.append(shift_frame_horizontally(frame, best_shift if abs(best_shift) >= 2 else 0))
    return aligned


def checkerboard(width: int, height: int, tile: int = 8) -> Image.Image:
    image = Image.new("RGBA", (width, height), (226, 226, 226, 255))
    draw = ImageDraw.Draw(image)
    for y in range(0, height, tile):
        for x in range(0, width, tile):
            if (x // tile + y // tile) % 2 == 0:
                draw.rectangle((x, y, x + tile - 1, y + tile - 1), fill=(248, 248, 248, 255))
    return image


def save_review(sheet: Image.Image, path: Path, *, frame_size: int, columns: int) -> None:
    scale = 2
    label_width = 68
    strip_height = frame_size * scale
    canvas = Image.new("RGB", (label_width + sheet.width * scale, strip_height * 4), (246, 247, 250))
    draw = ImageDraw.Draw(canvas)
    font = ImageFont.load_default()
    for row, name in enumerate(ROWS):
        y = row * strip_height
        draw.text((8, y + strip_height // 2 - 5), name, fill=(28, 32, 36), font=font)
        row_sheet = sheet.crop((0, row * frame_size, sheet.width, (row + 1) * frame_size))
        row_preview = row_sheet.resize((sheet.width * scale, strip_height), Image.Resampling.NEAREST)
        bg = checkerboard(row_preview.width, row_preview.height, 16)
        bg.alpha_composite(row_preview)
        canvas.paste(bg.convert("RGB"), (label_width, y))
        for index in range(columns + 1):
            x = label_width + index * frame_size * scale
            draw.line((x, y, x, y + strip_height), fill=(205, 60, 60), width=1)
    path.parent.mkdir(parents=True, exist_ok=True)
    canvas.save(path)


def save_feet_zoom(sheet: Image.Image, path: Path, *, frame_size: int, columns: int) -> None:
    scale = 4
    crop_y0 = round(frame_size * 0.49)
    crop_h = frame_size - crop_y0
    label_w = 64
    label_h = 22
    cell_w = frame_size * scale
    cell_h = crop_h * scale
    canvas = Image.new("RGB", (label_w + cell_w * columns, label_h + cell_h * 4), (246, 247, 250))
    draw = ImageDraw.Draw(canvas)
    font = ImageFont.load_default()
    for column in range(columns):
        draw.text((label_w + column * cell_w + 4, 5), f"F{column + 1}", fill=(72, 76, 82), font=font)
    for row, row_name in enumerate(ROWS):
        y = label_h + row * cell_h
        draw.text((6, y + cell_h // 2 - 5), row_name, fill=(28, 32, 36), font=font)
        for column in range(columns):
            cell = sheet.crop(
                (
                    column * frame_size,
                    row * frame_size + crop_y0,
                    (column + 1) * frame_size,
                    (row + 1) * frame_size,
                )
            ).resize((cell_w, cell_h), Image.Resampling.NEAREST)
            bg = checkerboard(cell_w, cell_h, 20)
            bg.alpha_composite(cell)
            x = label_w + column * cell_w
            canvas.paste(bg.convert("RGB"), (x, y))
            draw.rectangle((x, y, x + cell_w - 1, y + cell_h - 1), outline=(188, 194, 202), width=1)
    path.parent.mkdir(parents=True, exist_ok=True)
    canvas.save(path)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Assemble four WhaleTown direction strips into an 8x4 spritesheet.")
    parser.add_argument("--down", type=Path, required=True)
    parser.add_argument("--up", type=Path, required=True)
    parser.add_argument("--right", type=Path, required=True)
    parser.add_argument("--left", type=Path, required=True)
    parser.add_argument("--name", required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--frame-size", type=int, default=160)
    parser.add_argument("--columns", type=int, default=8)
    parser.add_argument("--target-body-height", type=int, default=116)
    parser.add_argument("--target-foot-y", type=int, default=137)
    parser.add_argument("--component-dilate", type=int, default=2)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    inputs = {"down": args.down, "up": args.up, "right": args.right, "left": args.left}
    for name, path in inputs.items():
        if not path.exists():
            raise FileNotFoundError(f"{name} strip not found: {path}")

    sheet = Image.new("RGBA", (args.frame_size * args.columns, args.frame_size * 4), (0, 0, 0, 0))
    for row, direction in enumerate(ROWS):
        frames = extract_row(inputs[direction], args)
        for column, frame in enumerate(frames):
            sheet.alpha_composite(frame, (column * args.frame_size, row * args.frame_size))

    processed_dir = args.output_dir / "processed"
    review_dir = args.output_dir / "review"
    processed_dir.mkdir(parents=True, exist_ok=True)
    review_dir.mkdir(parents=True, exist_ok=True)

    sheet_path = processed_dir / f"{args.name}_spritesheet.png"
    review_path = review_dir / f"{args.name}_review.png"
    feet_path = review_dir / f"{args.name}_feet_zoom.png"

    sheet.save(sheet_path)
    save_review(sheet, review_path, frame_size=args.frame_size, columns=args.columns)
    save_feet_zoom(sheet, feet_path, frame_size=args.frame_size, columns=args.columns)

    print(sheet_path)
    print(review_path)
    print(feet_path)


if __name__ == "__main__":
    main()