"""
Task(s) for assembling a browse movie
"""
from abc import ABC
from abc import abstractmethod
from typing import List
from typing import Literal

import moviepy.editor as mpy
import numpy as np
import pkg_resources
from matplotlib import cm
from PIL import Image
from PIL import ImageDraw
from PIL import ImageFont

from dkist_processing_common.models.fits_access import FitsAccessBase
from dkist_processing_common.models.tags import Tag
from dkist_processing_common.parsers.l0_fits_access import L1FitsAccess
from dkist_processing_common.tasks import WorkflowTaskBase
from dkist_processing_common.tasks.mixin.fits import FitsDataMixin


__all__ = ["AssembleMovie"]


class AssembleMovie(WorkflowTaskBase, FitsDataMixin, ABC):
    """

    Assemble all movie frames (tagged with Tag.movie_frame()) into an mp4 movie file.
    It is intended to be subclassed so, e.g., instrument-specific overlays can be added.

    Subclassing guide:
        Vocabulary:
            * frame - A single movie frame in the actual file. We don't really deal with these.
            * image - One of the sequential movie images. I.e., each file tagged with Tag.movie_frame(). This is different than "frame" because it is very likely that each image will exist in the movie formultiple frames.
        Required:
            `self.write_overlay(PIL.ImageDraw, FitsAccessBase) -> None`
                Can be used to draw whatever you want on each "image". See the PIL.ImageDraw documentation for all of the options.
                A convenience function `self.write_line` is provided for text overlays. If you don't want to draw anything you
                still need to `pass` this function. The second argument is a single movie "image", i.e., a single array tagged
                with `Tag.movie_frame`.
        Class Default Properties:
            * MOVIE_FRAME_SHAPE - `(x, y)` tuple of the final pixel size of the movie
            * IMAGES_PER_SEC - Int that defines how many "images" to display per second. Think of this as the useful information "framerate".
            * FPS - Int that defines the final movie frames-per-sec. Change this to conform the video for whatever service you need.
            * FONT_FILE - Str pointing to a ttf file with whatever font you want to use
            * TEXT_MARGIN_PX - Int with padding used by `write_line()`
            * MPL_COLOR_MAP - Str of matplotlib named colormap. Only used in the default `apply_colormap`. An easy and simple way to change the colormap
        Class @properties:
            * `fits_parsing_class` - The subclass of `L0FitsAccess` to use for reading images
            * `num_images` - Int that defines the total number of "images" in the final movie. Probably num_dsps_repeats, but maybe not for weird data.
            * `tags_for_image_n` - List that contains the tags used to grab image "n". Probably just the dsps_repeat num, but maybe not for weird data.
        Other Functions::
            * `apply_colormap(ndarray) -> ndarray`
              You can override the default if you need some fancy way to convert single floats to (RGB) color info. The output needs
              to have shape `(x, y, 3)`.
              Default is to use matplotlib to apply whatever colormap is defined as self.MPL_COLOR_MAP
            * `write_line(draw: PIL.ImageDraw, text: str, line: int, column: Literal["left", "middle", "right"], **text_kwargs) -> None`
                Useful for writing a line of text in `write_overlay`. ``line = 0`` is the bottom of the frame.

    """

    MOVIE_FRAME_SHAPE = (1024, 768)  # Is this a/THE standard size?
    IMAGES_PER_SEC = 3
    FPS = 15
    FONT_FILE = pkg_resources.resource_filename("dkist_processing_common", "fonts/Lato-Regular.ttf")
    TEXT_MARGIN_PX = 5
    MPL_COLOR_MAP = "viridis"

    @property
    def fits_parsing_class(self):
        return L1FitsAccess

    @property
    def num_images(self) -> int:
        """ The total number of images that will be used in the final movie """
        return self.constants.num_dsps_repeats

    def tags_for_image_n(self, n: int) -> List[str]:
        """ The tags to use to read the n'th image in the movie. It's OK to omit `Tag.movie_frame`. """
        return [Tag.dsps_repeat(n + 1)]

    def apply_colormap(self, array: np.ndarray) -> np.ndarray:
        """
        Convert floats to RGB colors.

        Default is to use matplotlib to apply a named colormap; whatever is named in `self.MPL_COLOR_MAP`.
        """
        color_mapper = cm.get_cmap(self.MPL_COLOR_MAP)
        scaled_array = array / array.max()
        return color_mapper(scaled_array, bytes=True)[
            :, :, :-1
        ]  # Drop the last (alpha) color dimension

    def pre_run(self) -> None:
        """
        Set up some movie and font constants
        """
        super().pre_run()
        self.font_15 = ImageFont.truetype(self.FONT_FILE, size=15)
        self.font_18 = ImageFont.truetype(self.FONT_FILE, size=18)
        self.duration = self.num_images / self.IMAGES_PER_SEC

        ## This is here to fix a floating point bug in moviepy
        #
        # The following line is what moviepy uses to iterate over the number of frames. It can be larger than we want
        # (duration * fps) due to floating point errors
        iter_nframes = np.arange(0, self.duration, 1.0 / self.FPS).size
        wanted_nframes = int(self.duration * self.FPS)
        if iter_nframes < wanted_nframes:
            raise ValueError(
                "The moviepy iterator will not produce enough frames and I don't know how to fix it"
            )
        while iter_nframes > wanted_nframes:
            self.duration -= 1.0 / self.FPS
            iter_nframes = np.arange(0, self.duration, 1.0 / self.FPS).size

    def run(self) -> None:
        """
        Setup moviepy's VideoClip mechanism, write the movie, and tag it as a movie()

        Unlike other files, the movie file has a specific name (because .mp4 extension is needed by ffmpeg). It is
        "{dataset_id}.mp4"
        """
        clip = mpy.VideoClip(self.make_frame, duration=self.duration)

        relative_movie_path = f"{self.constants.dataset_id}.mp4"
        absolute_movie_path = str(self.scratch.absolute_path(relative_movie_path))

        clip.write_videofile(absolute_movie_path, fps=self.FPS, codec="mpeg4")
        self.tag(path=absolute_movie_path, tags=[Tag.movie(), Tag.output()])

    def make_frame(self, t: float) -> np.ndarray:
        """
        Single-frame generator used by the moviepy machinery.

        This function takes in a numpy array, converts it to 8-bit int, and draws it into a movie frame canvas.
        It also calls the overlay function to write any metadata onto the frame.
        """
        image_num = int(t * self.IMAGES_PER_SEC)
        tags = self.tags_for_image_n(image_num)
        if Tag.movie_frame() not in tags:
            tags.append(Tag.movie_frame())

        frame_access_list = list(
            self.fits_data_read_fits_access(tags=tags, cls=self.fits_parsing_class)
        )
        if len(frame_access_list) == 0:
            raise ValueError(f"Did not find any frames for {image_num=} at time {t}")
        if len(frame_access_list) > 1:
            raise ValueError(
                f"Found {len(frame_access_list)} frames for {image_num=}. Expected only 1"
            )

        frame_access = frame_access_list[0]
        frame_data = frame_access.data
        colored_image = self.apply_colormap(frame_data)
        if colored_image.shape[-1] != 3:
            raise ValueError(
                f"Expected an RGB color image. Got shape {colored_image.shape} instead."
            )

        background = Image.new("RGB", self.MOVIE_FRAME_SHAPE, (0, 0, 0))
        frame_im = Image.fromarray(colored_image, mode="RGB")
        frame_im = frame_im.resize(self.MOVIE_FRAME_SHAPE)
        background.paste(frame_im)
        draw = ImageDraw.Draw(background)
        self.write_overlay(draw, frame_access)
        return np.array(background)

    @abstractmethod
    def write_overlay(self, draw: ImageDraw, fits_obj: FitsAccessBase) -> None:
        """
        Draw whatever you want on the data contained in a single `FitsAccess` object.

        See the documentation for PIL.ImageDraw for what you can do. If all that's needed is some text then use
        `write_line`.

        Example::

            def write_overlay(self, draw, fits_obj):
                self.write_line(draw, f"INSTRUMENT: FOO", 1, column="left")
                self.write_line(draw, f"WAVELENGTH: {fits_obj.wavelength}", 2, column="middle")
                self.write_line(draw, f"DATE BEGIN: {fits_obj.date_begin}", 3, column="right")

        """
        pass

    def write_line(
        self,
        draw: ImageDraw,
        text: str,
        line: int,
        column: Literal["left", "middle", "right"],
        **text_kwargs,
    ) -> None:
        """
        Draws one line of text on the video frame

        Parameters
        ----------
        draw
            The ImageDraw object of the current frame
        text
            The text to write
        line
            The row at which to write the text. 0 is the lowest row.
        column
            Must be either "left", "middle", or "right". The horizontal alignment of the text.
        text_kwargs
            Any additional kwargs to pass to `ImageDraw.ImageDraw.text`

        Returns
        -------
        None

        """
        text_width, text_height = draw.textsize(text, font=text_kwargs.get("font"))
        y = self.MOVIE_FRAME_SHAPE[1] - (self.TEXT_MARGIN_PX * line) - (text_height * line)

        if column == "right":
            anchor = "rd"
            x = self.MOVIE_FRAME_SHAPE[0] - self.TEXT_MARGIN_PX
        elif column == "middle":
            anchor = "md"
            x = self.MOVIE_FRAME_SHAPE[0] // 2
        elif column == "left":
            anchor = "ld"
            x = self.TEXT_MARGIN_PX
        else:
            raise ValueError(
                f"column argument must be in ['left', 'middle', 'right'], got {column=}"
            )

        draw.text((x, y), text, anchor=anchor, **text_kwargs)
