#!/usr/bin/env python3
"""Export top-level PSB groups as normalized app-ready raster layers."""
from __future__ import annotations

import json
from pathlib import Path
from typing import Iterable

import numpy as np
from PIL import Image, ImageDraw
from psd_tools import PSDImage

HERE = Path(__file__).resolve().parent
PSB = Path(r"C:\Users\encep\OneDrive\문서\카카오톡 받은 파일\라투디 여캐.psb")
OUT = HERE / "top-groups"
PARTS = OUT / "parts"

SOURCE_CANVAS = (1080, 1920)
TARGET_CANVAS = (900, 1125)
CROP_BOX = (0, 240, 1080, 1590)

LAYER_ORDER = [
    "hair-back",
    "body",
    "outfit",
    "face",
    "eye-left",
    "eye-right",
    "hair-front",
    "mouth",
]

GROUP_MAP = {
    "뒷머리": "hair-back",
    "몸": "body",
    "의상": "outfit",
    "얼굴": "face",
    "좌 눈": "eye-left",
    "우 눈": "eye-right",
    "머리카락": "hair-front",
    "입": "mouth",
}

EXTRA_BACK_HAIR = {"레이어 110", "레이어 110 복사"}


def alpha_bbox(im: Image.Image, threshold: int = 8) -> tuple[int, int, int, int] | None:
    alpha = np.array(im.convert("RGBA").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 resize_rgba(im: Image.Image, size: tuple[int, int]) -> Image.Image:
    arr = np.array(im.convert("RGBA")).astype(np.float32)
    alpha = arr[:, :, 3:4] / 255.0
    premultiplied = arr.copy()
    premultiplied[:, :, :3] *= alpha
    resized = Image.fromarray(np.clip(premultiplied, 0, 255).astype(np.uint8), "RGBA").resize(
        size,
        Image.Resampling.LANCZOS,
    )
    out = np.array(resized).astype(np.float32)
    out_alpha = out[:, :, 3:4] / 255.0
    out[:, :, :3] = np.where(out_alpha > 0.001, out[:, :, :3] / np.maximum(out_alpha, 0.001), 0)
    out[out[:, :, 3] < 1, :3] = 0
    return Image.fromarray(np.clip(out, 0, 255).astype(np.uint8), "RGBA")


def normalize(full_canvas: Image.Image) -> Image.Image:
    return resize_rgba(full_canvas.crop(CROP_BOX), TARGET_CANVAS)


def composite_layers(layers: Iterable[Image.Image]) -> Image.Image:
    canvas = Image.new("RGBA", SOURCE_CANVAS, (0, 0, 0, 0))
    for layer in layers:
        canvas.alpha_composite(layer.convert("RGBA"))
    return canvas


def make_contact(entries: list[dict[str, object]]) -> None:
    tile_w, tile_h = 180, 180
    sheet = Image.new("RGBA", (tile_w * 4, tile_h * 2), (30, 39, 36, 255))
    draw = ImageDraw.Draw(sheet)
    for i, entry in enumerate(entries):
        layer = Image.open(PARTS / str(entry["file"])).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 - 22, tile_h - 46), Image.Resampling.LANCZOS)
        x = (i % 4) * tile_w + (tile_w - thumb.width) // 2
        y = (i // 4) * tile_h + 12
        sheet.alpha_composite(thumb, (x, y))
        draw.text(((i % 4) * tile_w + 8, (i // 4) * tile_h + tile_h - 28), str(entry["id"]), fill=(210, 222, 216, 255))
        draw.text(((i % 4) * tile_w + 8, (i // 4) * tile_h + tile_h - 14), str(entry["source"]), fill=(150, 166, 160, 255))
    sheet.save(OUT / "top-groups-contact.png", optimize=True)


def main() -> int:
    psd = PSDImage.open(PSB)
    PARTS.mkdir(parents=True, exist_ok=True)

    top = {layer.name: layer for layer in psd}
    entries: list[dict[str, object]] = []

    extra_back = [top[name].composite(viewport=(0, 0, psd.width, psd.height)) for name in EXTRA_BACK_HAIR if name in top]

    for source_name, part_id in GROUP_MAP.items():
        if source_name not in top:
            raise KeyError(f"Missing top-level group: {source_name}")
        image = top[source_name].composite(viewport=(0, 0, psd.width, psd.height))
        if part_id == "hair-back":
            image = composite_layers([*extra_back, image])
        normalized = normalize(image)
        file_name = f"{part_id}.png"
        normalized.save(PARTS / file_name, optimize=True)
        entries.append(
            {
                "id": part_id,
                "source": source_name,
                "file": file_name,
                "alphaBox": list(alpha_bbox(normalized) or ()),
            }
        )

    preview = Image.new("RGBA", TARGET_CANVAS, (30, 39, 36, 255))
    for part_id in LAYER_ORDER:
        preview.alpha_composite(Image.open(PARTS / f"{part_id}.png").convert("RGBA"))
    preview.save(OUT / "preview-psb-groups.png", optimize=True)

    transparent = Image.new("RGBA", TARGET_CANVAS, (0, 0, 0, 0))
    for part_id in LAYER_ORDER:
        transparent.alpha_composite(Image.open(PARTS / f"{part_id}.png").convert("RGBA"))
    transparent.save(OUT / "preview-psb-groups-transparent.png", optimize=True)

    (OUT / "top-groups-manifest.json").write_text(
        json.dumps(
            {
                "source": str(PSB),
                "sourceCanvas": SOURCE_CANVAS,
                "targetCanvas": TARGET_CANVAS,
                "cropBox": CROP_BOX,
                "layerOrder": LAYER_ORDER,
                "parts": entries,
            },
            ensure_ascii=False,
            indent=2,
        ),
        encoding="utf-8",
    )
    make_contact(entries)
    print(f"exported {len(entries)} top-level groups")
    print(f"preview {OUT / 'preview-psb-groups.png'}")
    return 0


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