#!/usr/bin/env python3

import argparse
import json
from pathlib import Path

from PIL import Image


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Convert a JSON RGBA matrix back into a PNG image."
    )
    parser.add_argument("input_json", type=Path, help="Source JSON path")
    parser.add_argument("output_png", type=Path, help="Destination PNG path")
    return parser.parse_args()


def load_matrix(json_path: Path) -> tuple[int, int, list[tuple[int, int, int, int]]]:
    data = json.loads(json_path.read_text(encoding="utf-8"))
    width = int(data["width"])
    height = int(data["height"])
    rows = data["pixels"]

    if len(rows) != height:
        raise ValueError(f"Expected {height} rows, found {len(rows)}")

    flat_pixels: list[tuple[int, int, int, int]] = []
    for row_index, row in enumerate(rows):
        if len(row) != width:
            raise ValueError(
                f"Expected row {row_index} to contain {width} pixels, found {len(row)}"
            )
        for pixel in row:
            if len(pixel) != 4:
                raise ValueError("Each pixel must contain 4 RGBA channels")
            flat_pixels.append(tuple(int(channel) for channel in pixel))

    return width, height, flat_pixels


def main() -> None:
    args = parse_args()
    width, height, pixels = load_matrix(args.input_json)
    image = Image.new("RGBA", (width, height))
    image.putdata(pixels)
    args.output_png.parent.mkdir(parents=True, exist_ok=True)
    image.save(args.output_png, format="PNG")
    print(f"Wrote {args.output_png} with {width}x{height} RGBA pixels")


if __name__ == "__main__":
    main()