from __future__ import annotations

from collections import deque
from pathlib import Path

import numpy as np
from PIL import Image, ImageFilter


ROOT = Path(__file__).resolve().parents[1]
REFERENCE = ROOT / "03_Assets" / "Reference"
PARTS = ROOT / "03_Assets" / "Parts" / "Images"
RAW = REFERENCE / "sori_generated_apose_raw.png"

SRC_W, SRC_H = 520, 900


def load_source() -> Image.Image:
    return Image.open(RAW).convert("RGBA")


def checker_background_mask(img: Image.Image) -> np.ndarray:
    arr = np.array(img.convert("RGBA"))
    rgb = arr[:, :, :3].astype(np.int16)
    alpha = arr[:, :, 3]
    light_gray = (rgb.min(axis=2) > 218) & ((rgb.max(axis=2) - rgb.min(axis=2)) < 18)
    transparent = alpha < 8
    candidate = light_gray | transparent

    h, w = candidate.shape
    seen = np.zeros_like(candidate, dtype=bool)
    q: deque[tuple[int, int]] = deque()
    for x in range(w):
        if candidate[0, x]:
            q.append((x, 0))
        if candidate[h - 1, x]:
            q.append((x, h - 1))
    for y in range(h):
        if candidate[y, 0]:
            q.append((0, y))
        if candidate[y, w - 1]:
            q.append((w - 1, y))

    while q:
        x, y = q.popleft()
        if seen[y, x] or not candidate[y, x]:
            continue
        seen[y, x] = True
        if x > 0:
            q.append((x - 1, y))
        if x < w - 1:
            q.append((x + 1, y))
        if y > 0:
            q.append((x, y - 1))
        if y < h - 1:
            q.append((x, y + 1))
    return seen


def remove_checker_background(img: Image.Image) -> Image.Image:
    bg = checker_background_mask(img)
    alpha = Image.fromarray((~bg).astype(np.uint8) * 255, "L")
    alpha = alpha.filter(ImageFilter.GaussianBlur(0.65)).point(lambda p: 0 if p < 28 else min(255, int(p * 1.18)))
    out = img.copy()
    out.putalpha(alpha)
    return out


def normalize_canvas(img: Image.Image) -> Image.Image:
    bbox = img.getbbox()
    if bbox is None:
        raise RuntimeError("source image has no visible pixels")
    crop = img.crop(bbox)
    ratio = min(430 / crop.width, 850 / crop.height)
    resized = crop.resize((round(crop.width * ratio), round(crop.height * ratio)), Image.Resampling.LANCZOS)
    canvas = Image.new("RGBA", (SRC_W, SRC_H), (0, 0, 0, 0))
    canvas.alpha_composite(resized, ((SRC_W - resized.width) // 2, 24))
    return canvas


def save_part(master: Image.Image, name: str, box: tuple[int, int, int, int]) -> None:
    mask = Image.new("L", master.size, 0)
    part_alpha = master.getchannel("A").crop(box)
    mask.paste(part_alpha, box)
    out = master.copy()
    out.putalpha(mask.filter(ImageFilter.GaussianBlur(0.25)))
    out.save(PARTS / name)


def main() -> None:
    PARTS.mkdir(parents=True, exist_ok=True)
    master = normalize_canvas(remove_checker_background(load_source()))
    master.save(PARTS / "sori_part_master_apose.png")

    # Coarse Live2D source parts. Boxes are in the normalized 520x900 master coordinate system.
    boxes = {
        "sori_part_head.png": (135, 8, 385, 255),
        "sori_part_neck.png": (215, 220, 305, 295),
        "sori_part_chest.png": (112, 212, 408, 435),
        "sori_part_pelvis.png": (145, 385, 375, 505),
        "sori_part_upperarm_l.png": (342, 250, 465, 452),
        "sori_part_forearm_l.png": (370, 360, 500, 535),
        "sori_part_hand_l.png": (382, 428, 520, 590),
        "sori_part_upperarm_r.png": (55, 250, 178, 452),
        "sori_part_forearm_r.png": (20, 360, 150, 535),
        "sori_part_hand_r.png": (0, 428, 138, 590),
        "sori_part_thigh_l.png": (260, 438, 385, 675),
        "sori_part_thigh_r.png": (135, 438, 260, 675),
        "sori_part_shin_l.png": (262, 640, 370, 825),
        "sori_part_shin_r.png": (150, 640, 258, 825),
        "sori_part_foot_l.png": (258, 780, 375, 900),
        "sori_part_foot_r.png": (145, 780, 262, 900),
    }
    for name, box in boxes.items():
        save_part(master, name, box)
    print(PARTS / "sori_part_master_apose.png")


if __name__ == "__main__":
    main()