from __future__ import annotations

import json
from collections import OrderedDict
from pathlib import Path

from PIL import Image
from psd_tools import PSDImage
from psd_tools.api.layers import PixelLayer


ROOT = Path(__file__).resolve().parents[1]
LIVE2D = ROOT / "03_Assets" / "Live2D"
LAYER_BASE = LIVE2D / "LayerPNGs"
MANIFEST = LIVE2D / "layer_manifest.json"

MATERIAL_PSD = LIVE2D / "sori_live2d_material_separation.psd"
IMPORT_PSD = LIVE2D / "sori_live2d_import.psd"


def load_manifest() -> dict:
    return json.loads(MANIFEST.read_text(encoding="utf-8"))


def make_psd(name: str, size: tuple[int, int]) -> PSDImage:
    psd = PSDImage.new("RGB", size, color=(0, 0, 0), depth=8)
    return psd


def add_grouped_layer(psd: PSDImage, groups: OrderedDict[str, object], layer: dict, visible: bool) -> None:
    group_name = layer["group"]
    group = groups.get(group_name)
    if group is None:
        group = psd.create_group(name=group_name)
        groups[group_name] = group

    path = LAYER_BASE / layer["file"]
    if not path.exists():
        raise FileNotFoundError(path)

    image = Image.open(path).convert("RGBA")
    pixel_layer = PixelLayer.frompil(image, parent=group, name=layer["id"], top=0, left=0)
    pixel_layer.visible = visible
    group.append(pixel_layer)


def save_psd(path: Path, layers: list[dict], size: tuple[int, int], material: bool) -> None:
    psd = make_psd(path.stem, size)
    groups: OrderedDict[str, object] = OrderedDict()

    for layer in layers:
        if not material and not layer.get("import", False):
            continue
        visible = layer.get("group") not in ("Guide", "SwapParts")
        add_grouped_layer(psd, groups, layer, visible)

    psd.save(path)
    reopened = PSDImage.open(path)
    if reopened.size != size:
        raise RuntimeError(f"Unexpected PSD size for {path}: {reopened.size}")
    if not list(reopened.descendants()):
        raise RuntimeError(f"No layers saved in {path}")


def main() -> None:
    manifest = load_manifest()
    canvas = manifest["canvas"]
    size = (int(canvas["width"]), int(canvas["height"]))
    layers = manifest["layers"]

    save_psd(MATERIAL_PSD, layers, size, material=True)
    save_psd(IMPORT_PSD, layers, size, material=False)
    print(MATERIAL_PSD)
    print(IMPORT_PSD)


if __name__ == "__main__":
    main()