"""Minimal RGBA / indexed PNG reader and writer.

Only what the theme's assets need: 8-bit truecolour-alpha (type 6) and indexed
(type 3) images, non-interlaced. Written against zlib from the stdlib so the
theme can be rebuilt without Pillow or ImageMagick installed.
"""

import struct
import zlib


def _chunks(data):
    assert data[:8] == b"\x89PNG\r\n\x1a\n", "not a PNG"
    i = 8
    out = []
    while i < len(data):
        (ln,) = struct.unpack(">I", data[i:i + 4])
        out.append([data[i + 4:i + 8], data[i + 8:i + 8 + ln]])
        i += 8 + ln + 4
    return out


def _pack(chunks):
    out = bytearray(b"\x89PNG\r\n\x1a\n")
    for typ, body in chunks:
        out += struct.pack(">I", len(body)) + typ + body
        out += struct.pack(">I", zlib.crc32(typ + body) & 0xFFFFFFFF)
    return bytes(out)


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
    return b if pb <= pc else c


def _unfilter(raw, width, height, bpp):
    """Reverse the per-scanline filters, returning flat sample bytes."""
    stride = width * bpp
    out = bytearray(stride * height)
    pos = 0
    for y in range(height):
        ft = raw[pos]
        pos += 1
        line = bytearray(raw[pos:pos + stride])
        pos += stride
        base = y * stride
        prev = base - stride
        if ft == 0:
            pass
        elif ft == 1:
            for x in range(bpp, stride):
                line[x] = (line[x] + line[x - bpp]) & 0xFF
        elif ft == 2:
            if y:
                for x in range(stride):
                    line[x] = (line[x] + out[prev + x]) & 0xFF
        elif ft == 3:
            for x in range(stride):
                a = line[x - bpp] if x >= bpp else 0
                b = out[prev + x] if y else 0
                line[x] = (line[x] + ((a + b) >> 1)) & 0xFF
        elif ft == 4:
            for x in range(stride):
                a = line[x - bpp] if x >= bpp else 0
                b = out[prev + x] if y else 0
                c = out[prev + x - bpp] if (y and x >= bpp) else 0
                line[x] = (line[x] + _paeth(a, b, c)) & 0xFF
        else:
            raise ValueError("unknown filter %d" % ft)
        out[base:base + stride] = line
    return out


def _filter_up(flat, width, height, bpp):
    """Re-encode with the Up filter; cheap and compresses these assets well."""
    stride = width * bpp
    out = bytearray()
    for y in range(height):
        base = y * stride
        line = flat[base:base + stride]
        if y == 0:
            out.append(0)
            out += line
        else:
            prev = flat[base - stride:base]
            out.append(2)
            out += bytes((line[x] - prev[x]) & 0xFF for x in range(stride))
    return bytes(out)


class Png:
    def __init__(self, path):
        self.path = path
        self.chunks = _chunks(open(path, "rb").read())
        ihdr = dict((t, b) for t, b in self.chunks)[b"IHDR"]
        (self.width, self.height, self.depth, self.color_type,
         _, _, self.interlace) = struct.unpack(">IIBBBBB", ihdr)
        assert self.interlace == 0, "interlaced PNGs unsupported"

    def _get(self, typ):
        for t, b in self.chunks:
            if t == typ:
                return b
        return None

    def _set(self, typ, body):
        for c in self.chunks:
            if c[0] == typ:
                c[1] = body
                return
        raise KeyError(typ)

    # -- indexed ------------------------------------------------------------
    def palette(self):
        p = self._get(b"PLTE")
        return [tuple(p[i:i + 3]) for i in range(0, len(p), 3)]

    def set_palette(self, pal):
        self._set(b"PLTE", b"".join(bytes(c) for c in pal))

    # -- truecolour+alpha ---------------------------------------------------
    def rgba(self):
        assert self.color_type == 6 and self.depth == 8, "expected 8-bit RGBA"
        idat = b"".join(b for t, b in self.chunks if t == b"IDAT")
        flat = _unfilter(zlib.decompress(idat), self.width, self.height, 4)
        return [tuple(flat[i:i + 4]) for i in range(0, len(flat), 4)]

    def set_rgba(self, pixels):
        flat = bytearray()
        for px in pixels:
            flat += bytes(px)
        raw = _filter_up(flat, self.width, self.height, 4)
        body = zlib.compress(raw, 9)
        # collapse any split IDATs into one
        self.chunks = [c for c in self.chunks if c[0] != b"IDAT"]
        end = next(i for i, c in enumerate(self.chunks) if c[0] == b"IEND")
        self.chunks.insert(end, [b"IDAT", body])

    def save(self, path=None):
        open(path or self.path, "wb").write(_pack(self.chunks))