from __future__ import annotations

from dataclasses import dataclass
from io import BytesIO
from pathlib import Path

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


ALLOWED_FORMATS = {"PNG", "JPEG", "WEBP"}
MAX_FILE_BYTES = 50 * 1024 * 1024


class InvalidImage(ValueError):
    pass


@dataclass(slots=True)
class LoadedImage:
    rgba: np.ndarray
    icc_profile: bytes | None
    source_format: str

    @property
    def width(self) -> int:
        return int(self.rgba.shape[1])

    @property
    def height(self) -> int:
        return int(self.rgba.shape[0])


def load_image_bytes(data: bytes, max_pixels: int = 40_000_000) -> LoadedImage:
    if not data:
        raise InvalidImage("The image is empty")
    if len(data) > MAX_FILE_BYTES:
        raise InvalidImage("The image exceeds the 50 MB limit")

    try:
        with Image.open(BytesIO(data)) as source:
            if source.format not in ALLOWED_FORMATS:
                raise InvalidImage("Only PNG, JPEG, and WebP images are supported")
            if getattr(source, "is_animated", False):
                raise InvalidImage("Animated images are not supported")
            source.load()
            image = ImageOps.exif_transpose(source)
            if image.width * image.height > max_pixels:
                raise InvalidImage(f"The decoded image exceeds {max_pixels:,} pixels")
            rgba = np.asarray(image.convert("RGBA"), dtype=np.uint8).copy()
            return LoadedImage(
                rgba=rgba,
                icc_profile=source.info.get("icc_profile"),
                source_format=source.format,
            )
    except (UnidentifiedImageError, OSError) as exc:
        raise InvalidImage("The uploaded file is not a readable image") from exc


def load_image_path(path: Path, max_pixels: int = 40_000_000) -> LoadedImage:
    try:
        return load_image_bytes(path.read_bytes(), max_pixels=max_pixels)
    except OSError as exc:
        raise InvalidImage(f"Could not read {path}") from exc


def save_png(image: LoadedImage | np.ndarray, path: Path, icc_profile: bytes | None = None) -> None:
    rgba = image.rgba if isinstance(image, LoadedImage) else image
    if isinstance(image, LoadedImage):
        icc_profile = image.icc_profile
    target = Image.fromarray(rgba, "RGBA")
    target.save(path, format="PNG", icc_profile=icc_profile, optimize=True)