"""Resize, quantize, and pack a photo into the panel's raw 4bpp format."""

from __future__ import annotations

from PIL import Image, ImageOps

EPD_WIDTH = 800
EPD_HEIGHT = 480

# Approximate sRGB for each of the panel's 6 ink colors. These are
# reasonable placeholders, not measured values -- Waveshare doesn't publish
# exact color primaries for this panel. Tune them once you can compare a
# rendered test image against the real panel.
PALETTE_RGB = [
    (0, 0, 0),  # BLACK
    (255, 255, 255),  # WHITE
    (255, 219, 0),  # YELLOW
    (207, 0, 15),  # RED
    (0, 39, 133),  # BLUE
    (0, 133, 55),  # GREEN
]

# The panel's actual 4-bit color codes (see firmware/components/epd7in3e),
# in the same order as PALETTE_RGB. 0x4 is intentionally unused upstream.
PANEL_CODES = [0x0, 0x1, 0x2, 0x3, 0x5, 0x6]


def _build_palette_image() -> Image.Image:
    pal_img = Image.new("P", (1, 1))
    pal_img.putpalette([channel for rgb in PALETTE_RGB for channel in rgb])
    return pal_img


_PALETTE_IMAGE = _build_palette_image()


def render_frame(source: Image.Image) -> bytes:
    """Fits `source` to the panel's resolution, quantizes it to the 6-color
    palette with Floyd-Steinberg dithering, and packs 2 pixels/byte the way
    epd7in3e.c expects. Always returns exactly EPD_WIDTH*EPD_HEIGHT/2 bytes.
    """
    fitted = ImageOps.exif_transpose(source.convert("RGB"))
    fitted = ImageOps.fit(fitted, (EPD_WIDTH, EPD_HEIGHT), method=Image.LANCZOS)

    quantized = fitted.quantize(palette=_PALETTE_IMAGE, dither=Image.Dither.FLOYDSTEINBERG)
    pixels = quantized.load()

    out = bytearray(EPD_WIDTH * EPD_HEIGHT // 2)
    i = 0
    for y in range(EPD_HEIGHT):
        for x in range(0, EPD_WIDTH, 2):
            left = PANEL_CODES[pixels[x, y]]
            right = PANEL_CODES[pixels[x + 1, y]]
            out[i] = (left << 4) | right
            i += 1

    return bytes(out)