from __future__ import annotations

from pathlib import Path

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


ROOT = Path(__file__).resolve().parents[1]
WORK_DIR = ROOT / "tmp" / "liu-bei-handpaint-sample"
UNIT_DIR = ROOT / "src" / "assets" / "images" / "units"
DOCS_DIR = ROOT / "docs"

FRAME = 313
DIRECTIONS = ("south", "east", "north", "west")
BASE_FRAMES_PER_DIRECTION = 16
ACTION_COUNTS = {
    "attack": 10,
    "strategy": 8,
    "item": 8,
    "hurt": 4,
    "celebrate": 6,
}
ACTION_OFFSETS = {
    "attack": 0,
    "strategy": 10,
    "item": 18,
    "hurt": 26,
    "celebrate": 30,
}
ACTION_FRAMES_PER_DIRECTION = sum(ACTION_COUNTS.values())


def keyed_rgba(image: Image.Image) -> Image.Image:
    rgba = image.convert("RGBA")
    arr = np.array(rgba)
    rgb = arr[:, :, :3].astype(np.int16)
    r, g, b = rgb[:, :, 0], rgb[:, :, 1], rgb[:, :, 2]
    key = (r > 185) & (b > 185) & (g < 125)
    near_key = (r > 145) & (b > 135) & (g < 150) & ((r + b - g * 2) > 180)
    alpha = np.where(key | near_key, 0, 255).astype(np.uint8)

    body = alpha > 0
    magenta_fringe = body & (r > 155) & (b > 135) & (g < 130)
    arr[:, :, 0] = np.where(magenta_fringe, np.minimum(arr[:, :, 0], 70), arr[:, :, 0])
    arr[:, :, 2] = np.where(magenta_fringe, np.minimum(arr[:, :, 2], 70), arr[:, :, 2])
    arr[:, :, 3] = alpha
    return Image.fromarray(arr, "RGBA")


def crop_subject(image: Image.Image) -> Image.Image:
    rgba = keyed_rgba(image)
    rgba = keep_subject_components(rgba)
    alpha = np.array(rgba.getchannel("A"))
    ys, xs = np.where(alpha > 0)
    if len(xs) == 0:
        return Image.new("RGBA", (FRAME, FRAME), (0, 0, 0, 0))
    pad = 8
    left = max(int(xs.min()) - pad, 0)
    top = max(int(ys.min()) - pad, 0)
    right = min(int(xs.max()) + pad + 1, rgba.width)
    bottom = min(int(ys.max()) + pad + 1, rgba.height)
    return rgba.crop((left, top, right, bottom))


def keep_subject_components(image: Image.Image) -> Image.Image:
    arr = np.array(image.convert("RGBA"))
    alpha = arr[:, :, 3] > 0
    height, width = alpha.shape
    visited = np.zeros_like(alpha, dtype=bool)
    components: list[tuple[int, int, int, int, int, float, float]] = []

    for start_y in range(height):
        for start_x in range(width):
            if visited[start_y, start_x] or not alpha[start_y, start_x]:
                continue

            stack = [(start_x, start_y)]
            visited[start_y, start_x] = True
            xs: list[int] = []
            ys: list[int] = []
            while stack:
                x, y = stack.pop()
                xs.append(x)
                ys.append(y)
                for ny in (y - 1, y, y + 1):
                    if ny < 0 or ny >= height:
                        continue
                    for nx in (x - 1, x, x + 1):
                        if nx < 0 or nx >= width or visited[ny, nx] or not alpha[ny, nx]:
                            continue
                        visited[ny, nx] = True
                        stack.append((nx, ny))

            area = len(xs)
            if area < 24:
                continue
            left, right = min(xs), max(xs)
            top, bottom = min(ys), max(ys)
            components.append((area, left, top, right, bottom, (left + right) / 2, (top + bottom) / 2))

    if not components:
        return image

    image_center_x = width / 2
    image_center_y = height * 0.58
    anchor = max(
        components,
        key=lambda component: component[0]
        - abs(component[5] - image_center_x) * 3
        - abs(component[6] - image_center_y) * 1.4,
    )
    anchor_area, left, top, right, bottom, anchor_cx, anchor_cy = anchor
    expanded = (left - 70, top - 80, right + 70, bottom + 65)

    keep = np.zeros_like(alpha, dtype=bool)
    for component in components:
        area, c_left, c_top, c_right, c_bottom, cx, cy = component
        touches_cell_edge = c_left <= 2 or c_right >= width - 3 or c_top <= 2 or c_bottom >= height - 3
        if component != anchor and touches_cell_edge and area < anchor_area * 0.55:
            continue
        near_anchor = (
            c_right >= expanded[0]
            and c_left <= expanded[2]
            and c_bottom >= expanded[1]
            and c_top <= expanded[3]
        )
        central = 0.04 * width <= cx <= 0.96 * width
        substantial = area >= max(55, anchor_area * 0.025)
        if component == anchor or (near_anchor and central and substantial):
            keep[c_top : c_bottom + 1, c_left : c_right + 1] |= alpha[c_top : c_bottom + 1, c_left : c_right + 1]

    arr[:, :, 3] = np.where(keep, arr[:, :, 3], 0).astype(np.uint8)
    return Image.fromarray(arr, "RGBA")


def quantize_alpha(image: Image.Image) -> Image.Image:
    rgba = image.convert("RGBA")
    arr = np.array(rgba)
    arr[:, :, 3] = np.where(arr[:, :, 3] > 24, 255, 0).astype(np.uint8)
    return Image.fromarray(arr, "RGBA")


def fit_subject(
    source: Image.Image,
    max_width: int = 286,
    max_height: int = 300,
    bottom: int = 306,
    stretch_x: float = 1.18,
) -> Image.Image:
    subject = crop_subject(source)
    if subject.width <= 1 or subject.height <= 1:
        return Image.new("RGBA", (FRAME, FRAME), (0, 0, 0, 0))

    scale = min(max_width / subject.width, max_height / subject.height)
    size = (max(1, round(subject.width * scale)), max(1, round(subject.height * scale)))
    subject = subject.resize(size, Image.Resampling.LANCZOS)
    stretched_width = min(max_width, max(1, round(subject.width * stretch_x)))
    if stretched_width != subject.width:
        subject = subject.resize((stretched_width, subject.height), Image.Resampling.LANCZOS)
    subject = quantize_alpha(subject)

    frame = Image.new("RGBA", (FRAME, FRAME), (0, 0, 0, 0))
    x = (FRAME - subject.width) // 2
    y = bottom - subject.height
    frame.alpha_composite(subject, (x, y))
    frame = quantize_alpha(frame)
    return quantize_alpha(keep_subject_components(frame))


def crop_grid(image: Image.Image, row: int, col: int, rows: int, cols: int, x_pad: int = 0, y_pad: int = 0) -> Image.Image:
    cell_w = image.width / cols
    cell_h = image.height / rows
    left = max(0, int(round(col * cell_w)) - x_pad)
    top = max(0, int(round(row * cell_h)) - y_pad)
    right = min(image.width, int(round((col + 1) * cell_w)) + x_pad)
    bottom = min(image.height, int(round((row + 1) * cell_h)) + y_pad)
    return image.crop((left, top, right, bottom))


def crop_center(image: Image.Image, center_x: int, center_y: int, width: int, height: int) -> Image.Image:
    left = max(0, center_x - width // 2)
    top = max(0, center_y - height // 2)
    right = min(image.width, center_x + width // 2)
    bottom = min(image.height, center_y + height // 2)
    return image.crop((left, top, right, bottom))


def build_base_sheet(base_source: Image.Image) -> Image.Image:
    rows: list[list[Image.Image]] = []
    for row in range(4):
        frames = [fit_subject(crop_grid(base_source, row, col, 4, 8, 0, 2)) for col in range(8)]
        idle = [frames[0].copy() for _ in range(8)]
        walk = frames
        rows.append(idle + walk)

    sheet = Image.new("RGBA", (FRAME * BASE_FRAMES_PER_DIRECTION, FRAME * len(DIRECTIONS)), (0, 0, 0, 0))
    for row, frames in enumerate(rows):
        for col, frame in enumerate(frames):
            sheet.alpha_composite(frame, (col * FRAME, row * FRAME))
    return sheet


def build_action_source_frames(action_source: Image.Image) -> dict[str, list[Image.Image]]:
    attack_frames = [
        fit_subject(crop_grid(action_source, 0, col, 3, 10, 10, 28), 300, 300, 306)
        for col in range(10)
    ]

    mid = [fit_subject(crop_grid(action_source, 1, col, 3, 5, 18, 26), 286, 300, 306) for col in range(5)]
    strategy_seed = mid[:3]
    item_seed = mid[3:]

    lower_raw = [
        crop_center(action_source, 470, 805, 330, 310),
        crop_center(action_source, 790, 805, 330, 310),
        crop_center(action_source, 1100, 805, 330, 310),
    ]
    lower = [fit_subject(frame, 286, 300, 306) for frame in lower_raw]

    return {
        "attack": attack_frames,
        "strategy": [strategy_seed[index % len(strategy_seed)] for index in (0, 1, 2, 1, 0, 1, 2, 1)],
        "item": [item_seed[index % len(item_seed)] for index in range(8)],
        "hurt": [lower[0] for _ in range(4)],
        "celebrate": [lower[1], lower[2], lower[1], lower[2], lower[1], lower[2]],
    }


def direction_frame(frame: Image.Image, direction: str) -> Image.Image:
    if direction == "west":
        return frame.transpose(Image.Transpose.FLIP_LEFT_RIGHT)
    return frame


def build_action_sheet(action_source: Image.Image) -> Image.Image:
    source_frames = build_action_source_frames(action_source)
    sheet = Image.new("RGBA", (FRAME * ACTION_FRAMES_PER_DIRECTION, FRAME * len(DIRECTIONS)), (0, 0, 0, 0))
    for row, direction in enumerate(DIRECTIONS):
        for action, frames in source_frames.items():
            offset = ACTION_OFFSETS[action]
            for index, frame in enumerate(frames):
                sheet.alpha_composite(direction_frame(frame, direction), ((offset + index) * FRAME, row * FRAME))
    return sheet


def representative_base_frames(base_sheet: Image.Image) -> list[Image.Image]:
    picks = [
        ("south idle", 0, 0),
        ("south walk", 0, 10),
        ("east walk", 1, 10),
        ("north walk", 2, 10),
        ("west walk", 3, 10),
    ]
    return [base_sheet.crop((col * FRAME, row * FRAME, (col + 1) * FRAME, (row + 1) * FRAME)) for _, row, col in picks]


def representative_action_frames(action_sheet: Image.Image) -> list[Image.Image]:
    picks = [
        ("attack 1", 1, 0),
        ("attack 4", 1, 3),
        ("attack 7", 1, 6),
        ("strategy", 0, ACTION_OFFSETS["strategy"] + 2),
        ("item", 0, ACTION_OFFSETS["item"] + 1),
        ("hurt", 0, ACTION_OFFSETS["hurt"]),
        ("celebrate", 0, ACTION_OFFSETS["celebrate"] + 1),
    ]
    return [action_sheet.crop((col * FRAME, row * FRAME, (col + 1) * FRAME, (row + 1) * FRAME)) for _, row, col in picks]


def checker_background(size: tuple[int, int], tile: int = 16) -> Image.Image:
    image = Image.new("RGBA", size, (67, 88, 54, 255))
    draw = ImageDraw.Draw(image)
    for y in range(0, size[1], tile):
        for x in range(0, size[0], tile):
            color = (77, 99, 59, 255) if ((x // tile) + (y // tile)) % 2 == 0 else (58, 76, 48, 255)
            draw.rectangle((x, y, x + tile - 1, y + tile - 1), fill=color)
    return image


def draw_frame_on_bg(frame: Image.Image, size: int) -> Image.Image:
    bg = checker_background((size, size), 8)
    scaled = frame.resize((size, size), Image.Resampling.LANCZOS)
    bg.alpha_composite(scaled, (0, 0))
    return bg


def save_contact_sheet(base_sheet: Image.Image, action_sheet: Image.Image) -> None:
    thumbs = representative_base_frames(base_sheet) + representative_action_frames(action_sheet)
    labels = [
        "idle",
        "walk S",
        "walk E",
        "walk N",
        "walk W",
        "atk 1",
        "atk 4",
        "atk 7",
        "cmd",
        "item",
        "hurt",
        "win",
    ]
    thumb_size = 132
    pad = 18
    label_h = 24
    cols = 6
    rows = 2
    out = Image.new("RGBA", (pad + cols * (thumb_size + pad), pad + rows * (thumb_size + label_h + pad)), (20, 24, 22, 255))
    draw = ImageDraw.Draw(out)
    font = ImageFont.load_default()
    for index, frame in enumerate(thumbs):
        x = pad + (index % cols) * (thumb_size + pad)
        y = pad + (index // cols) * (thumb_size + label_h + pad)
        tile = draw_frame_on_bg(frame, thumb_size)
        out.alpha_composite(tile, (x, y))
        draw.text((x, y + thumb_size + 5), labels[index], fill=(238, 224, 178, 255), font=font)
    out.convert("RGB").save(DOCS_DIR / "liu-bei-handpaint-sample-contact.png", optimize=True)


def save_before_after(base_sheet: Image.Image) -> None:
    before = Image.open(WORK_DIR / "before-unit-liu-bei.png").convert("RGBA").crop((0, 0, FRAME, FRAME))
    after = base_sheet.crop((0, 0, FRAME, FRAME))
    out = Image.new("RGBA", (820, 360), (18, 21, 20, 255))
    draw = ImageDraw.Draw(out)
    font = ImageFont.load_default()
    draw.text((30, 18), "Before / After at 313px", fill=(238, 224, 178, 255), font=font)
    out.alpha_composite(draw_frame_on_bg(before, 160), (30, 48))
    out.alpha_composite(draw_frame_on_bg(after, 160), (220, 48))
    draw.text((30, 218), "Before / After at battle scale", fill=(238, 224, 178, 255), font=font)
    out.alpha_composite(draw_frame_on_bg(before, 68), (30, 252))
    out.alpha_composite(draw_frame_on_bg(after, 68), (128, 252))
    draw.text((30, 330), "left: previous checked-in Liu Bei, right: new hand-painted sample", fill=(172, 185, 166, 255), font=font)
    out.convert("RGB").save(DOCS_DIR / "liu-bei-handpaint-sample-before-after.png", optimize=True)


def save_animation_gif(base_sheet: Image.Image, action_sheet: Image.Image) -> None:
    frames: list[Image.Image] = []
    for index in range(10):
        canvas = checker_background((540, 190), 10)
        idle = base_sheet.crop(((index % 8) * FRAME, 0, ((index % 8) + 1) * FRAME, FRAME))
        walk = base_sheet.crop(((8 + index % 8) * FRAME, FRAME, (9 + index % 8) * FRAME, 2 * FRAME))
        attack = action_sheet.crop((index * FRAME, FRAME, (index + 1) * FRAME, 2 * FRAME))
        for x, frame in ((22, idle), (190, walk), (358, attack)):
            scaled = frame.resize((150, 150), Image.Resampling.LANCZOS)
            canvas.alpha_composite(scaled, (x, 25))
        frames.append(canvas.convert("P", palette=Image.Palette.ADAPTIVE, colors=128))
    frames[0].save(
        DOCS_DIR / "liu-bei-handpaint-sample-animation.gif",
        save_all=True,
        append_images=frames[1:],
        optimize=True,
        duration=110,
        loop=0,
    )


def validate_sheet(path: Path, expected_size: tuple[int, int]) -> None:
    image = Image.open(path).convert("RGBA")
    if image.size != expected_size:
        raise ValueError(f"{path.name}: expected {expected_size}, got {image.size}")
    alpha = np.array(image.getchannel("A"))
    partial = np.count_nonzero((alpha > 0) & (alpha < 255))
    if partial:
        raise ValueError(f"{path.name}: found {partial} partially transparent pixels")


def main() -> None:
    base_source = Image.open(WORK_DIR / "source-base-motion.png").convert("RGBA")
    action_source = Image.open(WORK_DIR / "source-action.png").convert("RGBA")

    base_sheet = build_base_sheet(base_source)
    action_sheet = build_action_sheet(action_source)

    base_path = UNIT_DIR / "unit-liu-bei.png"
    action_path = UNIT_DIR / "unit-liu-bei-actions.png"
    base_sheet.save(base_path, optimize=True)
    action_sheet.save(action_path, optimize=True)

    save_contact_sheet(base_sheet, action_sheet)
    save_before_after(base_sheet)
    save_animation_gif(base_sheet, action_sheet)

    validate_sheet(base_path, (FRAME * BASE_FRAMES_PER_DIRECTION, FRAME * len(DIRECTIONS)))
    validate_sheet(action_path, (FRAME * ACTION_FRAMES_PER_DIRECTION, FRAME * len(DIRECTIONS)))

    print(f"Wrote {base_path}")
    print(f"Wrote {action_path}")
    print(f"Wrote {DOCS_DIR / 'liu-bei-handpaint-sample-contact.png'}")
    print(f"Wrote {DOCS_DIR / 'liu-bei-handpaint-sample-before-after.png'}")
    print(f"Wrote {DOCS_DIR / 'liu-bei-handpaint-sample-animation.gif'}")


if __name__ == "__main__":
    main()