#!/usr/bin/env python3
"""발행된 변주 PNG(900x1125)를 Live2D식 파츠로 분리한다.

구조:
  parts/base-faceless.png     : 눈/코/입/눈썹을 지운 몸+머리+피부 베이스
  parts/shoulders.png         : 어깨/상의
  parts/neck.png              : 목
  parts/face-faceless.png     : 눈코입 없는 얼굴/피부
  parts/forehead.png          : 이마 피부
  parts/ear-left/right.png    : 귀
  parts/brow-<expr>.png       : 눈썹
  parts/eyes-<expr>.png       : 눈
  parts/eyelid-closed.png     : 닫힌 눈(깜빡임)
  parts/nose-neutral.png      : 코(정적)
  parts/mouth-<expr>.png      : 입 + mouth-open(발화)
  parts/hair-left/right/bangs.png : 찰랑임용 얇은 머리카락 오버레이

핵심은 base에 기존 눈코입을 남기지 않는 것. 원본 얼굴 위에 표정 패치를 얹으면
눈/입이 겹쳐 보여 불쾌해지므로, 달걀귀신 같은 베이스 위에 파츠를 올린다.
"""
import os
from PIL import Image, ImageDraw, ImageFilter
import numpy as np

SRC = r"D:/workspace/vignette/apps/web/public/avatar/seoyeon"
OUT = os.path.join(SRC, "parts")
os.makedirs(OUT, exist_ok=True)

# 현재 서연 v1 게시 에셋(900x1125) 기준 ROI.
BROW_ROI = (220, 300, 684, 372)
EYES_ROI = (214, 340, 696, 458)
EYELID_ROI = (214, 336, 696, 464)
NOSE_ROI = (360, 400, 546, 548)
MOUTH_ROI = (314, 486, 606, 650)
HAIR_BANGS_ROI = (210, 72, 696, 356)
HAIR_LEFT_ROI = (34, 88, 360, 1034)
HAIR_RIGHT_ROI = (548, 88, 866, 1034)
SHOULDERS_ROI = (0, 745, 900, 1125)
NECK_ROI = (300, 610, 604, 858)
FACE_ROI = (198, 160, 712, 690)
FOREHEAD_ROI = (278, 198, 628, 344)
EAR_LEFT_ROI = (126, 340, 278, 514)
EAR_RIGHT_ROI = (628, 340, 780, 514)

# diff seed → MaxFilter 팽창 → blur 페더. 보조/호환 파츠용.
UPPERFACE_ROI = (218, 314, 684, 456)
UPPERFACE_DILATE = 29
UPPERFACE_FEATHER = 6
EYELID_DILATE = 31
EYELID_FEATHER = 6
MOUTH_DILATE = 27
MOUTH_FEATHER = 6

EXPR_VARIANTS = ["neutral", "sad", "tired", "anxious", "warm", "startled"]


def odd(value: int) -> int:
    return value if value % 2 else value + 1


def empty_part(out_name: str, size: tuple[int, int]) -> None:
    im = Image.new("RGBA", size, (0, 0, 0, 0))
    im.save(f"{OUT}/{out_name}.png", optimize=True)
    print(f"  {out_name}.png  blank")


def save_part(arr: np.ndarray, out_name: str) -> None:
    arr = arr.astype(np.uint8)
    Image.fromarray(arr, "RGBA").save(f"{OUT}/{out_name}.png", optimize=True)
    al = arr[:, :, 3]
    ys, xs = np.where(al > 0)
    bbox = None if len(xs) == 0 else (int(xs.min()), int(ys.min()), int(xs.max() + 1), int(ys.max() + 1))
    print(f"  {out_name}.png  bbox={bbox} "
          f"opaque%={float((al == 255).mean()) * 100:.2f} "
          f"vis%={float((al > 0).mean()) * 100:.2f}")


def full_mask(size: tuple[int, int]) -> Image.Image:
    return Image.new("L", size, 0)


def soft_shapes(size: tuple[int, int], shapes: list[tuple[str, tuple[int, int, int, int]]], feather: int) -> Image.Image:
    mask = full_mask(size)
    draw = ImageDraw.Draw(mask)
    for kind, box in shapes:
        if kind == "ellipse":
            draw.ellipse(box, fill=255)
        elif kind == "round":
            draw.rounded_rectangle(box, radius=max(4, min(box[2] - box[0], box[3] - box[1]) // 3), fill=255)
        else:
            draw.rectangle(box, fill=255)
    if feather > 0:
        mask = mask.filter(ImageFilter.GaussianBlur(feather))
    return mask


def mask_roi(size: tuple[int, int], roi: tuple[int, int, int, int], feather: int) -> Image.Image:
    return soft_shapes(size, [("round", roi)], feather)


def alpha_part(im: Image.Image, mask: Image.Image, out_name: str) -> None:
    arr = np.array(im.convert("RGBA")).astype(np.float32)
    m = np.array(mask).astype(np.float32)
    arr[:, :, 3] = np.minimum(arr[:, :, 3], m)
    arr[arr[:, :, 3] < 1, :3] = 0
    save_part(arr, out_name)


def make_feature_masks(size: tuple[int, int]) -> dict[str, Image.Image]:
    return {
        "shoulders": soft_shapes(size, [("round", SHOULDERS_ROI)], 10),
        "neck": soft_shapes(size, [("ellipse", NECK_ROI)], 12),
        "face": soft_shapes(size, [("ellipse", FACE_ROI)], 12),
        "forehead": soft_shapes(size, [("ellipse", FOREHEAD_ROI)], 8),
        "ear-left": soft_shapes(size, [("ellipse", EAR_LEFT_ROI)], 7),
        "ear-right": soft_shapes(size, [("ellipse", EAR_RIGHT_ROI)], 7),
        "brow": soft_shapes(
            size,
            [
                ("ellipse", (238, 306, 424, 370)),
                ("ellipse", (484, 306, 670, 370)),
            ],
            7,
        ),
        "eyes": soft_shapes(
            size,
            [
                ("ellipse", (220, 328, 440, 452)),
                ("ellipse", (470, 328, 690, 452)),
            ],
            8,
        ),
        "eyelid": soft_shapes(
            size,
            [
                ("ellipse", (218, 326, 442, 458)),
                ("ellipse", (468, 326, 692, 458)),
            ],
            8,
        ),
        "nose": soft_shapes(size, [("ellipse", NOSE_ROI)], 10),
        "mouth": soft_shapes(size, [("ellipse", MOUTH_ROI)], 9),
    }


def make_faceless_base(base: Image.Image, masks: dict[str, Image.Image]) -> None:
    arr = np.array(base.convert("RGBA")).astype(np.float32)
    H, W = arr.shape[:2]
    feature = np.zeros((H, W), dtype=np.float32)
    for key in ["brow", "eyes", "nose", "mouth"]:
        feature = np.maximum(feature, np.array(masks[key]).astype(np.float32))

    # 주변 피부색으로 덮는다. 오버레이 파츠가 올라올 자리라 완벽한 인페인팅보다
    # 기존 눈코입 흔적 제거가 더 중요하다.
    skin = arr[:, :, :3]
    r, g, b = skin[:, :, 0], skin[:, :, 1], skin[:, :, 2]
    alpha = arr[:, :, 3]
    skin_pixels = (
        (alpha > 160)
        & (r > 145)
        & (g > 100)
        & (b > 75)
        & (r > b + 22)
        & (g > b + 8)
        & (feature < 16)
    )
    if skin_pixels.any():
        fill = np.median(skin[skin_pixels], axis=0)
    else:
        fill = np.array([224, 180, 145], dtype=np.float32)

    smooth = np.array(base.filter(ImageFilter.GaussianBlur(18)).convert("RGBA")).astype(np.float32)
    cover_rgb = smooth[:, :, :3] * 0.35 + fill.reshape(1, 1, 3) * 0.65
    weight = (feature / 255.0)[:, :, None]
    arr[:, :, :3] = arr[:, :, :3] * (1 - weight) + cover_rgb * weight
    arr[:, :, 3] = np.array(base.getchannel("A")).astype(np.float32)
    save_part(arr, "base-faceless")


def make_faceless_skin_part(base: Image.Image, masks: dict[str, Image.Image], part_key: str, out_name: str) -> None:
    arr = np.array(base.convert("RGBA")).astype(np.float32)
    H, W = arr.shape[:2]
    feature = np.zeros((H, W), dtype=np.float32)
    for key in ["brow", "eyes", "nose", "mouth"]:
        feature = np.maximum(feature, np.array(masks[key]).astype(np.float32))

    r, g, b = arr[:, :, 0], arr[:, :, 1], arr[:, :, 2]
    src_a = arr[:, :, 3]
    skin_pixels = (
        (src_a > 160)
        & (r > 145)
        & (g > 100)
        & (b > 75)
        & (r > b + 22)
        & (g > b + 8)
        & (feature < 16)
    )
    fill = np.median(arr[:, :, :3][skin_pixels], axis=0) if skin_pixels.any() else np.array([224, 180, 145])
    smooth = np.array(base.filter(ImageFilter.GaussianBlur(18)).convert("RGBA")).astype(np.float32)
    cover_rgb = smooth[:, :, :3] * 0.35 + fill.reshape(1, 1, 3) * 0.65
    remove_weight = (feature / 255.0)[:, :, None]
    arr[:, :, :3] = arr[:, :, :3] * (1 - remove_weight) + cover_rgb * remove_weight

    mask = np.array(masks[part_key]).astype(np.float32)
    arr[:, :, 3] = np.minimum(src_a, mask)
    arr[arr[:, :, 3] < 1, :3] = 0
    save_part(arr, out_name)


def hair_seed(im: Image.Image, roi: tuple[int, int, int, int]) -> Image.Image:
    arr = np.array(im.convert("RGBA")).astype(np.int16)
    r, g, b, a = arr[:, :, 0], arr[:, :, 1], arr[:, :, 2], arr[:, :, 3]
    # 갈색 머리 위주. 피부/셔츠/배경은 제외한다.
    hair = (
        (a > 80)
        & (r > 22)
        & (r < 155)
        & (g > 18)
        & (g < 130)
        & (b > 14)
        & (b < 120)
        & (r >= g - 8)
        & (g >= b - 12)
    )
    full = np.zeros(a.shape, dtype=np.uint8)
    x1, y1, x2, y2 = roi
    full[y1:y2, x1:x2] = hair[y1:y2, x1:x2].astype(np.uint8) * 255
    mask = Image.fromarray(full, "L").filter(ImageFilter.MaxFilter(9)).filter(ImageFilter.GaussianBlur(3))
    return mask


def make_hair_part(base: Image.Image, roi: tuple[int, int, int, int], out_name: str) -> None:
    alpha_part(base, hair_seed(base, roi), out_name)


def diff_seed(base: Image.Image, variant: Image.Image, roi: tuple[int, int, int, int], threshold: int) -> Image.Image:
    b = np.array(base.convert("RGBA")).astype(np.int16)
    v = np.array(variant.convert("RGBA")).astype(np.int16)
    rgb_diff = np.abs(v[:, :, :3] - b[:, :, :3]).max(axis=2)
    alpha_gate = (b[:, :, 3] > 96) | (v[:, :, 3] > 96)
    seed = ((rgb_diff >= threshold) & alpha_gate).astype(np.uint8) * 255
    full = np.zeros(seed.shape, dtype=np.uint8)
    x1, y1, x2, y2 = roi
    full[y1:y2, x1:x2] = seed[y1:y2, x1:x2]
    return Image.fromarray(full, "L")


def expand_mask(seed: Image.Image, dilate: int, feather: int) -> Image.Image:
    mask = seed.filter(ImageFilter.MaxFilter(odd(dilate)))
    if feather > 0:
        mask = mask.filter(ImageFilter.GaussianBlur(feather))
    return mask


def make_diff_part(
    base: Image.Image,
    variant: str,
    roi: tuple[int, int, int, int],
    out_name: str,
    *,
    threshold: int,
    dilate: int,
    feather: int,
) -> None:
    im = Image.open(f"{SRC}/{variant}.png").convert("RGBA")
    seed = diff_seed(base, im, roi, threshold)
    mask = expand_mask(seed, dilate, feather)
    arr = np.array(im).astype(np.float32)
    m = np.array(mask).astype(np.float32)
    src_a = arr[:, :, 3]
    alpha = np.minimum(src_a, m)
    arr[:, :, 3] = alpha
    arr[alpha < 1, :3] = 0
    Image.fromarray(arr.astype(np.uint8), "RGBA").save(f"{OUT}/{out_name}.png", optimize=True)
    al = arr[:, :, 3]
    ys, xs = np.where(al > 0)
    bbox = None if len(xs) == 0 else (int(xs.min()), int(ys.min()), int(xs.max() + 1), int(ys.max() + 1))
    print(f"  {out_name}.png  bbox={bbox} "
          f"opaque%={float((al == 255).mean()) * 100:.2f} "
          f"vis%={float((al > 0).mean()) * 100:.2f}")


def main() -> int:
    base = Image.open(f"{SRC}/neutral.png").convert("RGBA")
    masks = make_feature_masks(base.size)
    print("[base]")
    make_faceless_base(base, masks)
    print("[body]")
    alpha_part(base, masks["shoulders"], "shoulders")
    make_faceless_skin_part(base, masks, "neck", "neck")
    make_faceless_skin_part(base, masks, "face", "face-faceless")
    make_faceless_skin_part(base, masks, "forehead", "forehead")
    alpha_part(base, masks["ear-left"], "ear-left")
    alpha_part(base, masks["ear-right"], "ear-right")
    print("[hair]")
    make_hair_part(base, HAIR_LEFT_ROI, "hair-left")
    make_hair_part(base, HAIR_RIGHT_ROI, "hair-right")
    make_hair_part(base, HAIR_BANGS_ROI, "hair-bangs")
    print("[brow]")
    for v in EXPR_VARIANTS:
        alpha_part(Image.open(f"{SRC}/{v}.png").convert("RGBA"), masks["brow"], f"brow-{v}")
    print("[eyes]")
    for v in EXPR_VARIANTS:
        alpha_part(Image.open(f"{SRC}/{v}.png").convert("RGBA"), masks["eyes"], f"eyes-{v}")
    print("[nose]")
    alpha_part(base, masks["nose"], "nose-neutral")
    print("[upperface]")
    for v in EXPR_VARIANTS:
        # 구버전 렌더러 호환용. 새 렌더러는 brow/eyes를 사용한다.
        make_diff_part(
            base,
            v,
            UPPERFACE_ROI,
            f"upperface-{v}",
            threshold=18 if v != "neutral" else 1,
            dilate=UPPERFACE_DILATE,
            feather=UPPERFACE_FEATHER,
        )
    print("[eyelid]")
    alpha_part(Image.open(f"{SRC}/eyes-closed.png").convert("RGBA"), masks["eyelid"], "eyelid-closed")
    print("[mouth]")
    for v in EXPR_VARIANTS:
        alpha_part(Image.open(f"{SRC}/{v}.png").convert("RGBA"), masks["mouth"], f"mouth-{v}")
    alpha_part(Image.open(f"{SRC}/speaking.png").convert("RGBA"), masks["mouth"], "mouth-open")
    print("done ->", OUT)
    return 0


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