#!/usr/bin/env python3
"""Decode Compile .cns images to PNG.

This implements the LZ-style decoder used by cns110.exe and writes a
plain RGB PNG using only the Python standard library.
"""
from __future__ import annotations

import argparse
import struct
import zlib
from pathlib import Path


def decompress_cns(data: bytes) -> bytes:
    """Decompress a CNS byte stream.

    The stream is a command byte followed by literal bytes and/or a
    backwards copy. A zero command terminates the stream.
    """
    src = 0
    out = bytearray()

    def literal(count: int) -> None:
        nonlocal src
        if count < 0 or src + count > len(data):
            raise ValueError("literal run exceeds input")
        out.extend(data[src : src + count])
        src += count

    def copy_back(offset: int, count: int) -> None:
        if offset <= 0 or offset > len(out):
            raise ValueError(f"bad back-reference offset {offset} at output {len(out)}")
        start = len(out) - offset
        for i in range(count):
            out.append(out[start + i])

    while True:
        if src >= len(data):
            raise ValueError("missing CNS end marker")
        code = data[src]
        src += 1

        if code == 0:
            return bytes(out)

        if code >= 0x80:
            literal((code & 0x70) >> 4)
            if src >= len(data):
                raise ValueError("truncated one-byte back-reference")
            count = (code & 0x0F) + 2
            offset = data[src]
            src += 1
            copy_back(offset, count)
        elif code >= 0x60:
            if src >= len(data):
                raise ValueError("truncated long literal")
            count = ((code & 0x1F) << 8) + data[src]
            src += 1
            literal(count)
        elif code >= 0x40:
            literal(code & 0x1F)
        elif code >= 0x30:
            if src + 3 > len(data):
                raise ValueError("truncated long two-byte back-reference")
            count = ((code & 0x0F) << 8) + data[src]
            offset = data[src + 1] | (data[src + 2] << 8)
            src += 3
            copy_back(offset, count)
        elif code >= 0x20:
            if src + 2 > len(data):
                raise ValueError("truncated long one-byte back-reference")
            count = ((code & 0x0F) << 8) + data[src]
            offset = data[src + 1]
            src += 2
            copy_back(offset, count)
        elif code >= 0x10:
            if src + 2 > len(data):
                raise ValueError("truncated short two-byte back-reference")
            count = (code & 0x0F) + 2
            offset = data[src] | (data[src + 1] << 8)
            src += 2
            copy_back(offset, count)
        else:
            if src >= len(data):
                raise ValueError("truncated short one-byte back-reference")
            count = (code & 0x0F) + 2
            offset = data[src]
            src += 1
            copy_back(offset, count)


def parse_image(decoded: bytes) -> tuple[int, int, list[tuple[int, int, int]], bytes, int]:
    if len(decoded) < 8:
        raise ValueError("decoded CNS is too short")
    _unknown, width, height, palette_last = struct.unpack_from("<HHHH", decoded, 0)
    palette_count = palette_last + 1
    bpp = 8 if palette_count > 16 else 4
    palette_end = 8 + palette_count * 4
    if palette_end > len(decoded):
        raise ValueError("palette exceeds decoded data")

    palette: list[tuple[int, int, int]] = []
    for i in range(palette_count):
        b, g, r, _ = decoded[8 + i * 4 : 12 + i * 4]
        palette.append((r, g, b))
    pixels = decoded[palette_end:]
    return width, height, palette, pixels, bpp


def indices_to_rgb(
    width: int, height: int, palette: list[tuple[int, int, int]], pixels: bytes, bpp: int
) -> bytes:
    stride = ((width * bpp + 31) // 32) * 4
    expected = stride * height
    if len(pixels) < expected:
        raise ValueError(f"pixel data too short: got {len(pixels)}, expected {expected}")

    rgb = bytearray()
    for y in range(height):
        row = pixels[(height - 1 - y) * stride : (height - y) * stride]
        if bpp == 8:
            indices = row[:width]
        elif bpp == 4:
            indices = bytearray()
            for byte in row[: (width + 1) // 2]:
                indices.append(byte >> 4)
                if len(indices) < width:
                    indices.append(byte & 0x0F)
        else:
            raise ValueError(f"unsupported bit depth {bpp}")

        for idx in indices:
            if idx >= len(palette):
                rgb.extend((255, 0, 255))
            else:
                rgb.extend(palette[idx])
    return bytes(rgb)


def indices_to_rgba(
    width: int,
    height: int,
    palette: list[tuple[int, int, int]],
    pixels: bytes,
    bpp: int,
    transparent_index: int = 0,
) -> bytes:
    stride = ((width * bpp + 31) // 32) * 4
    expected = stride * height
    if len(pixels) < expected:
        raise ValueError(f"pixel data too short: got {len(pixels)}, expected {expected}")

    rgba = bytearray()
    for y in range(height):
        row = pixels[(height - 1 - y) * stride : (height - y) * stride]
        if bpp == 8:
            indices = row[:width]
        elif bpp == 4:
            indices = bytearray()
            for byte in row[: (width + 1) // 2]:
                indices.append(byte >> 4)
                if len(indices) < width:
                    indices.append(byte & 0x0F)
        else:
            raise ValueError(f"unsupported bit depth {bpp}")

        for idx in indices:
            if idx >= len(palette):
                rgba.extend((255, 0, 255, 255))
            else:
                r, g, b = palette[idx]
                alpha = 0 if idx == transparent_index else 255
                rgba.extend((r, g, b, alpha))
    return bytes(rgba)


def write_png(path: Path, width: int, height: int, rgb: bytes, transparent_rgb: tuple[int, int, int] | None = None) -> None:
    def chunk(kind: bytes, payload: bytes) -> bytes:
        return (
            struct.pack(">I", len(payload))
            + kind
            + payload
            + struct.pack(">I", zlib.crc32(kind + payload) & 0xFFFFFFFF)
        )

    scanlines = bytearray()
    channels = 4 if transparent_rgb is not None else 3
    row_len = width * 3
    for y in range(height):
        scanlines.append(0)
        row = rgb[y * row_len : (y + 1) * row_len]
        if transparent_rgb is None:
            scanlines.extend(row)
        else:
            for i in range(0, len(row), 3):
                r, g, b = row[i : i + 3]
                alpha = 0 if (r, g, b) == transparent_rgb else 255
                scanlines.extend((r, g, b, alpha))

    png = bytearray(b"\x89PNG\r\n\x1a\n")
    color_type = 6 if channels == 4 else 2
    png.extend(chunk(b"IHDR", struct.pack(">IIBBBBB", width, height, 8, color_type, 0, 0, 0)))
    png.extend(chunk(b"IDAT", zlib.compress(bytes(scanlines), 9)))
    png.extend(chunk(b"IEND", b""))
    path.write_bytes(png)


def write_rgba_png(path: Path, width: int, height: int, rgba: bytes) -> None:
    def chunk(kind: bytes, payload: bytes) -> bytes:
        return (
            struct.pack(">I", len(payload))
            + kind
            + payload
            + struct.pack(">I", zlib.crc32(kind + payload) & 0xFFFFFFFF)
        )

    scanlines = bytearray()
    row_len = width * 4
    for y in range(height):
        scanlines.append(0)
        scanlines.extend(rgba[y * row_len : (y + 1) * row_len])

    png = bytearray(b"\x89PNG\r\n\x1a\n")
    png.extend(chunk(b"IHDR", struct.pack(">IIBBBBB", width, height, 8, 6, 0, 0, 0)))
    png.extend(chunk(b"IDAT", zlib.compress(bytes(scanlines), 9)))
    png.extend(chunk(b"IEND", b""))
    path.write_bytes(png)


def decode_file(src: Path, dst: Path, transparent_palette0: bool = False) -> None:
    decoded = decompress_cns(src.read_bytes())
    width, height, palette, pixels, bpp = parse_image(decoded)
    dst.parent.mkdir(parents=True, exist_ok=True)
    if transparent_palette0:
        rgba = indices_to_rgba(width, height, palette, pixels, bpp, transparent_index=0)
        write_rgba_png(dst, width, height, rgba)
    else:
        rgb = indices_to_rgb(width, height, palette, pixels, bpp)
        write_png(dst, width, height, rgb)
    print(
        f"{src} -> {dst} ({width}x{height}, {bpp}bpp, "
        f"{len(palette)} colors, decoded {len(decoded)} bytes)"
    )


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("inputs", nargs="+", type=Path)
    parser.add_argument("-o", "--out-dir", type=Path, default=Path("out"))
    args = parser.parse_args()

    for src in args.inputs:
        dst = args.out_dir / (src.stem + ".png")
        decode_file(src, dst)


if __name__ == "__main__":
    main()