from scripts.configure import check_configuration

check_configuration()

import importlib
import traceback

from scripts import faceswaplab_globals
from scripts.faceswaplab_api import faceswaplab_api
from scripts.faceswaplab_postprocessing import upscaling
from scripts.faceswaplab_settings import faceswaplab_settings
from scripts.faceswaplab_swapping import swapper
from scripts.faceswaplab_ui import faceswaplab_tab, faceswaplab_unit_ui
from scripts.faceswaplab_utils import faceswaplab_logging, imgutils, models_utils
from scripts.faceswaplab_utils.models_utils import get_current_swap_model
from scripts.faceswaplab_utils.typing import *
from scripts.faceswaplab_utils.ui_utils import dataclasses_from_flat_list
from scripts.faceswaplab_utils.faceswaplab_logging import logger, save_img_debug

# Reload all the modules when using "apply and restart"
# This is mainly done for development purposes
import logging

if logger.getEffectiveLevel() <= logging.DEBUG:
    importlib.reload(swapper)
    importlib.reload(faceswaplab_logging)
    importlib.reload(faceswaplab_globals)
    importlib.reload(imgutils)
    importlib.reload(upscaling)
    importlib.reload(faceswaplab_settings)
    importlib.reload(models_utils)
    importlib.reload(faceswaplab_unit_ui)
    importlib.reload(faceswaplab_api)

import os
from pprint import pformat
from typing import Any, List, Optional, Tuple

import gradio as gr
import modules.scripts as scripts
from modules import script_callbacks, scripts, shared
from modules.images import save_image
from modules.processing import (
    Processed,
    StableDiffusionProcessing,
    StableDiffusionProcessingImg2Img,
)
from modules.shared import opts

from scripts.faceswaplab_globals import VERSION_FLAG
from scripts.faceswaplab_postprocessing.postprocessing import enhance_image
from scripts.faceswaplab_postprocessing.postprocessing_options import (
    PostProcessingOptions,
)
from scripts.faceswaplab_ui.faceswaplab_unit_settings import FaceSwapUnitSettings

EXTENSION_PATH = os.path.join("extensions", "sd-webui-faceswaplab")


# Register the tab, done here to prevent it from being added twice
script_callbacks.on_ui_tabs(faceswaplab_tab.on_ui_tabs)

try:
    import modules.script_callbacks as script_callbacks

    script_callbacks.on_app_started(faceswaplab_api.faceswaplab_api)
except:
    logger.error("Failed to register API")

    traceback.print_exc()


class FaceSwapScript(scripts.Script):
    def __init__(self) -> None:
        super().__init__()

    @property
    def units_count(self) -> int:
        return opts.data.get("faceswaplab_units_count", 3)

    @property
    def enabled(self) -> bool:
        """Return True if any unit is enabled and the state is not interupted"""
        return any([u.enable for u in self.units]) and not shared.state.interrupted

    @property
    def keep_original_images(self) -> bool:
        return opts.data.get("faceswaplab_keep_original", False)

    @property
    def swap_in_generated_units(self) -> List[FaceSwapUnitSettings]:
        return [u for u in self.units if u.swap_in_generated and u.enable]

    @property
    def swap_in_source_units(self) -> List[FaceSwapUnitSettings]:
        return [u for u in self.units if u.swap_in_source and u.enable]

    def title(self) -> str:
        return f"faceswaplab"

    def show(self, is_img2img: bool) -> bool:
        return scripts.AlwaysVisible  # type: ignore

    def ui(self, is_img2img: bool) -> List[gr.components.Component]:
        with gr.Accordion(f"FaceSwapLab {VERSION_FLAG}", open=False):
            components: List[gr.components.Component] = []
            for i in range(1, self.units_count + 1):
                components += faceswaplab_unit_ui.faceswap_unit_ui(is_img2img, i)
            post_processing = faceswaplab_tab.postprocessing_ui()
        # If the order is modified, the before_process should be changed accordingly.
        return components + post_processing

    def read_config(
        self, p: StableDiffusionProcessing, *components: Tuple[Any, ...]
    ) -> None:
        for i, c in enumerate(components):
            logger.debug("%s>%s", i, pformat(c))

        # The order of processing for the components is important
        # The method first process faceswap units then postprocessing units
        classes: List[Any] = dataclasses_from_flat_list(
            [FaceSwapUnitSettings] * self.units_count + [PostProcessingOptions],
            components,
        )
        self.units: List[FaceSwapUnitSettings] = []
        self.units += [u for u in classes if isinstance(u, FaceSwapUnitSettings)]
        self.postprocess_options = classes[-1]

        for i, u in enumerate(self.units):
            logger.debug("%s, %s", pformat(i), pformat(u))

        logger.debug("%s", pformat(self.postprocess_options))

        if self.enabled:
            p.do_not_save_samples = not self.keep_original_images

    def process(
        self, p: StableDiffusionProcessing, *components: Tuple[Any, ...]
    ) -> None:
        try:
            self.read_config(p, *components)

            # If is instance of img2img, we check if face swapping in source is required.
            if isinstance(p, StableDiffusionProcessingImg2Img):
                if self.enabled and len(self.swap_in_source_units) > 0:
                    init_images: List[Tuple[Optional[PILImage], Optional[str]]] = [
                        (img, None) for img in p.init_images
                    ]
                    new_inits = swapper.process_images_units(
                        get_current_swap_model(),
                        self.swap_in_source_units,
                        images=init_images,
                        force_blend=True,
                    )
                    logger.info(f"processed init images: {len(init_images)}")
                    if new_inits is not None:
                        p.init_images = [img[0] for img in new_inits]
        except Exception as e:
            logger.info("Failed to process : %s", e)
            traceback.print_exc()

    def postprocess(
        self, p: StableDiffusionProcessing, processed: Processed, *args: List[Any]
    ) -> None:
        try:
            if self.enabled:
                # Get the original images without the grid
                orig_images: List[PILImage] = processed.images[
                    processed.index_of_first_image :
                ]
                orig_infotexts: List[str] = processed.infotexts[
                    processed.index_of_first_image :
                ]

                keep_original = self.keep_original_images

                # These are were images and infos of swapped images will be stored
                images = []
                infotexts = []
                if (len(self.swap_in_generated_units)) > 0:
                    for i, (img, info) in enumerate(zip(orig_images, orig_infotexts)):
                        batch_index = i % p.batch_size
                        swapped_images = swapper.process_images_units(
                            get_current_swap_model(),
                            self.swap_in_generated_units,
                            images=[(img, info)],
                        )
                        if swapped_images is None:
                            continue

                        logger.info(f"{len(swapped_images)} images swapped")
                        for swp_img, new_info in swapped_images:
                            img = swp_img  # Will only swap the last image in the batch in next units (FIXME : hard to fix properly but not really critical)

                            if swp_img is not None:
                                save_img_debug(swp_img, "Before apply mask")
                                swp_img = imgutils.apply_mask(swp_img, p, batch_index)
                                save_img_debug(swp_img, "After apply mask")

                                try:
                                    if self.postprocess_options is not None:
                                        swp_img = enhance_image(
                                            swp_img, self.postprocess_options
                                        )
                                except Exception as e:
                                    logger.error("Failed to upscale : %s", e)

                                logger.info("Add swp image to processed")
                                images.append(swp_img)
                                infotexts.append(new_info)
                                if p.outpath_samples and opts.samples_save:
                                    save_image(
                                        swp_img,
                                        p.outpath_samples,
                                        "",
                                        p.all_seeds[batch_index],  # type: ignore
                                        p.all_prompts[batch_index],  # type: ignore
                                        opts.samples_format,
                                        info=new_info,
                                        p=p,
                                        suffix="-swapped",
                                    )
                            else:
                                logger.error("swp image is None")
                else:
                    keep_original = True

                # Generate grid :
                if opts.return_grid and len(images) > 1:
                    grid = imgutils.create_square_image(images)
                    text = processed.infotexts[0]
                    infotexts.insert(0, text)
                    if opts.enable_pnginfo:
                        grid.info["parameters"] = text  # type: ignore
                    images.insert(0, grid)

                    if opts.grid_save:
                        save_image(
                            grid,
                            p.outpath_grids,
                            "swapped-grid",
                            p.all_seeds[0],  # type: ignore
                            p.all_prompts[0],  # type: ignore
                            opts.grid_format,
                            info=text,
                            short_filename=not opts.grid_extended_filename,
                            p=p,
                            grid=True,
                        )

                if keep_original:
                    # If we want to keep original images, we add all existing (including grid this time)
                    images += processed.images
                    infotexts += processed.infotexts

                processed.images = images
                processed.infotexts = infotexts
        except Exception as e:
            logger.error("Failed to swap face in postprocess method : %s", e)
            traceback.print_exc()