#!/usr/bin/env python3
"""Losslessly extract and decode resources from the 16-bit TDKPIN NE image."""

from __future__ import annotations

import hashlib
import io
import json
import re
import struct
import wave
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from PIL import Image


ROOT = Path(__file__).resolve().parent.parent
SOURCE = ROOT / "TDKPIN.EXE"
OUTPUT = ROOT / "assets"

INTEGER_TYPES = {
    1: "CURSOR",
    2: "BITMAP",
    3: "ICON",
    4: "MENU",
    5: "DIALOG",
    6: "STRING",
    7: "FONTDIR",
    8: "FONT",
    9: "ACCELERATOR",
    10: "RCDATA",
    11: "MESSAGETABLE",
    12: "GROUP_CURSOR",
    14: "GROUP_ICON",
    15: "VERSION_OLD",
    16: "VERSION",
}


@dataclass
class Resource:
    type_name: str
    resource_id: str
    offset: int
    allocated_size: int
    flags: int
    data: bytes

    @property
    def stem(self) -> str:
        return f"{int(self.resource_id):05d}" if self.resource_id.isdigit() else self.resource_id


def u16(data: bytes, offset: int) -> int:
    return struct.unpack_from("<H", data, offset)[0]


def u32(data: bytes, offset: int) -> int:
    return struct.unpack_from("<I", data, offset)[0]


def read_pascal_string(image: bytes, base: int, offset: int) -> str:
    length = image[base + offset]
    return image[base + offset + 1 : base + offset + 1 + length].decode("cp1252")


def resource_name(image: bytes, base: int, value: int, is_type: bool) -> str:
    if value & 0x8000:
        number = value & 0x7FFF
        return INTEGER_TYPES.get(number, f"TYPE_{number}") if is_type else str(number)
    return read_pascal_string(image, base, value)


def parse_resources(image: bytes) -> tuple[list[Resource], dict[str, Any]]:
    ne_offset = u32(image, 0x3C)
    resource_offset = u16(image, ne_offset + 0x24)
    resident_names_offset = u16(image, ne_offset + 0x26)
    table_base = ne_offset + resource_offset
    pos = table_base
    shift = u16(image, pos)
    pos += 2
    resources: list[Resource] = []

    while True:
        type_id = u16(image, pos)
        if type_id == 0:
            break
        count = u16(image, pos + 2)
        pos += 8
        type_name = resource_name(image, table_base, type_id, True)
        for _ in range(count):
            offset_units, length_units, flags, name_id, _handle, _usage = struct.unpack_from(
                "<HHHHHH", image, pos
            )
            pos += 12
            offset = offset_units << shift
            allocated_size = length_units << shift
            data = image[offset : offset + allocated_size]
            if len(data) != allocated_size:
                raise ValueError(f"truncated resource {type_name}/{name_id:#x}")
            resources.append(
                Resource(
                    type_name=type_name,
                    resource_id=resource_name(image, table_base, name_id, False),
                    offset=offset,
                    allocated_size=allocated_size,
                    flags=flags,
                    data=data,
                )
            )

    metadata = {
        "ne_header_offset": ne_offset,
        "resource_table_offset": table_base,
        "resource_table_end": ne_offset + resident_names_offset,
        "alignment_shift": shift,
        "alignment_bytes": 1 << shift,
    }
    return resources, metadata


def dib_info(data: bytes) -> dict[str, int]:
    if len(data) < 40 or u32(data, 0) != 40:
        raise ValueError("not a BITMAPINFOHEADER DIB")
    width, height = struct.unpack_from("<ii", data, 4)
    planes, bits_per_pixel = struct.unpack_from("<HH", data, 12)
    compression = u32(data, 16)
    image_size = u32(data, 20)
    colors_used = u32(data, 32)
    return {
        "width": width,
        "height": height,
        "planes": planes,
        "bits_per_pixel": bits_per_pixel,
        "compression": compression,
        "image_size": image_size,
        "colors_used": colors_used,
    }


def decode_custom_dib(data: bytes, palette: bytes) -> tuple[Image.Image, int, dict[str, int]]:
    info = dib_info(data)
    width = info["width"]
    height = abs(info["height"])
    bits_per_pixel = info["bits_per_pixel"]
    if bits_per_pixel != 8 or info["compression"] != 0:
        raise ValueError(f"unsupported custom DIB: {bits_per_pixel} bpp compression={info['compression']}")
    stride = ((width * bits_per_pixel + 31) // 32) * 4
    pixel_size = info["image_size"] or stride * height
    pixel_offset = 40
    pixels = data[pixel_offset : pixel_offset + pixel_size]
    if len(pixels) != pixel_size:
        raise ValueError("truncated custom DIB pixels")
    rows = [pixels[i * stride : i * stride + width] for i in range(height)]
    if info["height"] > 0:
        rows.reverse()
    image = Image.frombytes("P", (width, height), b"".join(rows))
    image.putpalette(palette)
    return image, pixel_offset + pixel_size, info


def wrap_bitmap_dib(data: bytes) -> tuple[bytes, int, dict[str, int]]:
    info = dib_info(data)
    width = info["width"]
    height = abs(info["height"])
    bits_per_pixel = info["bits_per_pixel"]
    color_count = info["colors_used"] or (1 << bits_per_pixel if bits_per_pixel <= 8 else 0)
    pixel_offset = 40 + color_count * 4
    stride = ((width * bits_per_pixel + 31) // 32) * 4
    pixel_size = info["image_size"] or stride * height
    actual_size = pixel_offset + pixel_size
    if actual_size > len(data):
        raise ValueError("truncated bitmap DIB")
    header = struct.pack("<2sIHHI", b"BM", 14 + actual_size, 0, 0, 14 + pixel_offset)
    return header + data[:actual_size], actual_size, info


def decode_icon(data: bytes) -> tuple[bytes, int, dict[str, int]]:
    info = dib_info(data)
    width = info["width"]
    height = abs(info["height"]) // 2
    bits_per_pixel = info["bits_per_pixel"]
    color_count = info["colors_used"] or (1 << bits_per_pixel if bits_per_pixel <= 8 else 0)
    xor_stride = ((width * bits_per_pixel + 31) // 32) * 4
    and_stride = ((width + 31) // 32) * 4
    actual_size = 40 + color_count * 4 + xor_stride * height + and_stride * height
    if actual_size > len(data):
        raise ValueError("truncated icon DIB")
    directory = struct.pack(
        "<BBBBHHII",
        width if width < 256 else 0,
        height if height < 256 else 0,
        color_count if color_count < 256 else 0,
        0,
        info["planes"],
        bits_per_pixel,
        actual_size,
        22,
    )
    ico = struct.pack("<HHH", 0, 1, 1) + directory + data[:actual_size]
    normalized = dict(info)
    normalized["height"] = height
    return ico, actual_size, normalized


def read_ansi_z(data: bytes, pos: int) -> tuple[str, int]:
    end = data.index(0, pos)
    return data[pos:end].decode("cp1252", errors="replace"), end + 1


def read_dialog_field(data: bytes, pos: int) -> tuple[Any, int]:
    first = data[pos]
    if first == 0:
        return None, pos + 1
    if first == 0xFF:
        return {"ordinal": u16(data, pos + 1)}, pos + 3
    if first >= 0x80:
        atoms = {0x80: "BUTTON", 0x81: "EDIT", 0x82: "STATIC", 0x83: "LISTBOX", 0x84: "SCROLLBAR", 0x85: "COMBOBOX"}
        return {"class_atom": atoms.get(first, f"ATOM_{first:02X}")}, pos + 1
    return read_ansi_z(data, pos)


def decode_dialog(data: bytes) -> tuple[dict[str, Any], int]:
    style = u32(data, 0)
    count = data[4]
    x, y, cx, cy = struct.unpack_from("<hhhh", data, 5)
    pos = 13
    menu, pos = read_dialog_field(data, pos)
    window_class, pos = read_dialog_field(data, pos)
    caption, pos = read_dialog_field(data, pos)
    dialog: dict[str, Any] = {
        "style": f"0x{style:08x}",
        "bounds": {"x": x, "y": y, "width": cx, "height": cy},
        "menu": menu,
        "class": window_class,
        "caption": caption,
        "declared_control_count": count,
        "controls": [],
    }
    if style & 0x40:
        point_size = u16(data, pos)
        pos += 2
        font, pos = read_dialog_field(data, pos)
        dialog["font"] = {"point_size": point_size, "name": font}
    for _ in range(count):
        # Win16 templates pack controls byte-for-byte; unlike Win32 templates,
        # DLGITEMTEMPLATE records are not DWORD- or WORD-aligned.
        x, y, cx, cy, control_id = struct.unpack_from("<hhhhh", data, pos)
        style = u32(data, pos + 10)
        pos += 14
        control_class, pos = read_dialog_field(data, pos)
        title, pos = read_dialog_field(data, pos)
        extra_size = data[pos]
        pos += 1
        extra = data[pos : pos + extra_size]
        pos += extra_size
        dialog["controls"].append(
            {
                "id": control_id & 0xFFFF,
                "bounds": {"x": x, "y": y, "width": cx, "height": cy},
                "style": f"0x{style:08x}",
                "class": control_class,
                "title": title,
                "extra_hex": extra.hex(),
            }
        )
    return dialog, pos


def decode_version(data: bytes) -> dict[str, Any]:
    strings = [
        match.decode("cp1252", errors="replace")
        for match in re.findall(rb"[\x20-\x7e]{3,}", data)
    ]
    result: dict[str, Any] = {"printable_strings": strings}
    signature = data.find(b"\xbd\x04\xef\xfe")
    if signature >= 0 and signature + 52 <= len(data):
        fields = struct.unpack_from("<13I", data, signature)
        result["fixed_file_info"] = {
            "signature": f"0x{fields[0]:08x}",
            "structure_version": f"0x{fields[1]:08x}",
            "file_version_ms": f"0x{fields[2]:08x}",
            "file_version_ls": f"0x{fields[3]:08x}",
            "product_version_ms": f"0x{fields[4]:08x}",
            "product_version_ls": f"0x{fields[5]:08x}",
            "file_flags_mask": f"0x{fields[6]:08x}",
            "file_flags": f"0x{fields[7]:08x}",
            "file_os": f"0x{fields[8]:08x}",
            "file_type": f"0x{fields[9]:08x}",
            "file_subtype": f"0x{fields[10]:08x}",
            "file_date_ms": f"0x{fields[11]:08x}",
            "file_date_ls": f"0x{fields[12]:08x}",
        }
    return result


def wav_info(data: bytes) -> dict[str, Any]:
    with wave.open(io.BytesIO(data), "rb") as stream:
        frames = stream.getnframes()
        rate = stream.getframerate()
        return {
            "format": "PCM",
            "channels": stream.getnchannels(),
            "sample_rate": rate,
            "sample_width_bits": stream.getsampwidth() * 8,
            "frames": frames,
            "duration_seconds": frames / rate,
        }


def save_palette(palette: bytes, resource_id: str) -> list[str]:
    if len(palette) < 768:
        raise ValueError("palette is shorter than 256 RGB entries")
    palette = palette[:768]
    image = Image.new("RGB", (256, 256))
    pixels = image.load()
    for index in range(256):
        color = tuple(palette[index * 3 : index * 3 + 3])
        x0 = (index % 16) * 16
        y0 = (index // 16) * 16
        for y in range(y0, y0 + 16):
            for x in range(x0, x0 + 16):
                pixels[x, y] = color
    png_path = OUTPUT / "decoded" / "palettes" / f"palette_{int(resource_id):04d}.png"
    png_path.parent.mkdir(parents=True, exist_ok=True)
    image.save(png_path)
    gpl_path = png_path.with_suffix(".gpl")
    lines = ["GIMP Palette", f"Name: TDKPIN resource {resource_id}", "Columns: 16", "#"]
    for index in range(256):
        red, green, blue = palette[index * 3 : index * 3 + 3]
        lines.append(f"{red:3d} {green:3d} {blue:3d}\tIndex {index}")
    gpl_path.write_text("\n".join(lines) + "\n", encoding="utf-8")
    return [str(png_path.relative_to(ROOT)), str(gpl_path.relative_to(ROOT))]


def main() -> None:
    image = SOURCE.read_bytes()
    resources, ne_metadata = parse_resources(image)
    palette_resources = [resource for resource in resources if resource.type_name == "PAL"]
    if len(palette_resources) != 1:
        raise ValueError(f"expected one PAL resource, found {len(palette_resources)}")
    shared_palette = palette_resources[0].data[:768]
    manifest_entries: list[dict[str, Any]] = []
    extraction_errors: list[str] = []

    for resource in resources:
        raw_path = OUTPUT / "raw" / resource.type_name / f"{resource.stem}.bin"
        raw_path.parent.mkdir(parents=True, exist_ok=True)
        raw_path.write_bytes(resource.data)
        entry: dict[str, Any] = {
            "type": resource.type_name,
            "id": resource.resource_id,
            "file_offset": resource.offset,
            "allocated_size": resource.allocated_size,
            "flags": f"0x{resource.flags:04x}",
            "raw_sha256": hashlib.sha256(resource.data).hexdigest(),
            "raw_path": str(raw_path.relative_to(ROOT)),
            "decoded_paths": [],
        }
        try:
            if resource.type_name == "DAT":
                decoded, actual_size, info = decode_custom_dib(resource.data, shared_palette)
                path = OUTPUT / "decoded" / "images" / f"dat_{resource.stem}.png"
                path.parent.mkdir(parents=True, exist_ok=True)
                decoded.save(path)
                entry.update({"actual_size": actual_size, "image": info})
                entry["decoded_paths"].append(str(path.relative_to(ROOT)))
            elif resource.type_name == "PAL":
                entry["actual_size"] = 768
                entry["decoded_paths"].extend(save_palette(shared_palette, resource.resource_id))
            elif resource.type_name == "WAV":
                if resource.data[:4] != b"RIFF" or resource.data[8:12] != b"WAVE":
                    raise ValueError("WAV resource lacks RIFF/WAVE signature")
                actual_size = u32(resource.data, 4) + 8
                if actual_size > len(resource.data):
                    raise ValueError("truncated RIFF resource")
                path = OUTPUT / "decoded" / "audio" / f"wav_{resource.stem}.wav"
                path.parent.mkdir(parents=True, exist_ok=True)
                path.write_bytes(resource.data[:actual_size])
                entry["actual_size"] = actual_size
                entry["audio"] = wav_info(resource.data[:actual_size])
                entry["decoded_paths"].append(str(path.relative_to(ROOT)))
            elif resource.type_name == "BITMAP":
                bmp, actual_size, info = wrap_bitmap_dib(resource.data)
                png_path = OUTPUT / "decoded" / "images" / f"bitmap_{resource.stem}.png"
                bmp_path = png_path.with_suffix(".bmp")
                png_path.parent.mkdir(parents=True, exist_ok=True)
                bmp_path.write_bytes(bmp)
                with Image.open(bmp_path) as decoded:
                    decoded.save(png_path)
                entry.update({"actual_size": actual_size, "image": info})
                entry["decoded_paths"].extend(
                    [str(bmp_path.relative_to(ROOT)), str(png_path.relative_to(ROOT))]
                )
            elif resource.type_name == "ICON":
                ico, actual_size, info = decode_icon(resource.data)
                ico_path = OUTPUT / "decoded" / "icons" / f"icon_{resource.stem}.ico"
                png_path = ico_path.with_suffix(".png")
                ico_path.parent.mkdir(parents=True, exist_ok=True)
                ico_path.write_bytes(ico)
                with Image.open(ico_path) as decoded:
                    decoded.convert("RGBA").save(png_path)
                entry.update({"actual_size": actual_size, "image": info})
                entry["decoded_paths"].extend(
                    [str(ico_path.relative_to(ROOT)), str(png_path.relative_to(ROOT))]
                )
            elif resource.type_name == "DIALOG":
                dialog, actual_size = decode_dialog(resource.data)
                path = OUTPUT / "decoded" / "dialogs" / f"dialog_{resource.stem}.json"
                path.parent.mkdir(parents=True, exist_ok=True)
                path.write_text(json.dumps(dialog, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
                entry["actual_size"] = actual_size
                entry["decoded_paths"].append(str(path.relative_to(ROOT)))
            elif resource.type_name == "VERSION":
                decoded = decode_version(resource.data)
                path = OUTPUT / "decoded" / "metadata" / f"version_{resource.stem}.json"
                path.parent.mkdir(parents=True, exist_ok=True)
                path.write_text(json.dumps(decoded, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
                entry["actual_size"] = u16(resource.data, 0)
                entry["decoded_paths"].append(str(path.relative_to(ROOT)))
            elif resource.type_name == "GROUP_ICON":
                count = u16(resource.data, 4)
                entry["actual_size"] = 6 + count * 14
        except Exception as error:
            message = f"{resource.type_name}/{resource.resource_id}: {error}"
            entry["decode_error"] = str(error)
            extraction_errors.append(message)
        manifest_entries.append(entry)

    manifest = {
        "source": str(SOURCE.relative_to(ROOT)),
        "source_sha256": hashlib.sha256(image).hexdigest(),
        "resource_count": len(resources),
        "ne": ne_metadata,
        "resources": manifest_entries,
        "decode_errors": extraction_errors,
    }
    OUTPUT.mkdir(parents=True, exist_ok=True)
    (OUTPUT / "manifest.json").write_text(
        json.dumps(manifest, indent=2, ensure_ascii=False) + "\n", encoding="utf-8"
    )
    if extraction_errors:
        raise SystemExit("resource extraction completed with decode errors:\n" + "\n".join(extraction_errors))
    print(f"Extracted {len(resources)} resources to {OUTPUT}")


if __name__ == "__main__":
    main()