"""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

# How each orientation maps the logically-composed image onto the native
# 800x480 panel. "portrait"/"portrait_flipped" compose at 480x800 (so the
# crop ratio matches how the frame actually hangs) and rotate into native
# space afterwards -- rotation happens after dithering, which is lossless
# (a pure pixel permutation). Which of 90/270 is "portrait" vs
# "portrait_flipped" is a convention pick; whichever way the frame is
# hung, one of the two is right.
ORIENTATION_TRANSPOSE = {
    "landscape": None,
    "landscape_flipped": Image.Transpose.ROTATE_180,
    "portrait": Image.Transpose.ROTATE_90,
    "portrait_flipped": Image.Transpose.ROTATE_270,
}


def logical_render_size(orientation: str) -> tuple[int, int]:
    """(width, height) the photo is composed/cropped at for this
    orientation, before rotating into native panel space."""
    if orientation in ("portrait", "portrait_flipped"):
        return EPD_HEIGHT, EPD_WIDTH
    return EPD_WIDTH, EPD_HEIGHT


def logical_to_native(x: float, y: float, orientation: str) -> tuple[int, int]:
    """Maps a point in logical (pre-rotation) frame space to native
    800x480 panel space, applying the same rotation ORIENTATION_TRANSPOSE
    applies to the pixels -- anything positioned in logical coordinates
    (e.g. face labels) needs this to stay attached to the rotated
    content. PIL's ROTATE_90 is counterclockwise; ROTATE_270 clockwise."""
    logical_w, logical_h = logical_render_size(orientation)
    if orientation == "landscape_flipped":
        return int(logical_w - 1 - x), int(logical_h - 1 - y)
    if orientation == "portrait":  # ROTATE_90 (CCW)
        return int(y), int(logical_w - 1 - x)
    if orientation == "portrait_flipped":  # ROTATE_270 (CW)
        return int(logical_h - 1 - y), int(x)
    return int(x), int(y)

# 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 _plain_center_crop_box(
    img_width: int, img_height: int, target_width: int, target_height: int
) -> tuple[float, float, int, int]:
    """The largest target_width:target_height window centered in the
    source image -- the same box ImageOps.fit() computes internally when
    there's no face-aware shift to apply. Returns (left, top, crop_w,
    crop_h); left/top are floats (not yet rounded) since callers that go
    on to face-shift this box need the unrounded center point."""
    target_ratio = target_width / target_height
    if img_width / img_height > target_ratio:
        crop_h = img_height
        crop_w = int(crop_h * target_ratio)
    else:
        crop_w = img_width
        crop_h = int(crop_w / target_ratio)

    left = (img_width - crop_w) / 2
    top = (img_height - crop_h) / 2
    return left, top, crop_w, crop_h


def _face_aware_crop_box(
    img_width: int, img_height: int, target_width: int, target_height: int, faces: list[dict]
) -> tuple[int, int, int, int]:
    """Largest crop window matching target_width:target_height that fits
    inside the source image. Starts from the plain center crop and only
    shifts it the minimum amount needed to bring any faces that would
    otherwise be cut off back on screen -- an already-fine composition
    (faces already fully inside the center crop) is left untouched rather
    than re-centered on the faces. If the faces themselves span wider than
    the crop window allows, centers on their midpoint as best-effort,
    since there's no shift that fits them all regardless.

    Each face's box is given relative to its own imageWidth/imageHeight
    (the resolution Immich ran detection on), which may differ from the
    downloaded preview's resolution passed in here, so each box is scaled
    into img_width/img_height space before use.
    """
    min_x = min_y = float("inf")
    max_x = max_y = float("-inf")
    for face in faces:
        face_w = face.get("imageWidth") or img_width
        face_h = face.get("imageHeight") or img_height
        scale_x = img_width / face_w
        scale_y = img_height / face_h
        min_x = min(min_x, face["boundingBoxX1"] * scale_x)
        max_x = max(max_x, face["boundingBoxX2"] * scale_x)
        min_y = min(min_y, face["boundingBoxY1"] * scale_y)
        max_y = max(max_y, face["boundingBoxY2"] * scale_y)

    left, top, crop_w, crop_h = _plain_center_crop_box(img_width, img_height, target_width, target_height)

    if max_x - min_x <= crop_w:
        if min_x < left:
            left = min_x
        elif max_x > left + crop_w:
            left = max_x - crop_w
    else:
        left = (min_x + max_x) / 2 - crop_w / 2

    if max_y - min_y <= crop_h:
        if min_y < top:
            top = min_y
        elif max_y > top + crop_h:
            top = max_y - crop_h
    else:
        top = (min_y + max_y) / 2 - crop_h / 2

    left = max(0, min(left, img_width - crop_w))
    top = max(0, min(top, img_height - crop_h))

    return (int(left), int(top), int(left) + crop_w, int(top) + crop_h)


def render_frame(source: Image.Image, faces: list[dict] | None = None,
                 orientation: str = "landscape") -> 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.

    If `faces` (from ImmichClient.get_asset_faces) is non-empty, crops
    toward keeping them on screen instead of a plain center-crop.

    `orientation` (see ORIENTATION_TRANSPOSE) composes the photo for how
    the frame physically hangs, then rotates into native panel space --
    the output byte layout is identical either way.
    """
    logical_w, logical_h = logical_render_size(orientation)
    fitted = ImageOps.exif_transpose(source.convert("RGB"))

    if faces:
        box = _face_aware_crop_box(fitted.width, fitted.height, logical_w, logical_h, faces)
        fitted = fitted.crop(box).resize((logical_w, logical_h), Image.LANCZOS)
    else:
        fitted = ImageOps.fit(fitted, (logical_w, logical_h), method=Image.LANCZOS)

    quantized = fitted.quantize(palette=_PALETTE_IMAGE, dither=Image.Dither.FLOYDSTEINBERG)
    transpose = ORIENTATION_TRANSPOSE.get(orientation)
    if transpose is not None:
        quantized = quantized.transpose(transpose)
    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)