#!/usr/bin/env python3
import os
import re
import struct
import sys
import zlib


def write_png_rgb(path, width, height, rgb):
    def chunk(tag, data):
        body = tag + data
        return struct.pack(">I", len(data)) + body + struct.pack(">I", zlib.crc32(body) & 0xFFFFFFFF)

    rows = bytearray()
    stride = width * 3
    for y in range(height):
        rows.append(0)
        start = y * stride
        rows.extend(rgb[start:start + stride])

    ihdr = struct.pack(">IIBBBBB", width, height, 8, 2, 0, 0, 0)
    with open(path, "wb") as f:
        f.write(b"\x89PNG\r\n\x1a\n")
        f.write(chunk(b"IHDR", ihdr))
        f.write(chunk(b"IDAT", zlib.compress(bytes(rows), 6)))
        f.write(chunk(b"IEND", b""))


def write_png_gray(path, width, height, gray):
    def chunk(tag, data):
        body = tag + data
        return struct.pack(">I", len(data)) + body + struct.pack(">I", zlib.crc32(body) & 0xFFFFFFFF)

    rows = bytearray()
    for y in range(height):
        rows.append(0)
        start = y * width
        rows.extend(gray[start:start + width])

    ihdr = struct.pack(">IIBBBBB", width, height, 8, 0, 0, 0, 0)
    with open(path, "wb") as f:
        f.write(b"\x89PNG\r\n\x1a\n")
        f.write(chunk(b"IHDR", ihdr))
        f.write(chunk(b"IDAT", zlib.compress(bytes(rows), 6)))
        f.write(chunk(b"IEND", b""))


def clamp(v):
    if v < 0:
        return 0
    if v > 255:
        return 255
    return v


def yuv420sp_to_rgb(y_plane, uv_plane, width, height, nv21=False):
    rgb = bytearray(width * height * 3)
    for y in range(height):
        y_row = y * width
        uv_row = (y // 2) * width
        out_row = y_row * 3
        for x in range(width):
            yy = y_plane[y_row + x]
            uv_index = uv_row + (x & ~1)
            if nv21:
                v = uv_plane[uv_index] - 128
                u = uv_plane[uv_index + 1] - 128
            else:
                u = uv_plane[uv_index] - 128
                v = uv_plane[uv_index + 1] - 128

            c = yy - 16
            if c < 0:
                c = 0
            r = (298 * c + 409 * v + 128) >> 8
            g = (298 * c - 100 * u - 208 * v + 128) >> 8
            b = (298 * c + 516 * u + 128) >> 8

            out = out_row + x * 3
            rgb[out] = clamp(r)
            rgb[out + 1] = clamp(g)
            rgb[out + 2] = clamp(b)
    return rgb


def main():
    dump_dir = sys.argv[1] if len(sys.argv) > 1 else "/home/user/dump"
    out_dir = os.path.join(dump_dir, "png")
    os.makedirs(out_dir, exist_ok=True)

    pattern = re.compile(r"^(\d{8})_(\d+)X(\d+)_y\.raw$")
    converted = 0
    for name in sorted(os.listdir(dump_dir)):
        m = pattern.match(name)
        if not m:
            continue
        frame, width_s, height_s = m.groups()
        width = int(width_s)
        height = int(height_s)
        y_path = os.path.join(dump_dir, name)
        uv_path = os.path.join(dump_dir, f"{frame}_{width}X{height}_uv.raw")
        if not os.path.exists(uv_path):
            print(f"skip {frame}: missing {uv_path}")
            continue

        with open(y_path, "rb") as f:
            y_plane = f.read()
        with open(uv_path, "rb") as f:
            uv_plane = f.read()

        expected_y = width * height
        expected_uv = width * height // 2
        if len(y_plane) != expected_y or len(uv_plane) != expected_uv:
            print(f"skip {frame}: bad size y={len(y_plane)} uv={len(uv_plane)} expected {expected_y}/{expected_uv}")
            continue

        base = os.path.join(out_dir, f"{frame}_{width}x{height}")
        write_png_gray(base + "_gray.png", width, height, y_plane)
        write_png_rgb(base + "_nv12.png", width, height, yuv420sp_to_rgb(y_plane, uv_plane, width, height, nv21=False))
        write_png_rgb(base + "_nv21.png", width, height, yuv420sp_to_rgb(y_plane, uv_plane, width, height, nv21=True))
        converted += 1
        print(f"converted frame {frame}: gray/nv12/nv21")

    print(f"output: {out_dir}")
    print(f"frames: {converted}")


if __name__ == "__main__":
    main()