import os
import tempfile
from pprint import pformat, pprint

import dill as pickle
import gradio as gr
import modules.scripts as scripts
import numpy as np
import onnx
import pandas as pd
from scripts.faceswaplab_ui.faceswaplab_unit_ui import faceswap_unit_ui
from scripts.faceswaplab_ui.faceswaplab_upscaler_ui import upscaler_ui
from insightface.app.common import Face
from modules import script_callbacks, scripts
from PIL import Image
from modules.shared import opts

from scripts.faceswaplab_utils import imgutils
from scripts.faceswaplab_utils.imgutils import pil_to_cv2
from scripts.faceswaplab_utils.models_utils import get_models
from scripts.faceswaplab_utils.faceswaplab_logging import logger
import scripts.faceswaplab_swapping.swapper as swapper
from scripts.faceswaplab_postprocessing.postprocessing_options import (
    PostProcessingOptions,
)
from scripts.faceswaplab_postprocessing.postprocessing import enhance_image
from dataclasses import fields
from typing import List
from scripts.faceswaplab_ui.faceswaplab_unit_settings import FaceSwapUnitSettings
from scripts.faceswaplab_utils.models_utils import get_current_model


def compare(img1, img2):
    if img1 is not None and img2 is not None:
        return swapper.compare_faces(img1, img2)

    return "You need 2 images to compare"


def extract_faces(
    files,
    extract_path,
    face_restorer_name,
    face_restorer_visibility,
    codeformer_weight,
    upscaler_name,
    upscaler_scale,
    upscaler_visibility,
    inpainting_denoising_strengh,
    inpainting_prompt,
    inpainting_negative_prompt,
    inpainting_steps,
    inpainting_sampler,
    inpainting_when,
):
    if not extract_path:
        tempfile.mkdtemp()
    if files is not None:
        images = []
        for file in files:
            img = Image.open(file.name).convert("RGB")
            faces = swapper.get_faces(pil_to_cv2(img))
            if faces:
                face_images = []
                for face in faces:
                    bbox = face.bbox.astype(int)
                    x_min, y_min, x_max, y_max = bbox
                    face_image = img.crop((x_min, y_min, x_max, y_max))
                    if face_restorer_name or face_restorer_visibility:
                        scale = 1 if face_image.width > 512 else 512 // face_image.width
                        face_image = enhance_image(
                            face_image,
                            PostProcessingOptions(
                                face_restorer_name=face_restorer_name,
                                restorer_visibility=face_restorer_visibility,
                                codeformer_weight=codeformer_weight,
                                upscaler_name=upscaler_name,
                                upscale_visibility=upscaler_visibility,
                                scale=scale,
                                inpainting_denoising_strengh=inpainting_denoising_strengh,
                                inpainting_prompt=inpainting_prompt,
                                inpainting_steps=inpainting_steps,
                                inpainting_negative_prompt=inpainting_negative_prompt,
                                inpainting_when=inpainting_when,
                                inpainting_sampler=inpainting_sampler,
                            ),
                        )
                    path = tempfile.NamedTemporaryFile(
                        delete=False, suffix=".png", dir=extract_path
                    ).name
                    face_image.save(path)
                    face_images.append(path)
                images += face_images
        return images
    return None


def analyse_faces(image, det_threshold=0.5):
    try:
        faces = swapper.get_faces(imgutils.pil_to_cv2(image), det_thresh=det_threshold)
        result = ""
        for i, face in enumerate(faces):
            result += f"\nFace {i} \n" + "=" * 40 + "\n"
            result += pformat(face) + "\n"
            result += "=" * 40
        return result

    except Exception as e:
        logger.error("Analysis Failed : %s", e)
        return "Analysis Failed"


def build_face_checkpoint_and_save(batch_files, name):
    """
    Builds a face checkpoint, swaps faces, and saves the result to a file.

    Args:
        batch_files (list): List of image file paths.
        name (str): Name of the face checkpoint

    Returns:
        PIL.Image.Image or None: Resulting swapped face image if successful, otherwise None.
    """
    batch_files = batch_files or []
    print("Build", name, [x.name for x in batch_files])
    faces = swapper.get_faces_from_img_files(batch_files)
    blended_face = swapper.blend_faces(faces)
    preview_path = os.path.join(
        scripts.basedir(), "extensions", "sd-webui-faceswaplab", "references"
    )
    faces_path = os.path.join(scripts.basedir(), "models", "faceswaplab", "faces")
    if not os.path.exists(faces_path):
        os.makedirs(faces_path)

    target_img = None
    if blended_face:
        if blended_face["gender"] == 0:
            target_img = Image.open(os.path.join(preview_path, "woman.png"))
        else:
            target_img = Image.open(os.path.join(preview_path, "man.png"))

        if name == "":
            name = "default_name"
        pprint(blended_face)
        result = swapper.swap_face(
            blended_face, blended_face, target_img, get_models()[0]
        )
        result_image = enhance_image(
            result.image,
            PostProcessingOptions(
                face_restorer_name="CodeFormer", restorer_visibility=1
            ),
        )

        file_path = os.path.join(faces_path, f"{name}.pkl")
        file_number = 1
        while os.path.exists(file_path):
            file_path = os.path.join(faces_path, f"{name}_{file_number}.pkl")
            file_number += 1
        result_image.save(file_path + ".png")
        with open(file_path, "wb") as file:
            pickle.dump(
                {
                    "embedding": blended_face.embedding,
                    "gender": blended_face.gender,
                    "age": blended_face.age,
                },
                file,
            )
        try:
            with open(file_path, "rb") as file:
                data = Face(pickle.load(file))
                print(data)
        except Exception as e:
            print(e)
        return result_image

    print("No face found")

    return target_img


def explore_onnx_faceswap_model(model_path):
    data = {
        "Node Name": [],
        "Op Type": [],
        "Inputs": [],
        "Outputs": [],
        "Attributes": [],
    }
    if model_path:
        model = onnx.load(model_path)
        for node in model.graph.node:
            data["Node Name"].append(pformat(node.name))
            data["Op Type"].append(pformat(node.op_type))
            data["Inputs"].append(pformat(node.input))
            data["Outputs"].append(pformat(node.output))
            attributes = []
            for attr in node.attribute:
                attr_name = attr.name
                attr_value = attr.t
                attributes.append(
                    "{} = {}".format(pformat(attr_name), pformat(attr_value))
                )
            data["Attributes"].append(attributes)

    df = pd.DataFrame(data)
    return df


def batch_process(files, save_path, *components):
    try:
        if save_path is not None:
            os.makedirs(save_path, exist_ok=True)

        units_count = opts.data.get("faceswaplab_units_count", 3)
        units: List[FaceSwapUnitSettings] = []

        # Parse and convert units flat components into FaceSwapUnitSettings
        for i in range(0, units_count):
            units += [FaceSwapUnitSettings.get_unit_configuration(i, components)]

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

        # Parse the postprocessing options
        # We must first find where to start from (after face swapping units)
        len_conf: int = len(fields(FaceSwapUnitSettings))
        shift: int = units_count * len_conf
        postprocess_options = PostProcessingOptions(
            *components[shift : shift + len(fields(PostProcessingOptions))]
        )
        logger.debug("%s", pformat(postprocess_options))

        units = [u for u in units if u.enable]
        if files is not None:
            images = []
            for file in files:
                current_images = []
                src_image = Image.open(file.name).convert("RGB")
                swapped_images = swapper.process_images_units(
                    get_current_model(),
                    images=[(src_image, None)],
                    units=units,
                    upscaled_swapper=opts.data.get(
                        "faceswaplab_upscaled_swapper", False
                    ),
                )
                if len(swapped_images) > 0:
                    current_images += [img for img, info in swapped_images]

                logger.info("%s images generated", len(current_images))
                for i, img in enumerate(current_images):
                    current_images[i] = enhance_image(img, postprocess_options)

                for img in current_images:
                    path = tempfile.NamedTemporaryFile(
                        delete=False, suffix=".png", dir=save_path
                    ).name
                    img.save(path)

                images += current_images
            return images
    except Exception as e:
        logger.error("Batch Process error : %s", e)
        import traceback

        traceback.print_exc()
    return None


def tools_ui():
    models = get_models()
    with gr.Tab("Tools"):
        with gr.Tab("Build"):
            gr.Markdown(
                """Build a face based on a batch list of images. Will blend the resulting face and store the checkpoint in the faceswaplab/faces directory."""
            )
            with gr.Row():
                batch_files = gr.components.File(
                    type="file",
                    file_count="multiple",
                    label="Batch Sources Images",
                    optional=True,
                    elem_id="faceswaplab_build_batch_files",
                )
                preview = gr.components.Image(
                    type="pil",
                    label="Preview",
                    interactive=False,
                    elem_id="faceswaplab_build_preview_face",
                )
            name = gr.Textbox(
                value="Face",
                placeholder="Name of the character",
                label="Name of the character",
                elem_id="faceswaplab_build_character_name",
            )
            generate_checkpoint_btn = gr.Button(
                "Save", elem_id="faceswaplab_build_save_btn"
            )
        with gr.Tab("Compare"):
            gr.Markdown(
                """Give a similarity score between two images (only first face is compared)."""
            )

            with gr.Row():
                img1 = gr.components.Image(
                    type="pil", label="Face 1", elem_id="faceswaplab_compare_face1"
                )
                img2 = gr.components.Image(
                    type="pil", label="Face 2", elem_id="faceswaplab_compare_face2"
                )
            compare_btn = gr.Button("Compare", elem_id="faceswaplab_compare_btn")
            compare_result_text = gr.Textbox(
                interactive=False,
                label="Similarity",
                value="0",
                elem_id="faceswaplab_compare_result",
            )
        with gr.Tab("Extract"):
            gr.Markdown(
                """Extract all faces from a batch of images. Will apply enhancement in the tools enhancement tab."""
            )
            with gr.Row():
                extracted_source_files = gr.components.File(
                    type="file",
                    file_count="multiple",
                    label="Batch Sources Images",
                    optional=True,
                    elem_id="faceswaplab_extract_batch_images",
                )
                extracted_faces = gr.Gallery(
                    label="Extracted faces",
                    show_label=False,
                    elem_id="faceswaplab_extract_results",
                ).style(columns=[2], rows=[2])
            extract_save_path = gr.Textbox(
                label="Destination Directory",
                value="",
                elem_id="faceswaplab_extract_destination",
            )
            extract_btn = gr.Button("Extract", elem_id="faceswaplab_extract_btn")
        with gr.Tab("Explore Model"):
            model = gr.Dropdown(
                choices=models,
                label="Model not found, please download one and reload automatic 1111",
                elem_id="faceswaplab_explore_model",
            )
            explore_btn = gr.Button("Explore", elem_id="faceswaplab_explore_btn")
            explore_result_text = gr.Dataframe(
                interactive=False,
                label="Explored",
                elem_id="faceswaplab_explore_result",
            )
        with gr.Tab("Analyse Face"):
            img_to_analyse = gr.components.Image(
                type="pil", label="Face", elem_id="faceswaplab_analyse_face"
            )
            analyse_det_threshold = gr.Slider(
                0.1,
                1,
                0.5,
                step=0.01,
                label="Detection threshold",
                elem_id="faceswaplab_analyse_det_threshold",
            )
            analyse_btn = gr.Button("Analyse", elem_id="faceswaplab_analyse_btn")
            analyse_results = gr.Textbox(
                label="Results",
                interactive=False,
                value="",
                elem_id="faceswaplab_analyse_results",
            )

    with gr.Tab("Batch Process"):
        with gr.Tab("Source Images"):
            gr.Markdown(
                """Batch process images. Will apply enhancement in the tools enhancement tab."""
            )
            with gr.Row():
                batch_source_files = gr.components.File(
                    type="file",
                    file_count="multiple",
                    label="Batch Sources Images",
                    optional=True,
                    elem_id="faceswaplab_batch_images",
                )
                batch_results = gr.Gallery(
                    label="Batch result",
                    show_label=False,
                    elem_id="faceswaplab_batch_results",
                ).style(columns=[2], rows=[2])
            batch_save_path = gr.Textbox(
                label="Destination Directory",
                value="outputs/faceswap/",
                elem_id="faceswaplab_batch_destination",
            )
            batch_save_btn = gr.Button(
                "Process & Save", elem_id="faceswaplab_extract_btn"
            )
        unit_components = []
        for i in range(1, opts.data.get("faceswaplab_units_count", 3) + 1):
            unit_components += faceswap_unit_ui(False, i, id_prefix="faceswaplab_tab")

    upscale_options = upscaler_ui()

    explore_btn.click(
        explore_onnx_faceswap_model, inputs=[model], outputs=[explore_result_text]
    )
    compare_btn.click(compare, inputs=[img1, img2], outputs=[compare_result_text])
    generate_checkpoint_btn.click(
        build_face_checkpoint_and_save, inputs=[batch_files, name], outputs=[preview]
    )
    extract_btn.click(
        extract_faces,
        inputs=[extracted_source_files, extract_save_path] + upscale_options,
        outputs=[extracted_faces],
    )
    analyse_btn.click(
        analyse_faces,
        inputs=[img_to_analyse, analyse_det_threshold],
        outputs=[analyse_results],
    )
    batch_save_btn.click(
        batch_process,
        inputs=[batch_source_files, batch_save_path]
        + unit_components
        + upscale_options,
        outputs=[batch_results],
    )


def on_ui_tabs():
    with gr.Blocks(analytics_enabled=False) as ui_faceswap:
        tools_ui()
    return [(ui_faceswap, "FaceSwapLab", "faceswaplab_tab")]