#!/usr/bin/env python3
"""Build a stable raster-parts rig for Seoyeon.

The generated character sheet contains useful loose parts, but recomposing those
parts directly is fragile because each item is drawn at a different scale. This
script uses the approved assembled bust as the coordinate authority, creates a
faceless base from it, then adds small aligned raster parts for the eyes, brows,
nose, mouth, blink, and hair-sway overlays.

Output files are full-canvas 900x1125 PNG layers so the React renderer can stack
them without per-image layout math.
"""
from __future__ import annotations

from pathlib import Path
from typing import Iterable, Literal

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

ROOT = Path(__file__).resolve().parents[3]
SHEET = ROOT / "docs/avatar-art/seoyeon/character-sheet-v1.png"
PUB = ROOT / "apps/web/public/avatar/seoyeon"
NEUTRAL = PUB / "neutral.png"
OUT = PUB / "parts"
PREVIEW = ROOT / "docs/avatar-art/seoyeon/character-sheet-rig-preview-v2.png"
LEGACY_PREVIEW = ROOT / "docs/avatar-art/seoyeon/character-sheet-recompose-preview.png"
CANVAS = (900, 1125)

Box = tuple[int, int, int, int]
ShapeKind = Literal["ellipse", "round", "rect"]


def alpha_bbox(im: Image.Image, threshold: int = 10) -> Box | None:
    alpha = np.array(im.getchannel("A"))
    ys, xs = np.where(alpha > threshold)
    if len(xs) == 0:
        return None
    return int(xs.min()), int(ys.min()), int(xs.max() + 1), int(ys.max() + 1)


def median_skin(im: Image.Image, box: Box) -> tuple[int, int, int, int]:
    arr = np.array(im.crop(box).convert("RGBA"))
    rgb = arr[:, :, :3].astype(np.float32)
    alpha = arr[:, :, 3] > 180
    luma = rgb[:, :, 0] * 0.2126 + rgb[:, :, 1] * 0.7152 + rgb[:, :, 2] * 0.0722
    chroma = rgb.max(axis=2) - rgb.min(axis=2)
    # Keep face skin, reject hair/linework/shirt.
    mask = alpha & (luma > 145) & (chroma > 8)
    if int(mask.sum()) < 20:
        mask = alpha
    color = np.median(rgb[mask], axis=0)
    return int(color[0]), int(color[1]), int(color[2]), 255


def draw_shape(mask: Image.Image, box: Box, kind: ShapeKind, radius: int = 14) -> None:
    draw = ImageDraw.Draw(mask)
    if kind == "ellipse":
        draw.ellipse(box, fill=255)
    elif kind == "round":
        draw.rounded_rectangle(box, radius=radius, fill=255)
    else:
        draw.rectangle(box, fill=255)


def shape_mask(
    size: tuple[int, int],
    shapes: Iterable[tuple[Box, ShapeKind, int]],
    feather: float = 0,
) -> Image.Image:
    mask = Image.new("L", size, 0)
    for box, kind, radius in shapes:
        draw_shape(mask, box, kind, radius)
    if feather > 0:
        mask = mask.filter(ImageFilter.GaussianBlur(feather))
    return mask


def erase_features(neutral: Image.Image) -> Image.Image:
    base = neutral.convert("RGBA")
    patches = [
        ((318, 292, 410, 320), (418, 360, 476, 392), "round", 5),
        ((490, 292, 582, 320), (418, 360, 476, 392), "round", 5),
        ((306, 314, 412, 372), (418, 360, 476, 392), "ellipse", 5),
        ((488, 314, 594, 372), (418, 360, 476, 392), "ellipse", 5),
        ((412, 382, 490, 430), (392, 360, 512, 410), "ellipse", 5),
        ((382, 414, 518, 466), (388, 382, 512, 428), "ellipse", 6),
    ]

    for erase_box, sample_box, kind, feather in patches:
        color = median_skin(base, sample_box)
        fill = Image.new("RGBA", CANVAS, color)
        mask = Image.new("L", CANVAS, 0)
        draw_shape(mask, erase_box, kind, 18)
        mask = mask.filter(ImageFilter.GaussianBlur(feather))
        base = Image.composite(fill, base, mask)
    return base


def save_image(name: str, im: Image.Image) -> Image.Image:
    out = OUT / f"{name}.png"
    im.save(out, optimize=True)
    print(f"{name:18} bbox={alpha_bbox(im)}")
    return im


def save_masked_neutral(
    neutral: Image.Image,
    name: str,
    shapes: Iterable[tuple[Box, ShapeKind, int]],
    *,
    subtract_shapes: Iterable[tuple[Box, ShapeKind, int]] = (),
    feather: float = 0,
    mode: Literal["patch", "dark"] = "patch",
    dark_hi: float = 175,
    dark_lo: float = 80,
) -> Image.Image:
    src = neutral.convert("RGBA")
    mask_im = shape_mask(CANVAS, shapes, feather=0)
    if subtract_shapes:
        erase_im = shape_mask(CANVAS, subtract_shapes, feather=0)
        mask_arr = np.array(mask_im)
        mask_arr[np.array(erase_im) > 0] = 0
        mask_im = Image.fromarray(mask_arr, "L")
    if feather > 0:
        mask_im = mask_im.filter(ImageFilter.GaussianBlur(feather))
    mask = np.array(mask_im).astype(np.float32) / 255.0
    arr = np.array(src).astype(np.float32)
    alpha = arr[:, :, 3] * mask

    if mode == "dark":
        rgb = arr[:, :, :3]
        luma = rgb[:, :, 0] * 0.2126 + rgb[:, :, 1] * 0.7152 + rgb[:, :, 2] * 0.0722
        dark = np.clip((dark_hi - luma) / max(1.0, dark_hi - dark_lo), 0, 1)
        alpha *= dark

    arr[:, :, 3] = alpha
    arr[alpha < 1, :3] = 0
    return save_image(name, Image.fromarray(np.clip(arr, 0, 255).astype(np.uint8), "RGBA"))


def remove_sheet_bg(im: Image.Image, *, lo: float = 18.0, hi: float = 62.0) -> Image.Image:
    arr = np.array(im.convert("RGBA")).astype(np.float32)
    rgb = arr[:, :, :3]
    border = np.concatenate(
        [
            rgb[:5, :, :].reshape(-1, 3),
            rgb[-5:, :, :].reshape(-1, 3),
            rgb[:, :5, :].reshape(-1, 3),
            rgb[:, -5:, :].reshape(-1, 3),
        ],
        axis=0,
    )
    bg = np.median(border, axis=0)
    dist_bg = np.sqrt(((rgb - bg.reshape(1, 1, 3)) ** 2).sum(axis=2))
    dist_white = np.sqrt(((255.0 - rgb) ** 2).sum(axis=2))
    alpha_bg = np.clip((dist_bg - lo) * 255.0 / max(1.0, hi - lo), 0, 255)
    alpha_white = np.clip((dist_white - 15.0) * 255.0 / 50.0, 0, 255)
    alpha = np.minimum(alpha_bg, alpha_white)

    luma = rgb[:, :, 0] * 0.2126 + rgb[:, :, 1] * 0.7152 + rgb[:, :, 2] * 0.0722
    chroma = rgb.max(axis=2) - rgb.min(axis=2)
    keep = (luma < 205) | (chroma > 16)
    alpha = np.where(keep & (alpha > 35), np.maximum(alpha, 230), alpha)
    arr[:, :, 3] = np.minimum(alpha, arr[:, :, 3])
    arr[arr[:, :, 3] < 1, :3] = 0
    return Image.fromarray(np.clip(arr, 0, 255).astype(np.uint8), "RGBA")


def save_sheet_part(
    sheet: Image.Image,
    name: str,
    src_box: Box,
    dst_box: Box,
    *,
    matte_lo: float = 14.0,
    matte_hi: float = 54.0,
) -> Image.Image:
    part = remove_sheet_bg(sheet.crop(src_box), lo=matte_lo, hi=matte_hi)
    dst_w = dst_box[2] - dst_box[0]
    dst_h = dst_box[3] - dst_box[1]
    part = part.resize((dst_w, dst_h), Image.Resampling.LANCZOS)
    canvas = Image.new("RGBA", CANVAS, (0, 0, 0, 0))
    canvas.alpha_composite(part, (dst_box[0], dst_box[1]))
    return save_image(name, canvas)


def duplicate(src: str, dst: str) -> None:
    im = Image.open(OUT / f"{src}.png").convert("RGBA")
    save_image(dst, im)


def make_preview(parts: dict[str, Image.Image]) -> Image.Image:
    bg = Image.new("RGBA", CANVAS, (30, 39, 36, 255))
    order = [
        "base-faceless",
        "hair-left",
        "hair-right",
        "hair-bangs",
        "brow-neutral",
        "eyes-neutral",
        "nose-neutral",
        "mouth-neutral",
    ]
    for name in order:
        layer = parts.get(name)
        if layer is None:
            layer = Image.open(OUT / f"{name}.png").convert("RGBA")
        bg.alpha_composite(layer)

    # Add small comparison swatches for blink and mouth variants without touching
    # the app-visible neutral render.
    swatch = Image.new("RGBA", (360, 210), (30, 39, 36, 255))
    for i, name in enumerate(["eyelid-closed", "mouth-sad", "mouth-warm", "mouth-open"]):
        layer = Image.open(OUT / f"{name}.png").convert("RGBA")
        crop_box = alpha_bbox(layer) or (0, 0, 1, 1)
        crop = layer.crop(crop_box)
        crop.thumbnail((155, 85), Image.Resampling.LANCZOS)
        x = 18 + (i % 2) * 174
        y = 18 + (i // 2) * 98
        swatch.alpha_composite(crop, (x, y))
    bg.alpha_composite(swatch, (20, 895))
    return bg


def main() -> int:
    OUT.mkdir(parents=True, exist_ok=True)
    neutral = Image.open(NEUTRAL).convert("RGBA")
    sheet = Image.open(SHEET).convert("RGBA")

    if neutral.size != CANVAS:
        raise ValueError(f"Expected {CANVAS}, got {neutral.size} for {NEUTRAL}")

    parts: dict[str, Image.Image] = {}
    parts["base-faceless"] = save_image("base-faceless", erase_features(neutral))

    # Compatibility/static layers used by older manifests or future experiments.
    parts["face-faceless"] = save_masked_neutral(
        parts["base-faceless"],
        "face-faceless",
        [((250, 205, 650, 535), "ellipse", 40)],
        feather=3,
    )
    parts["forehead"] = save_masked_neutral(
        parts["base-faceless"],
        "forehead",
        [((340, 228, 560, 315), "round", 22)],
        feather=4,
    )
    parts["neck"] = save_masked_neutral(
        neutral,
        "neck",
        [((352, 470, 548, 618), "round", 28)],
        feather=2,
    )
    parts["shoulders"] = save_masked_neutral(
        neutral,
        "shoulders",
        [((160, 545, 740, 820), "round", 28)],
        feather=2,
    )

    parts["hair-back"] = save_masked_neutral(
        neutral,
        "hair-back",
        [((170, 80, 730, 520), "round", 80)],
        subtract_shapes=[
            ((262, 210, 638, 540), "ellipse", 72),
            ((292, 280, 608, 388), "round", 34),
        ],
        feather=2,
        mode="dark",
        dark_hi=172,
        dark_lo=70,
    )
    parts["hair-bangs"] = save_masked_neutral(
        neutral,
        "hair-bangs",
        [((326, 120, 574, 310), "round", 44)],
        subtract_shapes=[((292, 282, 608, 386), "round", 34)],
        feather=1.5,
        mode="dark",
        dark_hi=178,
        dark_lo=75,
    )
    parts["hair-left"] = save_masked_neutral(
        neutral,
        "hair-left",
        [
            ((168, 155, 318, 522), "round", 54),
            ((242, 430, 392, 522), "round", 42),
        ],
        subtract_shapes=[
            ((270, 228, 430, 538), "ellipse", 48),
            ((292, 282, 425, 390), "round", 32),
        ],
        feather=1.5,
        mode="dark",
        dark_hi=178,
        dark_lo=75,
    )
    parts["hair-right"] = save_masked_neutral(
        neutral,
        "hair-right",
        [
            ((582, 155, 732, 522), "round", 54),
            ((508, 430, 658, 522), "round", 42),
        ],
        subtract_shapes=[
            ((470, 228, 630, 538), "ellipse", 48),
            ((475, 282, 608, 390), "round", 32),
        ],
        feather=1.5,
        mode="dark",
        dark_hi=178,
        dark_lo=75,
    )

    # Ear files are kept as small static compatibility layers; the faceless base
    # already contains the ears in their approved position.
    parts["ear-left"] = save_masked_neutral(
        neutral,
        "ear-left",
        [((285, 318, 337, 392), "ellipse", 20)],
        feather=2,
    )
    parts["ear-right"] = save_masked_neutral(
        neutral,
        "ear-right",
        [((563, 318, 615, 392), "ellipse", 20)],
        feather=2,
    )

    parts["brow-neutral"] = save_sheet_part(
        sheet,
        "brow-neutral",
        (600, 424, 848, 482),
        (315, 272, 591, 310),
        matte_lo=10,
        matte_hi=42,
    )
    parts["eyes-neutral"] = save_sheet_part(
        sheet,
        "eyes-neutral",
        (594, 458, 856, 548),
        (311, 314, 599, 376),
        matte_lo=10,
        matte_hi=44,
    )
    parts["nose-neutral"] = save_masked_neutral(
        neutral,
        "nose-neutral",
        [((414, 385, 488, 426), "ellipse", 18)],
        feather=2,
    )
    parts["mouth-neutral"] = save_masked_neutral(
        neutral,
        "mouth-neutral",
        [((382, 414, 518, 468), "ellipse", 20)],
        feather=2,
    )

    parts["eyelid-closed"] = save_sheet_part(
        sheet,
        "eyelid-closed",
        (592, 552, 856, 620),
        (306, 315, 594, 369),
        matte_lo=11,
        matte_hi=46,
    )
    parts["mouth-sad"] = save_sheet_part(
        sheet,
        "mouth-sad",
        (888, 520, 1016, 586),
        (397, 414, 531, 468),
        matte_lo=10,
        matte_hi=44,
    )
    parts["mouth-warm"] = save_sheet_part(
        sheet,
        "mouth-warm",
        (888, 596, 1016, 662),
        (397, 414, 531, 468),
        matte_lo=10,
        matte_hi=44,
    )
    parts["mouth-open"] = save_sheet_part(
        sheet,
        "mouth-open",
        (888, 672, 1016, 752),
        (394, 407, 534, 475),
        matte_lo=10,
        matte_hi=44,
    )

    for variant in ["sad", "tired", "anxious", "warm", "startled"]:
        duplicate("brow-neutral", f"brow-{variant}")
        duplicate("eyes-neutral", f"eyes-{variant}")
    duplicate("mouth-sad", "mouth-tired")
    duplicate("mouth-sad", "mouth-anxious")
    duplicate("mouth-open", "mouth-startled")

    preview = make_preview(parts)
    preview.save(PREVIEW, optimize=True)
    preview.save(LEGACY_PREVIEW, optimize=True)
    print(f"preview {PREVIEW}")
    print(f"legacy-preview {LEGACY_PREVIEW}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())