"""Format-aware ithmb encode/decode for iPod artwork.

Ported from iOpenPod ArtworkDB_Writer/ithmb_codecs.py.
Supports RGB565_LE, RGB565_BE, RGB565_BE_90, RGB555_LE, RGB555_BE,
REC_RGB555_LE, UYVY, I420_LE, JPEG.
"""

from __future__ import annotations

import io
import logging

import numpy as np
from PIL import Image

from .presets import ARTWORK_FORMATS_BY_ID, ArtworkFormat
from .types import EncodedFormatPayload

logger = logging.getLogger(__name__)


def _fmt(format_id: int, fmt_override: ArtworkFormat | None = None) -> ArtworkFormat | None:
    return fmt_override if fmt_override is not None else ARTWORK_FORMATS_BY_ID.get(format_id)


def format_pixel_format(format_id: int, fmt_override: ArtworkFormat | None = None) -> str:
    fmt = _fmt(format_id, fmt_override)
    return fmt.pixel_format if fmt is not None else "UNKNOWN"


def format_dimensions(
    format_id: int,
    fallback_w: int,
    fallback_h: int,
    fmt_override: ArtworkFormat | None = None,
) -> tuple[int, int]:
    fmt = _fmt(format_id, fmt_override)
    if fmt is None:
        return fallback_w, fallback_h
    return int(fmt.width), int(fmt.height)


def default_stride_pixels(format_id: int, width: int, fmt_override: ArtworkFormat | None = None) -> int:
    fmt = _fmt(format_id, fmt_override)
    if fmt is None:
        return width
    pf = fmt.pixel_format
    if pf in ("RGB565_LE", "RGB565_BE", "RGB565_BE_90", "RGB555_LE", "RGB555_BE"):
        return max(width, int(fmt.row_bytes // 2) if fmt.row_bytes else width)
    if pf.startswith("REC_RGB555"):
        return max(width, int(fmt.row_bytes // 2) if fmt.row_bytes else width)
    if pf == "UYVY":
        return max(width, int(fmt.row_bytes // 2) if fmt.row_bytes else width)
    return width


def expected_size_bytes(
    format_id: int,
    width: int,
    height: int,
    stride_pixels: int | None = None,
    fmt_override: ArtworkFormat | None = None,
) -> int:
    pf = format_pixel_format(format_id, fmt_override=fmt_override)
    stride = stride_pixels if stride_pixels is not None else default_stride_pixels(
        format_id, width, fmt_override=fmt_override,
    )
    if pf in ("RGB565_LE", "RGB565_BE", "RGB565_BE_90", "RGB555_LE", "RGB555_BE", "UYVY") or pf.startswith("REC_RGB555"):
        return int(stride) * int(height) * 2
    if pf == "I420_LE":
        w = int(width) & ~1
        h = int(height) & ~1
        return (w * h * 3) // 2
    if pf in ("JPEG", "UNKNOWN"):
        return 0
    return int(stride) * int(height) * 2


def _rgb565_array_from_image(img: Image.Image) -> np.ndarray:
    arr = np.array(img.convert("RGB"), dtype=np.uint32)
    r = (arr[:, :, 0] >> 3) & 0x1F
    g = (arr[:, :, 1] >> 2) & 0x3F
    b = (arr[:, :, 2] >> 3) & 0x1F
    return ((r << 11) | (g << 5) | b).astype(np.uint16)


def _rgb565_to_rgb(arr16: np.ndarray) -> np.ndarray:
    r = ((arr16 >> 11) & 0x1F).astype(np.uint8)
    g = ((arr16 >> 5) & 0x3F).astype(np.uint8)
    b = (arr16 & 0x1F).astype(np.uint8)
    return np.stack(((r << 3) | (r >> 2), (g << 2) | (g >> 4), (b << 3) | (b >> 2)), axis=2)


def _rgb555_array_from_image(img: Image.Image) -> np.ndarray:
    arr = np.array(img.convert("RGB"), dtype=np.uint32)
    r = (arr[:, :, 0] >> 3) & 0x1F
    g = (arr[:, :, 1] >> 3) & 0x1F
    b = (arr[:, :, 2] >> 3) & 0x1F
    return ((r << 10) | (g << 5) | b).astype(np.uint16)


def _rgb555_to_rgb(arr16: np.ndarray) -> np.ndarray:
    r = ((arr16 >> 10) & 0x1F).astype(np.uint8)
    g = ((arr16 >> 5) & 0x1F).astype(np.uint8)
    b = (arr16 & 0x1F).astype(np.uint8)
    return np.stack(((r << 3) | (r >> 2), (g << 3) | (g >> 2), (b << 3) | (b >> 2)), axis=2)


def _pad_packed_rows(arr16: np.ndarray, stride_pixels: int) -> np.ndarray:
    stride = max(1, int(stride_pixels))
    height, width = arr16.shape[:2]
    if stride <= width:
        return arr16
    padded = np.zeros((height, stride), dtype=arr16.dtype)
    padded[:, :width] = arr16
    return padded


def encode_for_ipod(
    source_img: Image.Image,
    format_id: int,
    target_width: int | None = None,
    target_height: int | None = None,
    fmt_override: ArtworkFormat | None = None,
) -> EncodedFormatPayload:
    """Encode a PIL Image into an iPod-native format payload."""
    pf = format_pixel_format(format_id, fmt_override=fmt_override)
    w, h = format_dimensions(
        format_id,
        int(target_width or source_img.width),
        int(target_height or source_img.height),
        fmt_override=fmt_override,
    )
    stride = default_stride_pixels(format_id, w, fmt_override=fmt_override)
    base = source_img.convert("RGB").resize((w, h), Image.Resampling.LANCZOS)

    if pf == "RGB565_BE_90":
        rotated = base.transpose(Image.Transpose.ROTATE_270)
        arr16 = _rgb565_array_from_image(rotated)
        arr16 = _pad_packed_rows(arr16, stride)
        raw = arr16.astype(">u2").tobytes()
        return EncodedFormatPayload(raw, w, h, len(raw), stride, pixel_format=pf)

    if pf == "RGB565_BE":
        arr16 = _rgb565_array_from_image(base)
        arr16 = _pad_packed_rows(arr16, stride)
        raw = arr16.astype(">u2").tobytes()
        return EncodedFormatPayload(raw, w, h, len(raw), stride, pixel_format=pf)

    if pf == "RGB555_BE":
        arr16 = _rgb555_array_from_image(base)
        arr16 = _pad_packed_rows(arr16, stride)
        raw = arr16.astype(">u2").tobytes()
        return EncodedFormatPayload(raw, w, h, len(raw), stride, pixel_format=pf)

    if pf in ("RGB555_LE", "REC_RGB555_LE"):
        arr16 = _rgb555_array_from_image(base)
        arr16 = _pad_packed_rows(arr16, stride)
        raw = arr16.astype("<u2").tobytes()
        return EncodedFormatPayload(raw, w, h, len(raw), stride, pixel_format=pf)

    if pf == "JPEG":
        out = io.BytesIO()
        base.save(out, format="JPEG", quality=92, optimize=False)
        raw = out.getvalue()
        return EncodedFormatPayload(raw, w, h, len(raw), stride, pixel_format=pf)

    if pf == "UYVY":
        if w % 2 != 0:
            w -= 1
            base = base.resize((w, h), Image.Resampling.LANCZOS)
        arr = np.array(base, dtype=np.float32)
        r, g, b = arr[:, :, 0], arr[:, :, 1], arr[:, :, 2]
        y = np.clip(0.257 * r + 0.504 * g + 0.098 * b + 16, 0, 255).astype(np.uint8)
        u = np.clip(-0.148 * r - 0.291 * g + 0.439 * b + 128, 0, 255)
        v = np.clip(0.439 * r - 0.368 * g - 0.071 * b + 128, 0, 255)
        u2 = ((u[:, 0::2] + u[:, 1::2]) * 0.5).astype(np.uint8)
        v2 = ((v[:, 0::2] + v[:, 1::2]) * 0.5).astype(np.uint8)
        packed = np.empty((h, w * 2), dtype=np.uint8)
        packed[:, 0::4] = u2
        packed[:, 1::4] = y[:, 0::2]
        packed[:, 2::4] = v2
        packed[:, 3::4] = y[:, 1::2]
        raw = packed.tobytes()
        return EncodedFormatPayload(raw, w, h, len(raw), stride, pixel_format=pf)

    if pf == "I420_LE":
        w_even, h_even = w & ~1, h & ~1
        if w_even != w or h_even != h:
            w, h = w_even, h_even
            base = base.resize((w, h), Image.Resampling.LANCZOS)
        arr = np.array(base, dtype=np.float32)
        r, g, b = arr[:, :, 0], arr[:, :, 1], arr[:, :, 2]
        y = np.clip(0.257 * r + 0.504 * g + 0.098 * b + 16, 0, 255).astype(np.uint8)
        u = np.clip(-0.148 * r - 0.291 * g + 0.439 * b + 128, 0, 255)
        v = np.clip(0.439 * r - 0.368 * g - 0.071 * b + 128, 0, 255)
        u420 = ((u[0::2, 0::2] + u[0::2, 1::2] + u[1::2, 0::2] + u[1::2, 1::2]) * 0.25).astype(np.uint8)
        v420 = ((v[0::2, 0::2] + v[0::2, 1::2] + v[1::2, 0::2] + v[1::2, 1::2]) * 0.25).astype(np.uint8)
        raw = y.tobytes() + u420.tobytes() + v420.tobytes()
        return EncodedFormatPayload(raw, w, h, len(raw), stride, pixel_format=pf)

    arr16 = _rgb565_array_from_image(base)
    arr16 = _pad_packed_rows(arr16, stride)
    raw = arr16.astype("<u2").tobytes()
    return EncodedFormatPayload(raw, w, h, len(raw), stride, "RGB565_LE")


def decode_from_ipod(
    pixel_bytes: bytes,
    format_id: int,
    width: int,
    height: int,
    hpad: int = 0,
    vpad: int = 0,
    fmt_override: ArtworkFormat | None = None,
) -> Image.Image | None:
    """Decode iPod-native pixel data back to a PIL Image."""
    pf = format_pixel_format(format_id, fmt_override=fmt_override)
    width = max(1, int(width))
    height = max(1, int(height))

    if pf in ("RGB565_LE", "RGB565_BE", "RGB565_BE_90"):
        dtype = "<u2" if pf == "RGB565_LE" else ">u2"
        stride = default_stride_pixels(format_id, width, fmt_override=fmt_override)
        stored_h = height
        needed = stride * stored_h * 2
        if len(pixel_bytes) < needed:
            return None
        arr = np.frombuffer(pixel_bytes[:needed], dtype=dtype)
        if arr.size != stride * stored_h:
            return None
        arr = arr.reshape((stored_h, stride))
        arr = arr[:, :width]
        rgb = _rgb565_to_rgb(arr)
        if pf == "RGB565_BE_90":
            rgb = np.rot90(rgb, k=1)
        return Image.fromarray(rgb)

    if pf in ("RGB555_LE", "RGB555_BE", "REC_RGB555_LE"):
        dtype = "<u2" if pf != "RGB555_BE" else ">u2"
        stride = default_stride_pixels(format_id, width, fmt_override=fmt_override)
        stored_h = height
        needed = stride * stored_h * 2
        if len(pixel_bytes) < needed:
            return None
        arr = np.frombuffer(pixel_bytes[:needed], dtype=dtype)
        if arr.size != stride * stored_h:
            return None
        arr = arr.reshape((stored_h, stride))
        arr = arr[:, :width]
        rgb = _rgb555_to_rgb(arr)
        return Image.fromarray(rgb)

    if pf == "JPEG":
        try:
            return Image.open(io.BytesIO(pixel_bytes)).convert("RGB")
        except Exception:
            return None

    return None
