from __future__ import annotations

import json
import struct
from pathlib import Path

from PIL import Image


ROOT = Path(__file__).resolve().parents[1]
ASSET_DIR = ROOT / "client" / "dev" / "res" / "custom" / "47"
LAYER_DIR = ASSET_DIR / "wulin_juexue_psd_layers"
SOURCE_PSD = ASSET_DIR / "\u6b66\u6797\u7edd\u5b66.psd"
OUTPUT_PSD = ASSET_DIR / "wulin_juexue_layered.psd"
OUTPUT_PSDTOOLS_PSD = ASSET_DIR / "wulin_juexue_layered_psdtools.psd"


def u32(value: int) -> bytes:
    return struct.pack(">I", value)


def i32(value: int) -> bytes:
    return struct.pack(">i", value)


def u16(value: int) -> bytes:
    return struct.pack(">H", value)


def i16(value: int) -> bytes:
    return struct.pack(">h", value)


def pascal_name(name: str) -> bytes:
    raw = name.encode("ascii", "replace")[:255]
    data = bytes([len(raw)]) + raw
    while len(data) % 4:
        data += b"\0"
    return data


def unicode_layer_name(name: str) -> bytes:
    encoded = name.encode("utf-16-be")
    data = u32(len(name)) + encoded
    if len(data) % 2:
        data += b"\0"
    return b"8BIM" + b"luni" + u32(len(data)) + data


def layer_channels(image: Image.Image) -> list[tuple[int, bytes]]:
    rgba = image.convert("RGBA")
    r, g, b, a = rgba.split()
    return [
        (0, b"\0\0" + r.tobytes()),
        (1, b"\0\0" + g.tobytes()),
        (2, b"\0\0" + b.tobytes()),
        (-1, b"\0\0" + a.tobytes()),
    ]


def layer_record(name: str, image: Image.Image, x: int, y: int) -> tuple[bytes, bytes]:
    width, height = image.size
    channels = layer_channels(image)

    record = bytearray()
    record += i32(y) + i32(x) + i32(y + height) + i32(x + width)
    record += u16(len(channels))
    for channel_id, channel_data in channels:
        record += i16(channel_id) + u32(len(channel_data))

    record += b"8BIM" + b"norm"
    record += bytes([255, 0, 8, 0])

    extra = bytearray()
    extra += u32(0)
    extra += u32(0)
    extra += pascal_name(name)
    extra += unicode_layer_name(name)
    record += u32(len(extra)) + extra

    channel_image_data = b"".join(channel_data for _, channel_data in channels)
    return bytes(record), channel_image_data


def full_canvas_layer(path: Path, canvas_size: tuple[int, int], offset: tuple[int, int]) -> Image.Image:
    image = Image.open(path).convert("RGBA")
    canvas = Image.new("RGBA", canvas_size, (0, 0, 0, 0))
    canvas.alpha_composite(image, offset)
    return canvas


def main() -> None:
    width, height = 816, 582

    base = Image.open(SOURCE_PSD).convert("RGBA")
    if base.size != (width, height):
        raise ValueError(f"Unexpected source PSD size: {base.size}")

    layers: list[tuple[str, Image.Image, int, int]] = [
        ("Frame_WulinJuexue", base, 0, 0),
        (
            "Content_Juexue",
            Image.open(LAYER_DIR / "layer_01_main_content_733x447_at_51_104.png").convert("RGBA"),
            51,
            104,
        ),
        (
            "Modal_Dim",
            Image.open(LAYER_DIR / "layer_02_modal_dim_816x582_at_0_0.png").convert("RGBA"),
            0,
            0,
        ),
        (
            "Popup_UnlockSlot",
            Image.open(LAYER_DIR / "layer_03_unlock_popup_420x261_at_198_172.png").convert("RGBA"),
            198,
            172,
        ),
    ]

    records = bytearray()
    channel_data = bytearray()
    for name, image, x, y in layers:
        record, data = layer_record(name, image, x, y)
        records += record
        channel_data += data

    layer_info = i16(len(layers)) + records + channel_data
    if len(layer_info) % 2:
        layer_info += b"\0"

    layer_and_mask = u32(len(layer_info)) + layer_info + u32(0)

    composite = Image.new("RGBA", (width, height), (0, 0, 0, 0))
    for _, image, x, y in layers:
        composite.alpha_composite(image, (x, y))
    r, g, b, _ = composite.split()
    merged_image_data = b"\0\0" + r.tobytes() + g.tobytes() + b.tobytes()

    header = (
        b"8BPS"
        + u16(1)
        + b"\0" * 6
        + u16(3)
        + u32(height)
        + u32(width)
        + u16(8)
        + u16(3)
    )

    psd = header + u32(0) + u32(0) + u32(len(layer_and_mask)) + layer_and_mask + merged_image_data
    OUTPUT_PSD.write_bytes(psd)

    manifest = {
        "output": OUTPUT_PSD.name,
        "canvas": {"width": width, "height": height},
        "layers_bottom_to_top": [
            {"name": name, "x": x, "y": y, "width": image.size[0], "height": image.size[1]}
            for name, image, x, y in layers
        ],
    }
    (LAYER_DIR / "layered_psd_manifest.json").write_text(
        json.dumps(manifest, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )
    try:
        from psd_tools import PSDImage

        psd = PSDImage.new("RGB", (width, height), color=(0, 0, 0))
        for name, image, x, y in layers:
            psd.create_pixel_layer(image, name=name, top=y, left=x)
        psd.save(OUTPUT_PSDTOOLS_PSD)
        print(OUTPUT_PSDTOOLS_PSD)
    except ImportError:
        print("psd-tools is not installed; skipped psd-tools PSD output.")

    print(OUTPUT_PSD)


if __name__ == "__main__":
    main()