import io
import os
import tempfile
from collections import defaultdict
from collections.abc import Iterator
from typing import Any, NamedTuple, cast

import cv2
import imageio.v3 as iio
import numpy as np
from anydi import singleton
from pygifsicle import gifsicle

from ows_text_campaigns.assets.exceptions import (
    ImageProcessError,
    ImageValidationError,
    InvalidContentTypeError,
    InvalidFileSizeError,
    InvalidImageSizeError,
)
from ows_text_campaigns.assets.types import BinaryAsset, ImageAsset
from ows_text_campaigns.config import Settings

NDArray = np.ndarray[Any, np.dtype[np.unsignedinteger[Any]]]


class ImageNDArray(NamedTuple):
    image: ImageAsset
    ndarray: NDArray


@singleton
class ImageProcessor:
    default_content_types = [
        "image/jpg",
        "image/jpeg",
        "image/png",
        "image/gif",
    ]

    def __init__(self, settings: Settings) -> None:
        self.settings = settings

    def validate(
        self,
        asset: BinaryAsset,
        *,
        min_width: int,
        min_height: int,
        allowed_content_types: list[str] | None = None,
    ) -> ImageNDArray:
        # Ensure asset has valid image extension.
        allowed_content_types = allowed_content_types or self.default_content_types
        if asset.content_type not in allowed_content_types:
            raise InvalidContentTypeError(
                content_type=asset.content_type,
                allowed_content_types=allowed_content_types,
            )

        # Validate file size
        if not (
            self.settings.assets_image_min_allowed_file_size
            <= asset.file_size
            <= self.settings.assets_image_max_allowed_file_size
        ):
            raise InvalidFileSizeError(
                size=asset.file_size,
                min_size=self.settings.assets_image_min_allowed_file_size,
                max_size=self.settings.assets_image_max_allowed_file_size,
            )

        ndimage = self._read_image(asset)

        if asset.extension == ".gif":
            width, height = ndimage.shape[2], ndimage.shape[1]
        else:
            width, height = ndimage.shape[1], ndimage.shape[0]

        # Validate image size
        if width < min_width or height < min_height:
            raise InvalidImageSizeError(
                width=width,
                height=height,
                min_width=min_width,
                min_height=min_height,
            )

        # Create image asset
        image = ImageAsset.from_binary(asset, size=(width, height))
        return ImageNDArray(image=image, ndarray=ndimage)

    def process(
        self,
        image_ndarray: ImageNDArray,
        *,
        max_file_size: int,
        min_width: int,
        min_height: int,
        square_crop: bool = False,
    ) -> ImageAsset:
        image = image_ndarray.image

        if image.extension in (".jpg", ".jpeg"):
            return self._compress_jpeg(
                image_ndarray,
                max_file_size=max_file_size,
                min_width=min_width,
                min_height=min_height,
                square_crop=square_crop,
                resize_step=0.95,
                quality_step=5,
            )
        elif image.extension == ".png":
            return self._compress_png(
                image_ndarray,
                max_file_size=max_file_size,
                min_width=min_width,
                min_height=min_height,
                square_crop=square_crop,
                resize_step=0.9,
            )
        elif image.extension == ".gif":
            if square_crop:
                raise ImageProcessError("Cropping GIF images is not supported.")
            return self._compress_gif(
                image,
                max_file_size=max_file_size,
                min_width=min_width,
                min_height=min_height,
                scale_step=0.05,
            )

        return image

    def validate_and_process(
        self,
        asset: BinaryAsset,
        *,
        max_file_size: int,
        min_width: int,
        min_height: int,
        allowed_content_types: list[str] | None = None,
        square_crop: bool = False,
    ) -> ImageAsset:
        image_ndarray = self.validate(
            asset,
            min_width=min_width,
            min_height=min_height,
            allowed_content_types=allowed_content_types,
        )
        image = image_ndarray.image

        if image.file_size > max_file_size or square_crop:
            return self.process(
                image_ndarray,
                max_file_size=max_file_size,
                min_width=min_width,
                min_height=min_height,
                square_crop=square_crop,
            )
        return image

    @staticmethod
    def _read_image(asset: BinaryAsset) -> NDArray:
        if asset.extension == ".gif":
            frames = cast(
                Iterator[NDArray],
                iio.imiter(asset.data, extension=asset.extension),
            )
            frames_by_shape: dict[tuple[int, ...], list[NDArray]] = defaultdict(list)
            for frame in frames:
                frames_by_shape[frame.shape].append(frame)

            common_shape, _ = max(
                frames_by_shape.items(),
                key=lambda item: len(item[1]),
            )
            return np.array(frames_by_shape[common_shape])

        return iio.imread(asset.data, extension=asset.extension)

    @staticmethod
    def _crop(image_ndarray: ImageNDArray) -> ImageNDArray:
        image, ndarray = image_ndarray

        if image.extension == ".gif":
            raise ImageValidationError("Cropping is not supported for GIF images")

        width, height = image.size

        side = min(width, height)
        left = (width - side) // 2
        top = (height - side) // 2
        ndarray = ndarray[top : top + side, left : left + side]

        out = io.BytesIO()
        _ = iio.imwrite(out, ndarray, extension=image.extension)
        data = out.getvalue()
        out.close()

        return ImageNDArray(
            image=ImageAsset(
                data=data,
                content_type=image.content_type,
                extension=image.extension,
                file_size=len(data),
                filename=image.filename,
                size=(ndarray.shape[1], ndarray.shape[0]),
                is_cropped=True,
            ),
            ndarray=ndarray,
        )

    def _compress_png(
        self,
        image_ndarray: ImageNDArray,
        *,
        max_file_size: int,
        min_width: int,
        min_height: int,
        resize_step: float,
        square_crop: bool = False,
    ) -> ImageAsset:
        if square_crop:
            image_ndarray = self._crop(image_ndarray)

        image, ndarray = image_ndarray
        width, height = image.size
        file_size = image.file_size

        while width >= min_width and height >= min_height:
            resized = cv2.resize(
                ndarray,
                dsize=(width, height),
                interpolation=cv2.INTER_AREA,
            )

            # Encode to memory
            success, encoded = cv2.imencode(
                ext=image.extension,
                img=cv2.cvtColor(resized, cv2.COLOR_BGRA2RGBA),
                params=[cv2.IMWRITE_PNG_COMPRESSION, 6],
            )
            if not success:
                raise ImageProcessError

            file_size = len(encoded)
            if file_size <= max_file_size:
                return ImageAsset(
                    data=encoded.tobytes(),
                    content_type=image.content_type,
                    extension=image.extension,
                    file_size=file_size,
                    filename=image.filename,
                    size=(width, height),
                    is_cropped=image.is_cropped,
                    is_compressed=True,
                )

            width = int(width * resize_step)
            height = int(height * resize_step)

        raise ImageProcessError(
            f"Unable to compress PNG '{image.filename}' below {max_file_size} bytes "
            f"({max_file_size / 1024 / 1024:.2f} MB). "
            f"Attempted downscaling to {width}x{height} px, "
            f"but the final file size was {file_size / 1024 / 1024:.2f} MB. "
        )

    def _compress_jpeg(
        self,
        image_ndarray: ImageNDArray,
        *,
        max_file_size: int,
        min_width: int,
        min_height: int,
        resize_step: float,
        quality_step: int = 5,
        square_crop: bool = False,
    ) -> ImageAsset:
        if square_crop:
            image_ndarray = self._crop(image_ndarray)

        image, ndarray = image_ndarray
        width, height = image.size
        resized_ndarray = ndarray
        file_size = image.file_size

        while width >= min_width and height >= min_height:
            resized = cv2.resize(
                ndarray,
                dsize=(width, height),
                interpolation=cv2.INTER_AREA,
            )

            # Encode to memory
            success, encoded = cv2.imencode(
                ext=image.extension,
                img=cv2.cvtColor(resized, cv2.COLOR_BGR2RGB),
            )
            if not success:
                raise ImageProcessError

            if (file_size := len(encoded)) <= max_file_size:
                return ImageAsset(
                    data=encoded.tobytes(),
                    content_type=image.content_type,
                    extension=image.extension,
                    file_size=file_size,
                    filename=image.filename,
                    size=(width, height),
                    is_compressed=True,
                )

            width = int(width * resize_step)
            height = int(height * resize_step)
            resized_ndarray = encoded

        # Now try degrading quality at final (smallest) size
        start_quality = 95
        quality = start_quality
        for quality in range(start_quality, 5, -quality_step):
            success, encoded = cv2.imencode(
                ext=image.extension,
                img=resized_ndarray,
                params=[cv2.IMWRITE_JPEG_QUALITY, quality],
            )
            if not success:
                raise ImageProcessError

            if (file_size := len(encoded)) <= max_file_size:
                return ImageAsset(
                    data=encoded.tobytes(),
                    content_type=image.content_type,
                    extension=image.extension,
                    file_size=file_size,
                    filename=image.filename,
                    size=(width, height),
                    is_compressed=True,
                )

        raise ImageProcessError(
            f"Unable to compress JPEG below {max_file_size} bytes "
            f"({max_file_size / 1024 / 1024:.2f} MB). "
            f"Attempted downscaling to {width}x{height}px and reducing quality "
            f"to {quality + quality_step}%, but the final file size "
            f"was {file_size / 1024 / 1024:.2f} MB."
        )

    @staticmethod
    def _compress_gif(
        image: ImageAsset,
        *,
        max_file_size: int,
        min_width: int,
        min_height: int,
        scale_step: float,
        initial_scale: float = 0.9,
        min_colors: int = 32,
        color_step: int = 32,
    ) -> ImageAsset:
        width, height = image.size
        min_scale = max(min_width / width, min_height / height)
        scale = initial_scale
        colors = 256
        file_size = image.file_size

        with tempfile.NamedTemporaryFile(suffix=".gif", delete=True) as origin_file:
            origin_file.write(image.data)
            origin_file.flush()

            # First, try scaling down
            while scale >= min_scale:
                with tempfile.NamedTemporaryFile(
                    suffix=".gif", delete=True
                ) as optimized_file:
                    gifsicle(
                        sources=origin_file.name,
                        destination=optimized_file.name,
                        optimize=True,
                        colors=colors,
                        options=[f"--scale={scale}"],
                    )
                    file_size = os.path.getsize(optimized_file.name)
                    if file_size <= max_file_size:
                        return ImageAsset(
                            data=optimized_file.read(),
                            content_type=image.content_type,
                            extension=image.extension,
                            size=(int(width * scale), int(height * scale)),
                            file_size=file_size,
                            filename=image.filename,
                            is_compressed=True,
                        )
                    scale -= scale_step

            # Then try reducing color depth
            while colors > min_colors:
                with tempfile.NamedTemporaryFile(
                    suffix=".gif", delete=True
                ) as optimized_file:
                    gifsicle(
                        sources=origin_file.name,
                        destination=optimized_file.name,
                        optimize=True,
                        colors=colors,
                        options=[f"--scale={min_scale}"],
                    )
                    file_size = os.path.getsize(optimized_file.name)
                    if file_size <= max_file_size:
                        return ImageAsset(
                            data=optimized_file.read(),
                            content_type=image.content_type,
                            extension=image.extension,
                            size=(int(width * min_scale), int(height * min_scale)),
                            file_size=file_size,
                            filename=image.filename,
                            is_compressed=True,
                        )
                    colors -= color_step

        raise ImageProcessError(
            f"Unable to compress GIF below {max_file_size} bytes "
            f"({max_file_size / 1024 / 1024:.2f} MB). "
            f"Tried resizing down to {int(width * min_scale)}x{int(height * min_scale)}px "
            f"and reducing color depth to {colors + color_step} colors, "
            f"but the final file size was still {file_size / 1024 / 1024:.2f} MB. "
        )