#!/usr/bin/env python3
"""Dependency-free PNG read/write plus helpers for the Kissaten art pipeline.

No Pillow on this machine, so decode/encode are hand-rolled. Supports the
colour types PixelLab and friends actually emit: 2 (RGB), 6 (RGBA), 3
(palette), 0 (grey), 4 (grey+alpha), 8 bits per channel.
"""
import zlib, struct, binascii, sys


def read_png(path):
    """-> (w, h, pixels) where pixels is a list of rows of (r,g,b,a) tuples."""
    d = open(path, "rb").read()
    assert d[:8] == b"\x89PNG\r\n\x1a\x0a", f"{path}: not a PNG"
    pos, idat, plte, trns = 8, b"", None, None
    w = h = depth = ctype = None
    while pos < len(d):
        ln = struct.unpack(">I", d[pos:pos + 4])[0]
        typ = d[pos + 4:pos + 8]
        body = d[pos + 8:pos + 8 + ln]
        if typ == b"IHDR":
            w, h, depth, ctype = struct.unpack(">IIBB", body[:10])
        elif typ == b"PLTE":
            plte = body
        elif typ == b"tRNS":
            trns = body
        elif typ == b"IDAT":
            idat += body
        elif typ == b"IEND":
            break
        pos += 12 + ln
    assert depth == 8, f"{path}: only 8-bit channels supported (got {depth})"

    nch = {0: 1, 2: 3, 3: 1, 4: 2, 6: 4}[ctype]
    raw = zlib.decompress(idat)
    stride = w * nch
    out, prev = [], bytearray(stride)
    p = 0
    for _ in range(h):
        f = raw[p]; p += 1
        line = bytearray(raw[p:p + stride]); p += stride
        for i in range(stride):
            a = line[i - nch] if i >= nch else 0
            b = prev[i]
            c = prev[i - nch] if i >= nch else 0
            if f == 1:   line[i] = (line[i] + a) & 255
            elif f == 2: line[i] = (line[i] + b) & 255
            elif f == 3: line[i] = (line[i] + (a + b) // 2) & 255
            elif f == 4:
                pa, pb, pc = abs(b - c), abs(a - c), abs(a + b - 2 * c)
                pr = a if (pa <= pb and pa <= pc) else (b if pb <= pc else c)
                line[i] = (line[i] + pr) & 255
        prev = line

        row = []
        for x in range(w):
            v = line[x * nch:(x + 1) * nch]
            if ctype == 2:   row.append((v[0], v[1], v[2], 255))
            elif ctype == 6: row.append((v[0], v[1], v[2], v[3]))
            elif ctype == 0: row.append((v[0], v[0], v[0], 255))
            elif ctype == 4: row.append((v[0], v[0], v[0], v[1]))
            else:
                i = v[0]
                a = trns[i] if (trns and i < len(trns)) else 255
                row.append((plte[i * 3], plte[i * 3 + 1], plte[i * 3 + 2], a))
        out.append(row)
    return w, h, out


def write_png(path, rows):
    h, w = len(rows), len(rows[0])
    raw = b"".join(b"\x00" + bytes(c for px in r for c in px[:3]) for r in rows)

    def chunk(t, b):
        c = t + b
        return struct.pack(">I", len(b)) + c + struct.pack(">I", binascii.crc32(c) & 0xffffffff)

    open(path, "wb").write(
        b"\x89PNG\r\n\x1a\n"
        + chunk(b"IHDR", struct.pack(">IIBBBBB", w, h, 8, 2, 0, 0, 0))
        + chunk(b"IDAT", zlib.compress(raw, 9))
        + chunk(b"IEND", b""))


def zoom(rows, n):
    """Nearest-neighbour upscale, so pixels stay square and countable."""
    return [[px for px in r for _ in range(n)] for r in rows for _ in range(n)]


# The game palette, index-aligned with g_Palette in kissaten.c
PALETTE = [
    (0x00, 0x00, 0x00), (0x24, 0x24, 0x24), (0x49, 0x24, 0x24), (0x92, 0x49, 0x24),
    (0xB6, 0x6D, 0x49), (0xFF, 0xB6, 0x49), (0xFF, 0xDB, 0x92), (0x6D, 0x49, 0x49),
    (0x49, 0x24, 0x92), (0x92, 0x49, 0xB6), (0xDB, 0x6D, 0x6D), (0xFF, 0x92, 0x49),
    (0x6D, 0x24, 0x00), (0x49, 0x6D, 0xB6), (0xFF, 0xB6, 0x92), (0xFF, 0xFF, 0xDB),
]
SKY = {8, 9, 10, 11}


if __name__ == "__main__":
    for path in sys.argv[1:]:
        w, h, rows = read_png(path)
        flat = [px for r in rows for px in r]
        cols = {}
        for px in flat:
            cols[px[:3]] = cols.get(px[:3], 0) + 1
        alphas = sorted({px[3] for px in flat})
        exact = sum(n for c, n in cols.items() if c in PALETTE)
        print(f"\n{path}  {w}x{h}")
        print(f"  unique colours : {len(cols)}")
        print(f"  alpha values   : {alphas if len(alphas) < 6 else str(alphas[:5]) + '...'}")
        print(f"  already in our palette: {len(cols) - len([c for c in cols if c not in PALETTE])}"
              f" of {len(cols)}  ({100*exact//len(flat)}% of pixels)")
        print("  top 8 colours:")
        for c, n in sorted(cols.items(), key=lambda kv: -kv[1])[:8]:
            tag = "  <- in palette" if c in PALETTE else ""
            print(f"    #{c[0]:02X}{c[1]:02X}{c[2]:02X}  {n:5d} px  {100*n//len(flat):3d}%{tag}")