#!/usr/bin/env python3
"""Turn docs/branding/*.png into a C header the kernel can draw.

Why this exists rather than a PNG decoder in the kernel: the boot logo is one
fixed image that changes about never, so decoding it at runtime would mean
carrying an inflate implementation in the kernel forever to answer a question
that is already known at build time. Precedent is include/gfx/font.h, which is the
same trade for the same reason.

Output is an 8-bit coverage mask, not colour. The logo is a single-colour
silhouette (a monogram in a ring), so what matters is *how much* of each pixel
is ink -- which lets the kernel composite it in whatever the theme's accent
happens to be, anti-aliased, at one byte per pixel. Storing RGB would be four
times the size and would hardcode a colour into the kernel image.

    python3 tools/python/mklogo.py docs/branding/shift-mask.png include/gfx/logo.h 128
"""

import sys, zlib, struct


def read_png(path):
    d = open(path, 'rb').read()
    if d[:8] != b'\x89PNG\r\n\x1a\n':
        raise SystemExit(f'{path}: not a PNG')
    pos, idat, w, h, bd, ct = 8, b'', 0, 0, 0, 0
    while pos < len(d):
        ln, typ = struct.unpack('>I4s', d[pos:pos + 8])
        body = d[pos + 8:pos + 8 + ln]
        if typ == b'IHDR':
            w, h, bd, ct = struct.unpack('>IIBB', body[:10])
        elif typ == b'IDAT':
            idat += body
        elif typ == b'IEND':
            break
        pos += 12 + ln
    if bd != 8 or ct not in (2, 6):
        raise SystemExit(f'{path}: need 8-bit RGB or RGBA, got depth={bd} type={ct}')
    nch = 4 if ct == 6 else 3
    raw = zlib.decompress(idat)

    # Undo PNG's per-scanline filters. Each row is prefixed with a filter byte
    # and predicted from the row above and the pixel to the left; without this
    # step the data is not an image at all.
    stride = w * nch
    out = bytearray(h * stride)
    prev = bytearray(stride)
    p = 0
    for y in range(h):
        f = raw[p]; p += 1
        cur = bytearray(raw[p:p + stride]); p += stride
        if f == 1:
            for i in range(nch, stride):
                cur[i] = (cur[i] + cur[i - nch]) & 0xFF
        elif f == 2:
            for i in range(stride):
                cur[i] = (cur[i] + prev[i]) & 0xFF
        elif f == 3:
            for i in range(stride):
                a = cur[i - nch] if i >= nch else 0
                cur[i] = (cur[i] + ((a + prev[i]) >> 1)) & 0xFF
        elif f == 4:
            for i in range(stride):
                a = cur[i - nch] if i >= nch else 0
                b = prev[i]
                c = prev[i - nch] if i >= nch else 0
                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)
                cur[i] = (cur[i] + pr) & 0xFF
        elif f != 0:
            raise SystemExit(f'unknown PNG filter {f}')
        out[y * stride:(y + 1) * stride] = cur
        prev = cur
    return w, h, nch, out


def ink(px, nch, x, y, w):
    """How much ink is at (x, y), 0..255.

    The source is a dark mark on white with an alpha channel. Ink is therefore
    'opaque and not white': darkness weighted by alpha. Using luminance alone
    would make the transparent margin count as black.
    """
    i = (y * w + x) * nch
    r, g, b = px[i], px[i + 1], px[i + 2]
    a = px[i + 3] if nch == 4 else 255
    lum = (r * 299 + g * 587 + b * 114) // 1000
    return ((255 - lum) * a) // 255


def main():
    if len(sys.argv) < 3:
        raise SystemExit(__doc__)
    src, dst = sys.argv[1], sys.argv[2]
    size = int(sys.argv[3]) if len(sys.argv) > 3 else 128

    w, h, nch, px = read_png(src)

    # Crop to the mark. These files are 1920x1080 with the logo floating in a
    # large empty field; scaling that whole canvas down would waste most of the
    # output on margin and leave the mark tiny.
    x0, y0, x1, y1 = w, h, -1, -1
    for y in range(h):
        for x in range(w):
            if ink(px, nch, x, y, w) > 24:
                if x < x0: x0 = x
                if y < y0: y0 = y
                if x > x1: x1 = x
                if y > y1: y1 = y
    if x1 < 0:
        raise SystemExit(f'{src}: no non-white pixels found')

    # Square the crop so the aspect ratio survives the scale.
    cw, ch = x1 - x0 + 1, y1 - y0 + 1
    side = max(cw, ch)
    x0 -= (side - cw) // 2
    y0 -= (side - ch) // 2

    # Box filter. Averaging every source pixel that lands in an output cell is
    # what keeps the ring's edge smooth; nearest-neighbour at this reduction
    # ratio (roughly 8:1) would break it into dashes.
    mask = bytearray(size * size)
    for oy in range(size):
        sy0 = y0 + oy * side // size
        sy1 = max(sy0 + 1, y0 + (oy + 1) * side // size)
        for ox in range(size):
            sx0 = x0 + ox * side // size
            sx1 = max(sx0 + 1, x0 + (ox + 1) * side // size)
            tot = n = 0
            for sy in range(sy0, sy1):
                if sy < 0 or sy >= h: continue
                for sx in range(sx0, sx1):
                    if sx < 0 or sx >= w: continue
                    tot += ink(px, nch, sx, sy, w); n += 1
            mask[oy * size + ox] = tot // n if n else 0

    # Normalise to full strength. The source art is a mid-grey gradient on
    # white, so its darkest pixel is only about 60% "ink" -- composited as-is
    # against a near-black background the mark comes out muddy. Scaling by the
    # peak keeps the gradient's internal shape (the relative values are
    # unchanged) while letting the brightest part of the logo actually reach the
    # accent colour.
    peak = max(mask)
    if peak:
        for i, v in enumerate(mask):
            mask[i] = min(255, v * 255 // peak)

    # Pack two pixels per byte. This header is included by the *bootloader* as
    # well as the kernel, and stage2 has to fit between 0x7E00 and the 0x10000
    # kernel-search buffer -- so the mask's size is a boot-time constraint, not
    # just a matter of taste. 16 grey levels is visually indistinguishable from
    # 256 for a single-colour mark, and halves the cost.
    packed = bytearray(len(mask) // 2)
    for i in range(0, len(mask), 2):
        packed[i // 2] = ((mask[i] // 17) << 4) | (mask[i + 1] // 17)

    with open(dst, 'w', newline='\n') as f:
        f.write('#pragma once\n\n')
        f.write('// GENERATED by tools/mklogo.py -- do not edit. Run `make logo`.\n')
        f.write(f'//   source: {src}\n')
        f.write('//\n')
        f.write('// An ink-coverage mask, not colour: 0 is background and full scale is\n')
        f.write('// solid logo. Whoever draws it picks the colour, so the palette stays in\n')
        f.write('// include/gfx/theme.h and can change without regenerating anything.\n')
        f.write('//\n')
        f.write('// Two pixels per byte, high nibble first. Included by boot/stage2main.c\n')
        f.write('// as well as kernel/dev/vga.c, and stage2 must fit below the 0x10000\n')
        f.write('// kernel-search buffer -- hence 4 bits rather than 8.\n\n')
        f.write(f'#define LOGO_W {size}\n#define LOGO_H {size}\n\n')
        f.write('// Ink at linear pixel index i, 0..255.\n')
        f.write('#define LOGO_INK(i) ((unsigned char)((((i) & 1) '
                '? (logo_mask[(i) >> 1] & 0xF) \\\n')
        f.write('                                      : (logo_mask[(i) >> 1] >> 4)) * 17))\n\n')
        f.write('static const unsigned char logo_mask[LOGO_W * LOGO_H / 2] = {\n')
        for i in range(0, len(packed), 16):
            f.write('  ' + ','.join('0x%02x' % b for b in packed[i:i + 16]) + ',\n')
        f.write('};\n')

    nz = sum(1 for b in mask if b)
    print(f'{dst}: {size}x{size}, cropped from {cw}x{ch} at ({x0},{y0}), '
          f'{nz} inked pixels, {len(packed)} bytes packed (4bpp)')


main()