"""Load a modern image and prepare it for a given C64 display mode.

The output is always a plain HxWx3 uint8 sRGB numpy array sized to the mode's
*logical* pixel grid (e.g. 160x200 for multicolor, 320x200 for hires).  The 2:1
multicolor pixel aspect is handled here so the source image is never visually
squashed: we resize to the displayed shape and then sub-sample to the logical
grid, which keeps circles round.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np
from PIL import Image, ImageEnhance, ImageOps


@dataclass
class PrepOptions:
    aspect: str = "fit"          # "fit" (letterbox), "fill" (crop), "stretch"
    brightness: float = 1.0      # 1.0 = unchanged
    contrast: float = 1.0
    saturation: float = 1.0
    gamma: float = 1.0           # <1 brightens midtones, >1 darkens
    border_index: int = 0        # palette index used for letterbox bars


def _apply_enhancements(img: Image.Image, opt: PrepOptions) -> Image.Image:
    if opt.brightness != 1.0:
        img = ImageEnhance.Brightness(img).enhance(opt.brightness)
    if opt.contrast != 1.0:
        img = ImageEnhance.Contrast(img).enhance(opt.contrast)
    if opt.saturation != 1.0:
        img = ImageEnhance.Color(img).enhance(opt.saturation)
    if opt.gamma != 1.0:
        lut = [min(255, int((i / 255.0) ** opt.gamma * 255 + 0.5)) for i in range(256)]
        img = img.point(lut * 3)
    return img


def prepare(
    path_or_img,
    logical_w: int,
    logical_h: int,
    pixel_aspect: float,
    opt: PrepOptions,
    border_rgb=(0, 0, 0),
) -> np.ndarray:
    """Return an (logical_h, logical_w, 3) uint8 sRGB array.

    ``pixel_aspect`` is width/height of one logical pixel on screen (2.0 for
    multicolor, 1.0 for hires).  ``border_rgb`` fills letterbox bars for "fit".
    """
    if isinstance(path_or_img, Image.Image):
        img = path_or_img
    else:
        img = Image.open(path_or_img)
    img = ImageOps.exif_transpose(img).convert("RGB")
    img = _apply_enhancements(img, opt)

    # Displayed target shape (square display pixels): the logical grid stretched
    # by the pixel aspect ratio horizontally.
    disp_w = int(round(logical_w * pixel_aspect))
    disp_h = logical_h

    if opt.aspect == "stretch":
        fitted = img.resize((disp_w, disp_h), Image.LANCZOS)
    elif opt.aspect == "fill":
        fitted = ImageOps.fit(img, (disp_w, disp_h), Image.LANCZOS, centering=(0.5, 0.5))
    else:  # "fit" -> letterbox
        scaled = img.copy()
        scaled.thumbnail((disp_w, disp_h), Image.LANCZOS)
        fitted = Image.new("RGB", (disp_w, disp_h), tuple(int(c) for c in border_rgb))
        fitted.paste(scaled, ((disp_w - scaled.width) // 2, (disp_h - scaled.height) // 2))

    # Collapse displayed shape back to the logical grid (undo the aspect stretch).
    logical = fitted.resize((logical_w, logical_h), Image.LANCZOS)
    return np.asarray(logical, dtype=np.uint8)