#!/usr/bin/env python3
""" Tool to preview swaps and tweak configuration prior to running a convert """
from __future__ import annotations
import gettext
import logging
import random
import tkinter as tk
import typing as T

from tkinter import ttk
import os
import sys

from threading import Event, Lock, Thread

import numpy as np

from lib.align import DetectedFace
from lib.cli.args import ConvertArgs
from lib.gui.utils import get_images, get_config, initialize_config, initialize_images
from lib.convert import Converter
from lib.utils import FaceswapError
from lib.queue_manager import queue_manager
from scripts.fsmedia import Alignments, Images
from scripts.convert import Predict, ConvertItem

from plugins.extract.pipeline import ExtractMedia

from .control_panels import ActionFrame, ConfigTools, OptionsBook
from .viewer import FacesDisplay, ImagesCanvas

if T.TYPE_CHECKING:
    from argparse import Namespace
    from lib.queue_manager import EventQueue
    from .control_panels import BusyProgressBar

logger = logging.getLogger(__name__)  # pylint: disable=invalid-name

# LOCALES
_LANG = gettext.translation("tools.preview", localedir="locales", fallback=True)
_ = _LANG.gettext


class Preview(tk.Tk):  # pylint:disable=too-few-public-methods
    """ This tool is part of the Faceswap Tools suite and should be called from
    ``python tools.py preview`` command.

    Loads up 5 semi-random face swaps and displays them, cropped, in place in the final frame.
    Allows user to live tweak settings, before saving the final config to
    :file:`./config/convert.ini`

    Parameters
    ----------
    arguments: :class:`argparse.Namespace`
        The :mod:`argparse` arguments as passed in from :mod:`tools.py`
    """
    _w: str

    def __init__(self, arguments: Namespace) -> None:
        logger.debug("Initializing %s: (arguments: '%s'", self.__class__.__name__, arguments)
        super().__init__()
        self._config_tools = ConfigTools()
        self._lock = Lock()
        self._dispatcher = Dispatcher(self)
        self._display = FacesDisplay(self, 256, 64)
        self._samples = Samples(self, arguments, 5)
        self._patch = Patch(self, arguments)

        self._initialize_tkinter()
        self._image_canvas: ImagesCanvas | None = None
        self._opts_book: OptionsBook | None = None
        self._cli_frame: ActionFrame | None = None  # cli frame holds cli options
        logger.debug("Initialized %s", self.__class__.__name__)

    @property
    def config_tools(self) -> "ConfigTools":
        """ :class:`ConfigTools`: The object responsible for parsing configuration options and
        updating to/from the GUI """
        return self._config_tools

    @property
    def dispatcher(self) -> "Dispatcher":
        """ :class:`Dispatcher`: The object responsible for triggering events and variables and
        handling global GUI state """
        return self._dispatcher

    @property
    def display(self) -> FacesDisplay:
        """ :class:`~tools.preview.viewer.FacesDisplay`: The object that holds the sample,
        converted and patched faces """
        return self._display

    @property
    def lock(self) -> Lock:
        """ :class:`threading.Lock`: The threading lock object for the Preview GUI """
        return self._lock

    @property
    def progress_bar(self) -> BusyProgressBar:
        """ :class:`~tools.preview.control_panels.BusyProgressBar`: The progress bar that indicates
        a swap/patch thread is running """
        assert self._cli_frame is not None
        return self._cli_frame.busy_progress_bar

    def update_display(self):
        """ Update the images in the canvas and redraw """
        if not hasattr(self, "_image_canvas"):  # On first call object not yet created
            return
        assert self._image_canvas is not None
        self._image_canvas.reload()

    def _initialize_tkinter(self) -> None:
        """ Initialize a standalone tkinter instance. """
        logger.debug("Initializing tkinter")
        initialize_config(self, None, None)
        initialize_images()
        get_config().set_geometry(940, 600, fullscreen=False)
        self.title("Faceswap.py - Convert Settings")
        self.tk.call(
            "wm",
            "iconphoto",
            self._w,
            get_images().icons["favicon"])  # pylint:disable=protected-access
        logger.debug("Initialized tkinter")

    def process(self) -> None:
        """ The entry point for the Preview tool from :file:`lib.tools.cli`.

        Launch the tkinter preview Window and run main loop.
        """
        self._build_ui()
        self.mainloop()

    def _refresh(self, *args) -> None:
        """ Patch faces with current convert settings.

        Parameters
        ----------
        *args: tuple
            Unused, but required for tkinter callback.
        """
        logger.debug("Patching swapped faces. args: %s", args)
        self._dispatcher.set_busy()
        self._config_tools.update_config()
        with self._lock:
            assert self._cli_frame is not None
            self._patch.converter_arguments = self._cli_frame.convert_args

        self._dispatcher.set_needs_patch()
        logger.debug("Patched swapped faces")

    def _build_ui(self) -> None:
        """ Build the elements for displaying preview images and options panels. """
        container = ttk.PanedWindow(self,
                                    orient=tk.VERTICAL)
        container.pack(fill=tk.BOTH, expand=True)
        setattr(container, "preview_display", self._display)  # TODO subclass not setattr
        self._image_canvas = ImagesCanvas(self, container)
        container.add(self._image_canvas, weight=3)

        options_frame = ttk.Frame(container)
        self._cli_frame = ActionFrame(self, options_frame)
        self._opts_book = OptionsBook(options_frame,
                                      self._config_tools,
                                      self._refresh)
        container.add(options_frame, weight=1)
        self.update_idletasks()
        container.sashpos(0, int(400 * get_config().scaling_factor))


class Dispatcher():
    """ Handles the app level tk.Variables and the threading events. Dispatches events to the
    correct location and handles GUI state whilst events are handled

    Parameters
    ----------
    app: :class:`Preview`
        The main tkinter Preview app
    """
    def __init__(self, app: Preview):
        logger.debug("Initializing %s: (app: %s)", self.__class__.__name__, app)
        self._app = app
        self._tk_busy = tk.BooleanVar(value=False)
        self._evnt_needs_patch = Event()
        self._is_updating = False
        self._stacked_event = False
        logger.debug("Initialized %s", self.__class__.__name__)

    @property
    def needs_patch(self) -> Event:
        """:class:`threading.Event`. Set by the parent and cleared by the child. Informs the child
        patching thread that a run needs to be processed """
        return self._evnt_needs_patch

    # TKInter Variables
    def set_busy(self) -> None:
        """ Set the tkinter busy variable to ``True`` and display the busy progress bar """
        if self._tk_busy.get():
            logger.debug("Busy event is already set. Doing nothing")
            return
        if not hasattr(self._app, "progress_bar"):
            logger.debug("Not setting busy during initial startup")
            return

        logger.debug("Setting busy event to True")
        self._tk_busy.set(True)
        self._app.progress_bar.start()
        self._app.update_idletasks()

    def _unset_busy(self) -> None:
        """ Set the tkinter busy variable to ``False`` and hide the busy progress bar """
        self._is_updating = False
        if not self._tk_busy.get():
            logger.debug("busy unset when already unset. Doing nothing")
            return
        logger.debug("Setting busy event to False")
        self._tk_busy.set(False)
        self._app.progress_bar.stop()
        self._app.update_idletasks()

    # Threading Events
    def _wait_for_patch(self) -> None:
        """ Wait for a patch thread to complete before triggering a display refresh and unsetting
        the busy indicators """
        logger.debug("Checking for patch completion...")
        if self._evnt_needs_patch.is_set():
            logger.debug("Samples not patched. Waiting...")
            self._app.after(1000, self._wait_for_patch)
            return

        logger.debug("Patch completion detected")
        self._app.update_display()
        self._unset_busy()

        if self._stacked_event:
            logger.debug("Processing last stacked event")
            self.set_busy()
            self._stacked_event = False
            self.set_needs_patch()
            return

    def set_needs_patch(self) -> None:
        """ Sends a trigger to the patching thread that it needs to be run. Waits for the patching
        to complete prior to triggering a display refresh and unsetting the busy indicators """
        if self._is_updating:
            logger.debug("Request to run patch when it is already running. Adding stacked event.")
            self._stacked_event = True
            return
        self._is_updating = True
        logger.debug("Triggering patch")
        self._evnt_needs_patch.set()
        self._wait_for_patch()


class Samples():
    """ The display samples.

    Obtains and holds :attr:`sample_size` semi random test faces for displaying in the
    preview GUI.

    The file list is split into evenly sized groups of :attr:`sample_size`. When a display set is
    generated, a random image from each of the groups is selected to provide an array of images
    across the length of the video.

    Parameters
    ----------
    app: :class:`Preview`
        The main tkinter Preview app
    arguments: :class:`argparse.Namespace`
        The :mod:`argparse` arguments as passed in from :mod:`tools.py`
    sample_size: int
        The number of samples to take from the input video/images
    """

    def __init__(self, app: Preview, arguments: Namespace, sample_size: int) -> None:
        logger.debug("Initializing %s: (app: %s, arguments: '%s', sample_size: %s)",
                     self.__class__.__name__, app, arguments, sample_size)
        self._sample_size = sample_size
        self._app = app
        self._input_images: list[ConvertItem] = []
        self._predicted_images: list[tuple[ConvertItem, np.ndarray]] = []

        self._images = Images(arguments)
        self._alignments = Alignments(arguments,
                                      is_extract=False,
                                      input_is_video=self._images.is_video)
        if self._alignments.version == 1.0:
            logger.error("The alignments file format has been updated since the given alignments "
                         "file was generated. You need to update the file to proceed.")
            logger.error("To do this run the 'Alignments Tool' > 'Extract' Job.")
            sys.exit(1)
        if not self._alignments.have_alignments_file:
            logger.error("Alignments file not found at: '%s'", self._alignments.file)
            sys.exit(1)
        self._filelist = self._get_filelist()
        self._indices = self._get_indices()

        self._predictor = Predict(queue_manager.get_queue("preview_predict_in"),
                                  self._sample_size,
                                  arguments)
        self._app._display.set_centering(self._predictor.centering)
        self.generate()

        logger.debug("Initialized %s", self.__class__.__name__)

    @property
    def available_masks(self) -> list[str]:
        """ list: The mask names that are available for every face in the alignments file """
        retval = [key
                  for key, val in self.alignments.mask_summary.items()
                  if val == self.alignments.faces_count]
        return retval

    @property
    def sample_size(self) -> int:
        """ int: The number of samples to take from the input video/images """
        return self._sample_size

    @property
    def predicted_images(self) -> list[tuple[ConvertItem, np.ndarray]]:
        """ list: The predicted faces output from the Faceswap model """
        return self._predicted_images

    @property
    def alignments(self) -> Alignments:
        """ :class:`~lib.align.Alignments`: The alignments for the preview faces """
        return self._alignments

    @property
    def predictor(self) -> Predict:
        """ :class:`~scripts.convert.Predict`: The Predictor for the Faceswap model """
        return self._predictor

    @property
    def _random_choice(self) -> list[int]:
        """ list: Random indices from the :attr:`_indices` group """
        retval = [random.choice(indices) for indices in self._indices]
        logger.debug(retval)
        return retval

    def _get_filelist(self) -> list[str]:
        """ Get a list of files for the input, filtering out those frames which do
        not contain faces.

        Returns
        -------
        list
            A list of filenames of frames that contain faces.
        """
        logger.debug("Filtering file list to frames with faces")
        if isinstance(self._images.input_images, str):
            filelist = [f"{os.path.splitext(self._images.input_images)[0]}_{frame_no:06d}.png"
                        for frame_no in range(1, self._images.images_found + 1)]
        else:
            filelist = self._images.input_images

        retval = [filename for filename in filelist
                  if self._alignments.frame_has_faces(os.path.basename(filename))]
        logger.debug("Filtered out frames: %s", self._images.images_found - len(retval))
        try:
            assert retval
        except AssertionError as err:
            msg = ("No faces were found in any of the frames passed in. Make sure you are passing "
                   "in a frames source rather than extracted faces, and that you have provided "
                   "the correct alignments file.")
            raise FaceswapError(msg) from err
        return retval

    def _get_indices(self) -> list[list[int]]:
        """ Get indices for each sample group.

        Obtain :attr:`self.sample_size` evenly sized groups of indices
        pertaining to the filtered :attr:`self._file_list`

        Returns
        -------
        list
            list of indices relating to the filtered file list, split into groups
        """
        # Remove start and end values to get a list divisible by self.sample_size
        no_files = len(self._filelist)
        self._sample_size = min(self._sample_size, no_files)
        crop = no_files % self._sample_size
        top_tail = list(range(no_files))[
            crop // 2:no_files - (crop - (crop // 2))]
        # Partition the indices
        size = len(top_tail)
        retval = [top_tail[start:start + size // self._sample_size]
                  for start in range(0, size, size // self._sample_size)]
        logger.debug("Indices pools: %s", [f"{idx}: (start: {min(pool)}, "
                                           f"end: {max(pool)}, size: {len(pool)})"
                                           for idx, pool in enumerate(retval)])
        return retval

    def generate(self) -> None:
        """ Generate a sample set.

        Selects :attr:`sample_size` random faces. Runs them through prediction to obtain the
        swap, then trigger the patch event to run the faces through patching.
        """
        logger.debug("Generating new random samples")
        self._app.dispatcher.set_busy()
        self._load_frames()
        self._predict()
        self._app.dispatcher.set_needs_patch()
        logger.debug("Generated new random samples")

    def _load_frames(self) -> None:
        """ Load a sample of random frames.

        * Picks a random face from each indices group.

        * Takes the first face from the image (if there are multiple faces). Adds the images to \
        :attr:`self._input_images`.

        * Sets :attr:`_display.source` to the input images and flags that the display should be \
        updated
        """
        self._input_images = []
        for selection in self._random_choice:
            filename = os.path.basename(self._filelist[selection])
            image = self._images.load_one_image(self._filelist[selection])
            # Get first face only
            face = self._alignments.get_faces_in_frame(filename)[0]
            detected_face = DetectedFace()
            detected_face.from_alignment(face, image=image)
            inbound = ExtractMedia(filename=filename, image=image, detected_faces=[detected_face])
            self._input_images.append(ConvertItem(inbound=inbound))
        self._app.display.source = self._input_images
        self._app.display.update_source = True
        logger.debug("Selected frames: %s",
                     [frame.inbound.filename for frame in self._input_images])

    def _predict(self) -> None:
        """ Predict from the loaded frames.

        With a threading lock (to prevent stacking), run the selected faces through the Faceswap
        model predict function and add the output to :attr:`predicted`
        """
        with self._app.lock:
            self._predicted_images = []
            for frame in self._input_images:
                self._predictor.in_queue.put(frame)
            idx = 0
            while idx < self._sample_size:
                logger.debug("Predicting face %s of %s", idx + 1, self._sample_size)
                items: (T.Literal["EOF"] |
                        list[tuple[ConvertItem, np.ndarray]]) = self._predictor.out_queue.get()
                if items == "EOF":
                    logger.debug("Received EOF")
                    break
                for item in items:
                    self._predicted_images.append(item)
                    logger.debug("Predicted face %s of %s", idx + 1, self._sample_size)
                    idx += 1
        logger.debug("Predicted faces")


class Patch():  # pylint:disable=too-few-public-methods
    """ The Patch pipeline

    Runs in it's own thread. Takes the output from the Faceswap model predictor and runs the faces
    through the convert pipeline using the currently selected options.

    Parameters
    ----------
    app: :class:`Preview`
        The main tkinter Preview app
    arguments: :class:`argparse.Namespace`
        The :mod:`argparse` arguments as passed in from :mod:`tools.py`

    Attributes
    ----------
    converter_arguments: dict
        The currently selected converter command line arguments for the patch queue
    """
    def __init__(self, app: Preview, arguments: Namespace) -> None:
        logger.debug("Initializing %s: (app: %s, arguments: '%s')",
                     self.__class__.__name__, app, arguments)
        self._app = app
        self._queue_patch_in = queue_manager.get_queue("preview_patch_in")
        self.converter_arguments: dict[str, T.Any] | None = None  # Updated converter args

        configfile = arguments.configfile if hasattr(arguments, "configfile") else None
        self._converter = Converter(output_size=app._samples.predictor.output_size,
                                    coverage_ratio=app._samples.predictor.coverage_ratio,
                                    centering=app._samples.predictor.centering,
                                    draw_transparent=False,
                                    pre_encode=None,
                                    arguments=self._generate_converter_arguments(
                                        arguments,
                                        app._samples.available_masks),
                                    configfile=configfile)
        self._thread = Thread(target=self._process,
                              name="patch_thread",
                              args=(self._queue_patch_in,
                                    self._app.dispatcher.needs_patch,
                                    app._samples),
                              daemon=True)
        self._thread.start()
        logger.debug("Initializing %s", self.__class__.__name__)

    @property
    def converter(self) -> Converter:
        """ :class:`lib.convert.Converter`: The converter to use for patching the images. """
        return self._converter

    @staticmethod
    def _generate_converter_arguments(arguments: Namespace,
                                      available_masks: list[str]) -> Namespace:
        """ Add the default converter arguments to the initial arguments. Ensure the mask selection
        is available.

        Parameters
        ----------
        arguments: :class:`argparse.Namespace`
            The :mod:`argparse` arguments as passed in from :mod:`tools.py`
        available_masks: list
            The masks that are available for convert
        Returns
        ----------
        arguments: :class:`argparse.Namespace`
            The :mod:`argparse` arguments as passed in with converter default
            arguments added
        """
        valid_masks = available_masks + ["none"]
        converter_arguments = ConvertArgs(None, "convert").get_optional_arguments()  # type: ignore
        for item in converter_arguments:
            value = item.get("default", None)
            # Skip options without a default value
            if value is None:
                continue
            option = item.get("dest", item["opts"][1].replace("--", ""))
            if option == "mask_type" and value not in valid_masks:
                logger.debug("Amending default mask from '%s' to '%s'", value, valid_masks[0])
                value = valid_masks[0]
            # Skip options already in arguments
            if hasattr(arguments, option):
                continue
            # Add option to arguments
            setattr(arguments, option, value)
        logger.debug(arguments)
        return arguments

    def _process(self,
                 patch_queue_in: EventQueue,
                 trigger_event: Event,
                 samples: Samples) -> None:
        """ The face patching process.

        Runs in a thread, and waits for an event to be set. Once triggered, runs a patching
        cycle and sets the :class:`Display` destination images.

        Parameters
        ----------
        patch_queue_in: :class:`~lib.queue_manager.EventQueue`
            The input queue for the patching process
        trigger_event: :class:`threading.Event`
            The event that indicates a patching run needs to be processed
        samples: :class:`Samples`
            The Samples for display.
        """
        logger.debug("Launching patch process thread: (patch_queue_in: %s, trigger_event: %s, "
                     "samples: %s)", patch_queue_in, trigger_event, samples)
        patch_queue_out = queue_manager.get_queue("preview_patch_out")
        while True:
            trigger = trigger_event.wait(1)
            if not trigger:
                continue
            logger.debug("Patch Triggered")
            queue_manager.flush_queue("preview_patch_in")
            self._feed_swapped_faces(patch_queue_in, samples)
            with self._app.lock:
                self._update_converter_arguments()
                self._converter.reinitialize(config=self._app.config_tools.config)
            swapped = self._patch_faces(patch_queue_in, patch_queue_out, samples.sample_size)
            with self._app.lock:
                self._app.display.destination = swapped

            logger.debug("Patch complete")
            trigger_event.clear()

        logger.debug("Closed patch process thread")

    def _update_converter_arguments(self) -> None:
        """ Update the converter arguments to the currently selected values. """
        logger.debug("Updating Converter cli arguments")
        if self.converter_arguments is None:
            logger.debug("No arguments to update")
            return
        for key, val in self.converter_arguments.items():
            logger.debug("Updating %s to %s", key, val)
            setattr(self._converter.cli_arguments, key, val)
        logger.debug("Updated Converter cli arguments")

    @staticmethod
    def _feed_swapped_faces(patch_queue_in: EventQueue, samples: Samples) -> None:
        """ Feed swapped faces to the converter's in-queue.

        Parameters
        ----------
        patch_queue_in: :class:`~lib.queue_manager.EventQueue`
            The input queue for the patching process
        samples: :class:`Samples`
            The Samples for display.
        """
        logger.debug("feeding swapped faces to converter")
        for item in samples.predicted_images:
            patch_queue_in.put(item)
        logger.debug("fed %s swapped faces to converter",
                     len(samples.predicted_images))
        logger.debug("Putting EOF to converter")
        patch_queue_in.put("EOF")

    def _patch_faces(self,
                     queue_in: EventQueue,
                     queue_out: EventQueue,
                     sample_size: int) -> list[np.ndarray]:
        """ Patch faces.

        Run the convert process on the swapped faces and return the patched faces.

        patch_queue_in: :class:`~lib.queue_manager.EventQueue`
            The input queue for the patching process
        queue_out: :class:`~lib.queue_manager.EventQueue`
            The output queue from the patching process
        sample_size: int
            The number of samples to be displayed

        Returns
        -------
        list
            The swapped faces patched with the selected convert settings
        """
        logger.debug("Patching faces")
        self._converter.process(queue_in, queue_out)
        swapped = []
        idx = 0
        while idx < sample_size:
            logger.debug("Patching image %s of %s", idx + 1, sample_size)
            item = queue_out.get()
            swapped.append(item[1])
            logger.debug("Patched image %s of %s", idx + 1, sample_size)
            idx += 1
        logger.debug("Patched faces")
        return swapped
