#!/usr/bin/env python3
"""Extract a Live2D-style raster parts set from sheet-candidate-a.png.

This is intentionally offline-only. It writes to docs/avatar-art/seoyeon/live2d-v2
and does not publish anything to apps/web/public until the recomposition preview
is acceptable.
"""
from __future__ import annotations

from pathlib import Path
from typing import Iterable

import numpy as np
from PIL import Image, ImageDraw, ImageOps
from scipy import ndimage as ndi

HERE = Path(__file__).resolve().parent
SHEET = HERE / "sheet-candidate-a.png"
OUT = HERE / "parts"
CANVAS = (900, 1125)

Box = tuple[int, int, int, int]


def alpha_bbox(im: Image.Image, threshold: int = 8) -> 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 remove_white_bg(im: Image.Image, *, lo: float = 12.0, hi: float = 58.0) -> Image.Image:
    arr = np.array(im.convert("RGBA")).astype(np.float32)
    rgb = arr[:, :, :3]
    border = np.concatenate(
        [
            rgb[:6, :, :].reshape(-1, 3),
            rgb[-6:, :, :].reshape(-1, 3),
            rgb[:, :6, :].reshape(-1, 3),
            rgb[:, -6:, :].reshape(-1, 3),
        ],
        axis=0,
    )
    bg = np.median(border, axis=0)
    dist = np.sqrt(((rgb - bg.reshape(1, 1, 3)) ** 2).sum(axis=2))
    white_dist = np.sqrt(((255.0 - rgb) ** 2).sum(axis=2))
    alpha = np.minimum(
        np.clip((dist - lo) * 255.0 / max(1.0, hi - lo), 0, 255),
        np.clip((white_dist - 12.0) * 255.0 / 48.0, 0, 255),
    )

    luma = rgb[:, :, 0] * 0.2126 + rgb[:, :, 1] * 0.7152 + rgb[:, :, 2] * 0.0722
    chroma = rgb.max(axis=2) - rgb.min(axis=2)
    keep = (luma < 225) | (chroma > 14)
    alpha = np.where(keep & (alpha > 28), np.maximum(alpha, 235), 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 crop_part(sheet: Image.Image, src: Box, *, lo: float = 12, hi: float = 58) -> Image.Image:
    return remove_white_bg(sheet.crop(src), lo=lo, hi=hi)


def filter_components(im: Image.Image, *, min_area: int = 0, keep_largest: bool = False) -> Image.Image:
    if min_area <= 0 and not keep_largest:
        return im
    arr = np.array(im.convert("RGBA"))
    alpha = arr[:, :, 3] > 8
    labels, count = ndi.label(alpha)
    if count == 0:
        return im
    areas = np.bincount(labels.reshape(-1))
    areas[0] = 0
    if keep_largest:
        keep = labels == int(areas.argmax())
    else:
        keep_labels = np.where(areas >= min_area)[0]
        keep = np.isin(labels, keep_labels)
    arr[:, :, 3] = np.where(keep, arr[:, :, 3], 0)
    arr[arr[:, :, 3] == 0, :3] = 0
    return Image.fromarray(arr, "RGBA")


def place_exact(part: Image.Image, dst: Box) -> Image.Image:
    w, h = dst[2] - dst[0], dst[3] - dst[1]
    resized = part.resize((w, h), Image.Resampling.LANCZOS)
    canvas = Image.new("RGBA", CANVAS, (0, 0, 0, 0))
    canvas.alpha_composite(resized, (dst[0], dst[1]))
    return canvas


def save_part(
    sheet: Image.Image,
    name: str,
    src: Box,
    dst: Box,
    *,
    lo: float = 12,
    hi: float = 58,
    min_area: int = 0,
    keep_largest: bool = False,
    flip_x: bool = False,
) -> Image.Image:
    part = filter_components(crop_part(sheet, src, lo=lo, hi=hi), min_area=min_area, keep_largest=keep_largest)
    if flip_x:
        part = ImageOps.mirror(part)
    im = place_exact(part, dst)
    im.save(OUT / f"{name}.png", optimize=True)
    print(f"{name:20} src={src} dst={dst} bbox={alpha_bbox(im)}")
    return im


def save_shape_crop(sheet: Image.Image, name: str, src: Box, dst: Box, *, shape: str = "ellipse") -> Image.Image:
    part = sheet.crop(src).convert("RGBA")
    mask = Image.new("L", part.size, 0)
    draw = ImageDraw.Draw(mask)
    if shape == "ellipse":
        draw.ellipse((1, 1, part.width - 2, part.height - 2), fill=255)
    else:
        draw.rounded_rectangle((1, 1, part.width - 2, part.height - 2), radius=max(2, part.height // 3), fill=255)
    arr = np.array(part)
    arr[:, :, 3] = np.minimum(arr[:, :, 3], np.array(mask))
    arr[arr[:, :, 3] == 0, :3] = 0
    im = place_exact(Image.fromarray(arr, "RGBA"), dst)
    im.save(OUT / f"{name}.png", optimize=True)
    print(f"{name:20} src={src} dst={dst} bbox={alpha_bbox(im)}")
    return im


def save_solid_ellipse(name: str, dst: Box, color: tuple[int, int, int, int]) -> Image.Image:
    w, h = dst[2] - dst[0], dst[3] - dst[1]
    part = Image.new("RGBA", (w, h), (0, 0, 0, 0))
    draw = ImageDraw.Draw(part)
    draw.ellipse((0, 0, w - 1, h - 1), fill=color)
    canvas = Image.new("RGBA", CANVAS, (0, 0, 0, 0))
    canvas.alpha_composite(part, (dst[0], dst[1]))
    canvas.save(OUT / f"{name}.png", optimize=True)
    print(f"{name:20} solid dst={dst} bbox={alpha_bbox(canvas)}")
    return canvas


def save_duplicate(name: str, src_name: str) -> Image.Image:
    im = Image.open(OUT / f"{src_name}.png").convert("RGBA")
    im.save(OUT / f"{name}.png", optimize=True)
    print(f"{name:20} duplicate={src_name} bbox={alpha_bbox(im)}")
    return im


def composite(names: Iterable[str], out_path: Path) -> Image.Image:
    bg = Image.new("RGBA", CANVAS, (30, 39, 36, 255))
    for name in names:
        bg.alpha_composite(Image.open(OUT / f"{name}.png").convert("RGBA"))
    bg.save(out_path, optimize=True)
    print("preview", out_path)
    return bg


def make_contact() -> None:
    names = [
        "head-faceless",
        "torso",
        "hair-back",
        "hair-front",
        "hair-left",
        "hair-right",
        "brow-left",
        "brow-right",
        "eye-white-left",
        "eye-white-right",
        "iris-left",
        "iris-right",
        "pupil-left",
        "pupil-right",
        "highlight-left",
        "highlight-right",
        "lash-left",
        "lash-right",
        "eyelid-left-closed",
        "eyelid-right-closed",
        "nose",
        "mouth-neutral",
        "mouth-sad",
        "mouth-warm",
        "mouth-open",
    ]
    tile_w, tile_h = 180, 160
    cols = 5
    rows = (len(names) + cols - 1) // cols
    sheet = Image.new("RGBA", (cols * tile_w, rows * tile_h), (30, 39, 36, 255))
    draw = ImageDraw.Draw(sheet)
    for i, name in enumerate(names):
        layer = Image.open(OUT / f"{name}.png").convert("RGBA")
        box = alpha_bbox(layer)
        thumb = Image.new("RGBA", (1, 1), (0, 0, 0, 0)) if box is None else layer.crop(box)
        thumb.thumbnail((tile_w - 24, tile_h - 42), Image.Resampling.LANCZOS)
        x = (i % cols) * tile_w + (tile_w - thumb.width) // 2
        y = (i // cols) * tile_h + 16
        sheet.alpha_composite(thumb, (x, y))
        draw.text(((i % cols) * tile_w + 8, (i // cols) * tile_h + tile_h - 22), name, fill=(200, 214, 208, 255))
    sheet.save(HERE / "parts-contact.png", optimize=True)


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

    parts: dict[str, Image.Image] = {}

    # Body/base layers.
    parts["hair-back"] = save_part(sheet, "hair-back", (1145, 585, 1450, 925), (230, 90, 670, 560), lo=10, hi=46, min_area=900)
    parts["torso"] = save_part(sheet, "torso", (540, 720, 970, 920), (165, 585, 735, 825), lo=12, hi=54, min_area=900)
    parts["head-faceless"] = save_part(sheet, "head-faceless", (595, 40, 875, 465), (285, 98, 615, 600), lo=10, hi=48, min_area=900)

    # Hair layers. Front hair is intentionally above the face, as in a real PSD.
    parts["hair-left"] = save_part(sheet, "hair-left", (1118, 305, 1265, 600), (150, 200, 360, 555), lo=10, hi=46, min_area=900)
    parts["hair-right"] = save_part(sheet, "hair-right", (1362, 302, 1510, 600), (540, 200, 750, 555), lo=10, hi=46, min_area=900)
    parts["hair-front"] = save_part(sheet, "hair-front", (1135, 45, 1518, 260), (230, 84, 670, 320), lo=10, hi=46, min_area=900)

    # Brows and eye components. These are separate so a later renderer can move
    # irises/pupils within eye-white masks instead of scaling the whole eye.
    parts["brow-left"] = save_part(sheet, "brow-left", (920, 112, 1010, 140), (322, 292, 410, 314), lo=8, hi=38, min_area=18)
    parts["brow-right"] = save_part(sheet, "brow-right", (1040, 112, 1130, 140), (490, 292, 578, 314), lo=8, hi=38, min_area=18)
    parts["eye-white-left"] = save_shape_crop(sheet, "eye-white-left", (920, 154, 1015, 202), (323, 326, 397, 356), shape="ellipse")
    parts["eye-white-right"] = save_shape_crop(sheet, "eye-white-right", (920, 154, 1015, 202), (503, 326, 577, 356), shape="ellipse")
    parts["iris-left"] = save_part(sheet, "iris-left", (944, 232, 994, 282), (345, 326, 377, 358), lo=8, hi=42, keep_largest=True)
    parts["iris-right"] = save_part(sheet, "iris-right", (944, 232, 994, 282), (523, 326, 555, 358), lo=8, hi=42, keep_largest=True, flip_x=True)
    parts["pupil-left"] = save_part(sheet, "pupil-left", (954, 314, 990, 350), (353, 334, 373, 354), lo=8, hi=40, keep_largest=True)
    parts["pupil-right"] = save_part(sheet, "pupil-right", (954, 314, 990, 350), (527, 334, 547, 354), lo=8, hi=40, keep_largest=True, flip_x=True)
    parts["highlight-left"] = save_solid_ellipse("highlight-left", (360, 328, 368, 336), (248, 246, 238, 235))
    parts["highlight-right"] = save_solid_ellipse("highlight-right", (534, 328, 542, 336), (248, 246, 238, 235))
    parts["lash-left"] = save_part(sheet, "lash-left", (918, 420, 1018, 463), (318, 318, 402, 360), lo=8, hi=38, min_area=18)
    parts["lash-right"] = save_part(sheet, "lash-right", (918, 420, 1018, 463), (498, 318, 582, 360), lo=8, hi=38, min_area=18, flip_x=True)
    parts["eyelid-left-closed"] = save_part(sheet, "eyelid-left-closed", (918, 480, 1018, 510), (320, 338, 400, 360), lo=8, hi=38, min_area=18)
    parts["eyelid-right-closed"] = save_part(sheet, "eyelid-right-closed", (918, 480, 1018, 510), (500, 338, 580, 360), lo=8, hi=38, min_area=18, flip_x=True)

    parts["nose"] = save_part(sheet, "nose", (985, 550, 1045, 595), (424, 388, 476, 428), lo=8, hi=40, keep_largest=True)
    parts["mouth-neutral"] = save_part(sheet, "mouth-neutral", (965, 612, 1065, 650), (398, 430, 502, 458), lo=8, hi=40, keep_largest=True)
    parts["mouth-sad"] = save_part(sheet, "mouth-sad", (965, 672, 1065, 712), (398, 430, 502, 460), lo=8, hi=40, keep_largest=True)
    parts["mouth-warm"] = save_part(sheet, "mouth-warm", (965, 735, 1065, 775), (398, 430, 502, 460), lo=8, hi=40, keep_largest=True)
    parts["mouth-open"] = save_part(sheet, "mouth-open", (970, 790, 1065, 850), (400, 418, 500, 472), lo=8, hi=42, keep_largest=True)

    # Compatibility aliases for expression-level renderers.
    save_duplicate("brow-neutral", "brow-left")
    save_duplicate("eyes-neutral", "eye-white-left")

    neutral_order = [
        "hair-back",
        "torso",
        "head-faceless",
        "hair-front",
        "eye-white-left",
        "eye-white-right",
        "iris-left",
        "iris-right",
        "pupil-left",
        "pupil-right",
        "highlight-left",
        "highlight-right",
        "lash-left",
        "lash-right",
        "brow-left",
        "brow-right",
        "nose",
        "mouth-neutral",
    ]
    composite(neutral_order, HERE / "preview-neutral.png")

    blink_order = [
        "hair-back",
        "torso",
        "head-faceless",
        "hair-front",
        "eyelid-left-closed",
        "eyelid-right-closed",
        "brow-left",
        "brow-right",
        "nose",
        "mouth-neutral",
    ]
    composite(blink_order, HERE / "preview-blink.png")

    speaking_order = [
        *neutral_order[:-1],
        "mouth-open",
    ]
    composite(speaking_order, HERE / "preview-speaking.png")
    make_contact()
    return 0


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