#!/usr/bin/env python3
"""Generate the SUB-OS boot logo asset from a source image.

Emits init/logo_data.c: a 4-bit palette-indexed bitmap plus the ARGB palette
and an ASCII-art rendition used before the framebuffer is up.

    python3 scripts/mklogo.py ~/Pictures/logo_SUB.png

Re-run this whenever the source artwork changes; init/logo_data.c is generated
and should not be hand-edited.
"""

import sys
import os
from PIL import Image

SIZE       = 160          # Bitmap edge length in pixels
COLORS     = 16           # Palette entries, one nibble per pixel
BACKGROUND = (11, 19, 43) # 0xFF0B132B, the desktop wallpaper's top gradient stop
ASCII_W    = 46           # Text-mode logo width in characters
ASCII_RAMP = " .:-=+*#%@"


def build_bitmap(path):
    img = Image.open(path).convert("RGBA")

    # Composite onto the wallpaper colour so the logo's transparent surround
    # blends into the splash instead of turning black.
    canvas = Image.new("RGB", img.size, BACKGROUND)
    canvas.paste(img, (0, 0), img)

    canvas = canvas.resize((SIZE, SIZE), Image.LANCZOS)

    # Dithering would scatter isolated palette indices through flat areas and
    # look like noise at this size, so quantize with it disabled.
    quant = canvas.quantize(colors=COLORS, dither=Image.NONE)

    raw_palette = quant.getpalette()[: COLORS * 3]
    palette = [tuple(raw_palette[i * 3 : i * 3 + 3]) for i in range(COLORS)]
    indices = list(quant.getdata())
    return palette, indices


def build_ascii(path):
    img = Image.open(path).convert("RGBA")
    canvas = Image.new("RGB", img.size, BACKGROUND)
    canvas.paste(img, (0, 0), img)

    # Console cells are about twice as tall as they are wide.
    height = max(1, ASCII_W // 2)
    small = canvas.convert("L").resize((ASCII_W, height), Image.LANCZOS)

    pixels = list(small.getdata())
    lo, hi = min(pixels), max(pixels)
    span = max(1, hi - lo)

    rows = []
    for y in range(height):
        row = []
        for x in range(ASCII_W):
            level = (pixels[y * ASCII_W + x] - lo) * (len(ASCII_RAMP) - 1) // span
            row.append(ASCII_RAMP[level])
        rows.append("".join(row).rstrip())
    while rows and not rows[0]:
        rows.pop(0)
    while rows and not rows[-1]:
        rows.pop()
    return rows


def transparent_mask(palette):
    """Bitmask of palette entries close enough to the background to skip.

    The source art has a transparent surround, which the compositing step turns
    into the background colour. Quantization then spreads that flat area across
    several near-identical entries, and drawing them would stamp a visible
    rectangle over whatever the logo sits on. Marking them lets the renderer
    skip those pixels so only the logo itself is painted.
    """
    mask = 0
    br, bg, bb = BACKGROUND
    for i, (r, g, b) in enumerate(palette):
        if abs(r - br) + abs(g - bg) + abs(b - bb) <= 12:
            mask |= 1 << i
    return mask


def emit(path, palette, indices, ascii_rows, out_path):
    packed = []
    for i in range(0, len(indices), 2):
        hi = indices[i] & 0x0F
        lo = indices[i + 1] & 0x0F if i + 1 < len(indices) else 0
        packed.append((hi << 4) | lo)

    lines = []
    lines.append("// Generated by scripts/mklogo.py -- do not edit by hand.")
    lines.append("// Source artwork: %s" % os.path.basename(path))
    lines.append("")
    lines.append('#include <init/bootlogo.h>')
    lines.append("")
    lines.append("const int boot_logo_width  = %d;" % SIZE)
    lines.append("const int boot_logo_height = %d;" % SIZE)
    lines.append("")
    lines.append("// Palette entries that are background rather than artwork; the renderer")
    lines.append("// skips these so the logo composites onto whatever is behind it.")
    lines.append("const uint16_t boot_logo_transparent_mask = 0x%04X;"
                 % transparent_mask(palette))
    lines.append("")
    lines.append("// 16-entry ARGB palette; index 0 is the wallpaper background.")
    lines.append("const uint32_t boot_logo_palette[%d] = {" % COLORS)
    for r, g, b in palette:
        lines.append("    0xFF%02X%02X%02X," % (r, g, b))
    lines.append("};")
    lines.append("")
    lines.append("// %dx%d pixels, two 4-bit indices per byte, row major." % (SIZE, SIZE))
    lines.append("const uint8_t boot_logo_bitmap[%d] = {" % len(packed))
    for i in range(0, len(packed), 16):
        chunk = ", ".join("0x%02X" % b for b in packed[i : i + 16])
        lines.append("    %s," % chunk)
    lines.append("};")
    lines.append("")
    lines.append("// ASCII rendition for the text console, before the framebuffer exists.")
    lines.append("const char* const boot_logo_ascii[] = {")
    for row in ascii_rows:
        lines.append('    "%s",' % row.replace("\\", "\\\\").replace('"', '\\"'))
    lines.append("    NULL")
    lines.append("};")
    lines.append("")

    with open(out_path, "w") as fh:
        fh.write("\n".join(lines))

    return len(packed), len(ascii_rows)


def main():
    if len(sys.argv) < 2:
        print(__doc__)
        return 1

    src = os.path.expanduser(sys.argv[1])
    out = sys.argv[2] if len(sys.argv) > 2 else "init/logo_data.c"

    palette, indices = build_bitmap(src)
    ascii_rows = build_ascii(src)
    nbytes, nrows = emit(src, palette, indices, ascii_rows, out)

    print("%s: %dx%d bitmap, %d palette entries, %d bytes packed, %d ASCII rows"
          % (out, SIZE, SIZE, len(palette), nbytes, nrows))
    return 0


if __name__ == "__main__":
    sys.exit(main())