"""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
    tint: float = 0.0            # hue rotation in degrees (-180..180), 0 = none
    red: float = 1.0             # per-channel level (gain), 1.0 = unchanged
    green: float = 1.0
    blue: float = 1.0
    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)
    # per-channel colour levels (RGB gain / colour balance)
    if opt.red != 1.0 or opt.green != 1.0 or opt.blue != 1.0:
        def _gain(g):
            return [min(255, max(0, int(i * g + 0.5))) for i in range(256)]
        img = img.point(_gain(opt.red) + _gain(opt.green) + _gain(opt.blue))
    # tint = rotate every hue around the colour wheel
    if opt.tint != 0.0:
        h, s, v = img.convert("HSV").split()
        shift = int(round(opt.tint / 360.0 * 255))
        h = h.point(lambda i, sh=shift: (i + sh) % 256)
        img = Image.merge("HSV", (h, s, v)).convert("RGB")
    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)