"""OpenFlexure Microscope Camera

This module defines the interface for cameras. Any compatible Thing
should enabe the server to work.

See repository root for licensing information.
"""

from __future__ import annotations
from typing import Literal, Optional, Tuple, Any
import json

from pydantic import RootModel
from PIL import Image
import piexif

from labthings_fastapi.thing import Thing
from labthings_fastapi.decorators import thing_action, thing_property
from labthings_fastapi.dependencies.metadata import GetThingStates
from labthings_fastapi.dependencies.blocking_portal import BlockingPortal
from labthings_fastapi.dependencies.thing import direct_thing_client_dependency
from labthings_fastapi.dependencies.raw_thing import raw_thing_dependency
from labthings_fastapi.dependencies.invocation import InvocationLogger
from labthings_fastapi.outputs.mjpeg_stream import MJPEGStreamDescriptor
from labthings_fastapi.outputs.blob import Blob
from labthings_fastapi.types.numpy import NDArray


class JPEGBlob(Blob):
    media_type: str = "image/jpeg"


class PNGBlob(Blob):
    """A class representing a PNG image as a LabThings FastAPI Blob"""

    media_type: str = "image/png"


class ArrayModel(RootModel):
    """A model for an array"""

    root: NDArray


class CaptureError(RuntimeError):
    """An error trying to capture from a CameraThing"""


class NoImageInMemoryError(RuntimeError):
    """An error called if no image in in memory when an method is called to use that image"""


class CameraMemoryBuffer:
    """
    A class that holds images in memory. The images are by default PIL images.

    However subclasses of BaseCamera can use this class to store other object types
    """

    _storage: dict[int, tuple[Any, Optional[dict]]]

    def __init__(self):
        self._storage = {}
        self._latest_id: int = 0

    def add_image(
        self, image: Any, metadata: Optional[dict], buffer_max: int = 1
    ) -> None:
        self._latest_id += 1
        self._create_space(buffer_max)
        self._storage[self._latest_id] = (image, metadata)
        return self._latest_id

    def get_image(
        self, buffer_id: Optional[int] = None, remove: bool = True
    ) -> tuple[Any, Optional[dict]]:
        """
        If the buffer ID is not set, always clear whole buffer
        """

        # No id given
        if buffer_id is None:
            # Get the latest image and metadata tuple from storage
            try:
                image_tuple = list(self._storage.value())[-1]
            except IndexError as e:
                raise NoImageInMemoryError("No image in memory to save.") from e
            # Clear the storage so images don't get retrieved out of order
            self._storage.clear()
            return image_tuple

        if remove:
            return self._storage.pop(buffer_id)
        return self._storage[buffer_id]

    def clear(self):
        self._storage.clear()

    def _create_space(self, buffer_max: int) -> None:
        """
        Create space to add an image.
        """
        # If only one image to be stored just clear the storage and return
        if buffer_max <= 1:
            self._storage.clear()
            return

        # Number to remove to get the storage down to 1 less than the buffer length
        to_remove = len(self._storage) - (buffer_max - 1)
        # If if there is space. Nothing to do, just return
        if to_remove < 1:
            return

        keys_to_remove = list(self._storage.keys())[:to_remove]
        for key in keys_to_remove:
            del self._storage[key]


class BaseCamera(Thing):
    """The base class for all cameras. All cameras must directly inherit from this class"""

    mjpeg_stream = MJPEGStreamDescriptor()
    lores_mjpeg_stream = MJPEGStreamDescriptor()
    _memory_buffer = CameraMemoryBuffer()

    def __enter__(self) -> None:
        raise NotImplementedError("CameraThings must define their own __enter__ method")

    def __exit__(self, _exc_type, _exc_value, _traceback) -> None:
        raise NotImplementedError("CameraThings must define their own __exit__ method")

    @thing_action
    def start_streaming(
        self, main_resolution: tuple[int, int], buffer_count: int
    ) -> None:
        """Start (or stop and restart) the camera with the given resolution
        for the main stream, and buffer_count number of images in the buffer"""
        raise NotImplementedError(
            "CameraThings must define their own start_streaming method"
        )

    @thing_property
    def stream_active(self) -> bool:
        "Whether the MJPEG stream is active"
        raise NotImplementedError(
            "CameraThings must define their own stream_active method"
        )

    @thing_action
    def capture_array(
        self,
        stream_name: Literal["main", "lores", "raw", "full"] = "main",
        wait: Optional[float] = 5,
    ) -> NDArray:
        raise NotImplementedError(
            "CameraThings must define their own capture_array method"
        )

    @thing_action
    def capture_jpeg(
        self,
        metadata_getter: GetThingStates,
        resolution: Literal["lores", "main", "full"] = "main",
        wait: Optional[float] = 5,
    ) -> JPEGBlob:
        """Acquire one image from the camera and return as a JPEG blob"""
        raise NotImplementedError(
            "CameraThings must define their own capture_jpeg method"
        )

    @thing_action
    def grab_jpeg(
        self,
        portal: BlockingPortal,
        stream_name: Literal["main", "lores"] = "main",
    ) -> JPEGBlob:
        """Acquire one image from the preview stream and return as an array

        This differs from `capture_jpeg` in that it does not pause the MJPEG
        preview stream. Instead, we simply return the next frame from that
        stream (either "main" for the preview stream, or "lores" for the low
        resolution preview). No metadata is returned.
        """
        stream = (
            self.lores_mjpeg_stream if stream_name == "lores" else self.mjpeg_stream
        )
        frame = portal.call(stream.grab_frame)
        return JPEGBlob.from_bytes(frame)

    @thing_action
    def grab_jpeg_size(
        self,
        portal: BlockingPortal,
        stream_name: Literal["main", "lores"] = "main",
    ) -> int:
        """Acquire one image from the preview stream and return its size"""
        stream = (
            self.lores_mjpeg_stream if stream_name == "lores" else self.mjpeg_stream
        )
        return portal.call(stream.next_frame_size)

    @thing_action
    def capture_image(
        self,
        stream_name: Literal["main", "lores", "raw"],
        wait: Optional[float],
    ) -> None:
        """Capture a PIL image from stream stream_name with timeout wait"""
        raise NotImplementedError(
            "CameraThings must define their own capture_image method"
        )

    @thing_action
    def capture_and_save(
        self,
        jpeg_path: str,
        logger: InvocationLogger,
        metadata_getter: GetThingStates,
        save_resolution: Optional[Tuple[int, int]] = None,
    ) -> None:
        """Capture an image and save it to disk


        save_resolution can be set to resize the image before saving. By default this
            is None meaning that the image is saved at original resoltion.
        """
        image, metadata = self._robust_image_capture(
            metadata_getter,
            logger=logger,
        )

        self._save_capture(
            jpeg_path,
            image,
            metadata,
            logger,
            save_resolution,
        )

    @thing_action
    def capture_to_memory(
        self, logger: InvocationLogger, metadata_getter: GetThingStates, buffer_max: int
    ) -> None:
        """
        Capture an image to memory. This can be saved later with `save_from_memory`

        Note that only one image is held in memory so this will overwrite any image
        in memory.

        Return the buffer id of the image captured
        """
        image, metadata = self._robust_image_capture(metadata_getter, logger)
        return self._memory_buffer.add_image(image, metadata, buffer_max=buffer_max)

    @thing_action
    def save_from_memory(
        self,
        jpeg_path: str,
        logger: InvocationLogger,
        save_resolution: Optional[Tuple[int, int]] = None,
        buffer_id: Optional[int] = None,
    ) -> None:
        """
        Save an image that has been captured to memory.
        """
        image, metadata = self._memory_buffer.get_image(buffer_id)

        self._save_capture(
            jpeg_path=jpeg_path,
            image=image,
            metadata=metadata,
            logger=logger,
            save_resolution=save_resolution,
        )

    @thing_action
    def clear_buffers(self) -> None:
        self._memory_buffer.clear()

    def _robust_image_capture(
        self,
        metadata_getter: GetThingStates,
        logger: InvocationLogger,
    ) -> Image:
        """Capture an image in memory and return it with metadata
        CaptureError raised if the capture fails for any reason
        returns tuple with PIL Image, and dict of metadata.

        This robust capturing method attempts to capture the image five times
        each time with a 5 second timeout set.
        """
        for capture_attempts in range(5):
            try:
                metadata = metadata_getter()
                image = self.capture_image(stream_name="main", wait=5)
                return image, metadata
            except TimeoutError:
                logger.warning(
                    f"Attempt {capture_attempts + 1} to capture image timed out. Do you have enough RAM?"
                )
        raise CaptureError("An error occurred while capturing after 5 attempts")

    def _save_capture(
        self,
        jpeg_path: str,
        image: Image,
        metadata: dict,
        logger: InvocationLogger,
        save_resolution: Optional[Tuple[int, int]] = None,
    ) -> None:
        """Saving the captured image and metadata to disk
        logger warning (via InvocationLogger) is raised if metadata is failed to be added
        IOError is raised if the file cannot be saved
        nothing is returned on success"""
        if save_resolution is not None and image.size != save_resolution:
            image = image.resize(save_resolution, Image.BOX)
        try:
            # Per PIL documentation,
            # (https://pillow.readthedocs.io/en/stable/handbook/image-file-formats.html#jpeg)
            # there are two factors when saving a JPEG. Subsampling affects the colour,
            # quality affects the pixels.
            # subsampling = 0 disables subsampling of colour
            # quality = 95 is the maximum recommended - above this, JPEG compression is
            # disabled, file size increases and quality is barely or not affected
            image.save(jpeg_path, quality=95, subsampling=0)
            try:
                # Load EXIF metadata from image so it can be added to.
                exif_dict = piexif.load(jpeg_path)
                exif_dict["Exif"][piexif.ExifIFD.UserComment] = json.dumps(
                    metadata
                ).encode("utf-8")
                piexif.insert(piexif.dump(exif_dict), jpeg_path)
            except:  # noqa: E722
                # We need to capture any exception as there are many reasons metadata
                # might not be added. We warn rather than log the error.
                logger.warning(f"Failed to add metadata to {jpeg_path}")
        except Exception as e:
            raise IOError(f"An error occurred while saving {jpeg_path}") from e


CameraDependency = direct_thing_client_dependency(BaseCamera, "/camera/")
RawCameraDependency = raw_thing_dependency(BaseCamera)