#!/usr/bin/env python3
""" Animated GIF writer for faceswap.py converter """
import os

import cv2
import imageio

from ._base import Output, logger


class Writer(Output):
    """ Video output writer using imageio """
    def __init__(self, output_folder, total_count, frame_ranges, **kwargs):
        logger.debug("total_count: %s, frame_ranges: %s", total_count, frame_ranges)
        super().__init__(output_folder, **kwargs)
        self.frame_order = self.set_frame_order(total_count, frame_ranges)
        self.output_dimensions = None  # Fix dims of 1st frame in case of different sized images
        self.writer = None  # Need to know dimensions of first frame, so set writer then
        self.gif_file = None  # Set filename based on first file seen

    @property
    def gif_params(self):
        """ Format the gif params """
        kwargs = {key: int(val) for key, val in self.config.items()}
        logger.debug(kwargs)
        return kwargs

    @staticmethod
    def set_frame_order(total_count, frame_ranges):
        """ Return the full list of frames to be converted in order """
        if frame_ranges is None:
            retval = list(range(1, total_count + 1))
        else:
            retval = list()
            for rng in frame_ranges:
                retval.extend(list(range(rng[0], rng[1] + 1)))
        logger.debug("frame_order: %s", retval)
        return retval

    def get_writer(self):
        """ Add the requested encoding options and return the writer """
        logger.debug("writer config: %s", self.config)
        return imageio.get_writer(self.gif_file,
                                  mode="i",
                                  **self.config)

    def write(self, filename, image):
        """ Frames come from the pool in arbitrary order, so cache frames
            for writing out in correct order """
        logger.trace("Received frame: (filename: '%s', shape: %s", filename, image.shape)
        if not self.gif_file:
            self.set_gif_filename(filename)
            self.set_dimensions(image.shape[:2])
            self.writer = self.get_writer()
        if (image.shape[1], image.shape[0]) != self.output_dimensions:
            image = cv2.resize(image, self.output_dimensions)  # pylint: disable=no-member
        self.cache_frame(filename, image)
        self.save_from_cache()

    def set_gif_filename(self, filename):
        """ Set the gif output filename """
        logger.debug("sample filename: '%s'", filename)
        filename = os.path.splitext(os.path.basename(filename))[0]
        idx = len(filename)
        for char in list(filename[::-1]):
            if not char.isdigit() and char not in ("_", "-"):
                break
            idx -= 1
        self.gif_file = os.path.join(self.output_folder, "{}_converted.gif".format(filename[:idx]))
        logger.info("Outputting to: '%s'", self.gif_file)

    def set_dimensions(self, frame_dims):
        """ Set the dimensions based on a given frame frame. This protects against different
            sized images coming in and ensure all images go out at the same size for writers
            that require it """
        logger.debug("input dimensions: %s", frame_dims)
        self.output_dimensions = (frame_dims[1], frame_dims[0])
        logger.debug("Set dimensions: %s", self.output_dimensions)

    def save_from_cache(self):
        """ Save all the frames that are ready to be output from cache """
        while self.frame_order:
            if self.frame_order[0] not in self.cache:
                logger.trace("Next frame not ready. Continuing")
                break
            save_no = self.frame_order.pop(0)
            save_image = self.cache.pop(save_no)
            logger.trace("Rendering from cache. Frame no: %s", save_no)
            self.writer.append_data(save_image[:, :, ::-1])
        logger.trace("Current cache size: %s", len(self.cache))

    def close(self):
        """ Close the ffmpeg writer and mux the audio """
        self.writer.close()