#!/usr/bin/env python3

"""
Convert between custom bottom‑up raw files (.r8g8b8a8 / .r8b8g8 / .r8 / .a8)
and normal PNG using Pillow.

Format
------
uint32 width  (little‑endian)
uint32 height (little‑endian)
pixel data rows bottom‑first:
  * .r8g8b8a8 : R G B A   (4× uint8)
  * .r8b8g8   : R G B     (3× uint8)
  * .r8       : R         (1× uint8 <- grayscale)
  * .a8       :       A   (can convert from it, but can't convert to it)

CLI
---
  raw -> png :  python raw_png_conv.py to_png  input.raw  output.png
  png -> raw :  python raw_png_conv.py to_raw  input.png  output.raw
"""
import argparse
import struct
from pathlib import Path
from PIL import Image

# --------------------------------------------------------------------- #
# Helpers
# --------------------------------------------------------------------- #


RAW_FORMATS = {
    ".r8g8b8a8": {"pix_size": 4, "mode": "RGBA"},
    ".r8g8b8":   {"pix_size": 3, "mode": "RGB"},
    ".r8":       {"pix_size": 1, "mode": "L"},
    ".a8":       {"pix_size": 1, "mode": None},  # special-cased (alpha-only)
}


def get_fmt(path: Path):
    fmt = RAW_FORMATS.get(path.suffix.lower())
    if not fmt:
        raise ValueError(f"Unknown raw extension: {path.suffix}")
    return fmt


def read_header_and_data(path: Path, pix_size: int):
    with path.open("rb") as f:
        header = f.read(8)
        if len(header) != 8:
            raise ValueError("File too short for header")
        w, h = struct.unpack("<II", header)
        expected = w * h * pix_size
        data = f.read()
        if len(data) != expected:
            raise ValueError(
                f"Pixel data length mismatch: expected {expected}, got {len(data)}"
            )
    return w, h, data


def read_raw(path: Path) -> Image.Image:
    """Load any supported raw file -> Pillow Image."""
    spec = get_fmt(path)
    w, h, data = read_header_and_data(path, spec["pix_size"])

    row_len = w * spec["pix_size"]
    rows = [data[i : i + row_len] for i in range(0, len(data), row_len)]
    top_down = b"".join(reversed(rows))  # flip bottom‑up -> top‑down

    # Special handling for .a8 (alpha-only -> LA with black color)
    if path.suffix.lower() == ".a8":
        big = bytearray(4 * len(top_down))
        big[3::4] = top_down
        big = bytes(big)
        return Image.frombytes("RGBA", (w, h), big)

    # Normal cases can be constructed directly
    return Image.frombytes(spec["mode"], (w, h), top_down)


def write_raw(img: Image.Image, path: Path) -> None:
    """Write Pillow Image -> raw file chosen by path suffix."""
    ext = path.suffix.lower()
    spec = get_fmt(path)

    if ext == ".a8":
        raise ValueError("Don't convert to .a8 format")

    target_mode = spec["mode"]
    if img.mode != target_mode:
        img = img.convert(target_mode)

    w, h = img.size
    data = img.tobytes()
    row_len = w * spec["pix_size"]
    rows = [data[i : i + row_len] for i in range(0, len(data), row_len)]
    bottom_first = b"".join(reversed(rows))  # top‑down -> bottom‑up

    with path.open("wb") as f:
        f.write(struct.pack("<II", w, h))
        f.write(bottom_first)


# --------------------------------------------------------------------- #
# CLI
# --------------------------------------------------------------------- #

def to_png(src: Path, dst: Path):
    read_raw(src).save(dst, "PNG")


def to_raw(src: Path, dst: Path):
    write_raw(Image.open(src), dst)


def main():
    ap = argparse.ArgumentParser(description="Convert raw <-> PNG")
    sub = ap.add_subparsers(dest="cmd", required=True)
    p_png = sub.add_parser("to_png", help="raw -> png")
    p_png.add_argument("src", type=Path)
    p_png.add_argument("dst", type=Path)
    p_raw = sub.add_parser("to_raw", help="png -> raw")
    p_raw.add_argument("src", type=Path)
    p_raw.add_argument("dst", type=Path)
    args = ap.parse_args()

    if args.cmd == "to_png":
        to_png(args.src, args.dst)
    else:
        to_raw(args.src, args.dst)


if __name__ == "__main__":
    main()
