#!/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
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   = 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 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 reserved black)' if RESERVE else ''} -> {out}")