#!/usr/bin/env python3
"""Frame -> flat (RGB888 palette + index plane) blob for the MAME Lua loader.

Packing into the X68000 palette word is done Lua-side on purpose: the exact
channel order is a hardware fact we intend to CONFIRM BY EYE, not assume, so it
has to be cheap to change without regenerating the blob.
"""
import sys, struct, glob
import numpy as np
from PIL import Image

argv = [a for a in sys.argv[1:] if not a.startswith("--")]
# --reserve-black: quantise to 255 colours and reserve index 0 as black.
# Needed for any mode that letterboxes (256x192 inside 256x256): GVRAM cleared
# to 0 displays palette entry 0, and a free mediancut palette puts a real image
# colour there. Costs one of 256 entries; measured quality cost is negligible.
RESERVE = "--reserve-black" in sys.argv
# --pack-transparent: the layout FINDINGS 46.6 needs.  The packed scheme puts
# the TOP graphics page's index 0 to work as a transparency key, so index 0 must
# never appear in the picture -- and black therefore cannot live there.  So:
# quantise to 254, place them at 1..254, put black at 255, leave 0 UNUSED.
# Costs two of 256 entries against --reserve-black's one.
PACKT = "--pack-transparent" in sys.argv
src, out = argv[0], argv[1]
f = sorted(glob.glob(f"{src}/*.png"))[int(argv[2]) if len(argv) > 2 else 0]
im = Image.open(f).convert("RGB")
W, H = im.size

n   = 254 if PACKT else (255 if RESERVE else 256)
q   = im.quantize(colors=n, method=Image.MEDIANCUT, dither=Image.NONE)
pal = np.array(q.getpalette()[:n*3], dtype=np.uint8).reshape(n, 3)
idx = np.asarray(q, dtype=np.uint8)
if PACKT:
    # 0 unused (transparency key), 1..254 picture, 255 black
    pal = np.vstack([np.zeros((1, 3), np.uint8), pal, np.zeros((1, 3), np.uint8)])
    idx = idx + 1
    assert idx.min() >= 1 and idx.max() <= 254, "index 0/255 must stay free"
elif RESERVE:
    pal = np.vstack([np.zeros((1, 3), np.uint8), pal])   # index 0 = black
    idx = idx + 1

with open(out, "wb") as fh:
    fh.write(b"DLXR")
    fh.write(struct.pack(">HH", W, H))
    fh.write(pal.tobytes())
    fh.write(idx.tobytes())

# reference PNG of exactly what the X68000 should display
Image.fromarray(pal[idx]).save(out.replace(".bin", "_ref.png"))
print(f"src={f} {W}x{H} colors={len(np.unique(idx))}"
      f"{' (idx 0 unused/transparent, 255 black)' if PACKT else (' (idx 0 reserved black)' if RESERVE else '')} -> {out}")