"""Minimal pure-Python PNG decoder (stdlib only: zlib + struct).

Supports 8-bit, non-interlaced grayscale / grayscale+alpha / RGB / RGBA PNGs -
which covers ordinary exported heightmap/relief images. Deliberately does not
depend on Pillow or any other package, since Unreal's embedded Python has no
pip packages installed by default.
"""
import struct, zlib

_CHANNELS_BY_COLOR_TYPE = {0: 1, 2: 3, 4: 2, 6: 4}  # 3 (palette) is unsupported


def _paeth(a, b, c):
    p = a + b - c
    pa, pb, pc = abs(p - a), abs(p - b), abs(p - c)
    if pa <= pb and pa <= pc:
        return a
    if pb <= pc:
        return b
    return c


def decode_png(path):
    """Returns (width, height, channels, pixels) where pixels is a flat
    bytearray of width*height*channels bytes, row-major, top row first."""
    with open(path, "rb") as f:
        data = f.read()
    if data[:8] != b"\x89PNG\r\n\x1a\n":
        raise ValueError(f"Not a PNG file: {path}")

    width = height = bit_depth = color_type = interlace = None
    idat = bytearray()
    pos = 8
    while pos < len(data):
        length = struct.unpack(">I", data[pos:pos + 4])[0]
        ctype = data[pos + 4:pos + 8]
        chunk = data[pos + 8:pos + 8 + length]
        if ctype == b"IHDR":
            width, height, bit_depth, color_type, _comp, _filt, interlace = struct.unpack(">IIBBBBB", chunk)
        elif ctype == b"IDAT":
            idat.extend(chunk)
        elif ctype == b"IEND":
            break
        pos += 8 + length + 4  # length + type + data + crc

    if bit_depth != 8:
        raise ValueError(f"Only 8-bit PNGs are supported (got bit depth {bit_depth}) - re-export as 8-bit: {path}")
    if interlace != 0:
        raise ValueError(f"Interlaced PNGs are not supported - re-export as non-interlaced: {path}")
    if color_type not in _CHANNELS_BY_COLOR_TYPE:
        raise ValueError(f"Unsupported PNG color type {color_type} (palette?) - re-export as grayscale/RGB/RGBA: {path}")
    channels = _CHANNELS_BY_COLOR_TYPE[color_type]

    raw = zlib.decompress(bytes(idat))
    stride = width * channels
    pixels = bytearray(width * height * channels)

    prev_row = bytearray(stride)
    row_start = 0
    for y in range(height):
        filter_type = raw[row_start]
        row = bytearray(raw[row_start + 1: row_start + 1 + stride])
        row_start += 1 + stride

        if filter_type == 1:  # Sub
            for i in range(channels, stride):
                row[i] = (row[i] + row[i - channels]) & 0xFF
        elif filter_type == 2:  # Up
            for i in range(stride):
                row[i] = (row[i] + prev_row[i]) & 0xFF
        elif filter_type == 3:  # Average
            for i in range(stride):
                a = row[i - channels] if i >= channels else 0
                row[i] = (row[i] + ((a + prev_row[i]) >> 1)) & 0xFF
        elif filter_type == 4:  # Paeth
            for i in range(stride):
                a = row[i - channels] if i >= channels else 0
                b = prev_row[i]
                c = prev_row[i - channels] if i >= channels else 0
                row[i] = (row[i] + _paeth(a, b, c)) & 0xFF
        # filter_type == 0 (None): no-op

        pixels[y * stride:(y + 1) * stride] = row
        prev_row = row

    return width, height, channels, pixels