import glob
import os
from typing import *
from insightface.app.common import Face
from safetensors.torch import save_file, safe_open
import torch

import modules.scripts as scripts
from modules import scripts
from scripts.faceswaplab_swapping.upcaled_inswapper_options import InswappperOptions
from scripts.faceswaplab_utils.faceswaplab_logging import logger
from scripts.faceswaplab_utils.typing import *
from scripts.faceswaplab_utils import imgutils
from scripts.faceswaplab_utils.models_utils import get_swap_models
import traceback

import dill as pickle  # will be removed in future versions
from scripts.faceswaplab_swapping import swapper
from pprint import pformat
import re
from client_api import api_utils
import tempfile


def sanitize_name(name: str) -> str:
    """
    Sanitize the input name by removing special characters and replacing spaces with underscores.

    Parameters:
        name (str): The input name to be sanitized.

    Returns:
        str: The sanitized name with special characters removed and spaces replaced by underscores.
    """
    name = re.sub("[^A-Za-z0-9_. ]+", "", name)
    name = name.replace(" ", "_")
    return name[:255]


def build_face_checkpoint_and_save(
    images: List[PILImage],
    name: str,
    overwrite: bool = False,
    path: Optional[str] = None,
) -> Optional[PILImage]:
    """
    Builds a face checkpoint using the provided image files, performs face swapping,
    and saves the result to a file. If a blended face is successfully obtained and the face swapping
    process succeeds, the resulting image is returned. Otherwise, None is returned.

    Args:
        batch_files (list): List of image file paths used to create the face checkpoint.
        name (str): The name assigned to the face checkpoint.

    Returns:
        PIL.PILImage or None: The resulting swapped face image if the process is successful; None otherwise.
    """

    try:
        name = sanitize_name(name)
        images = images or []
        logger.info("Build %s with %s images", name, len(images))
        faces: List[Face] = swapper.get_faces_from_img_files(images=images)
        if faces is None or len(faces) == 0:
            logger.error("No source faces found")
            return None

        blended_face: Optional[Face] = swapper.blend_faces(faces)
        preview_path = os.path.join(
            scripts.basedir(), "extensions", "sd-webui-faceswaplab", "references"
        )

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

            if name == "":
                name = "default_name"
            logger.debug("Face %s", pformat(blended_face))
            target_face = swapper.get_or_default(
                swapper.get_faces(imgutils.pil_to_cv2(reference_preview_img)), 0, None
            )
            if target_face is None:
                logger.error(
                    "Failed to open reference image, cannot create preview : That should not happen unless you deleted the references folder or change the detection threshold."
                )
            else:
                result: swapper.ImageResult = swapper.swap_face(
                    target_faces=[target_face],
                    source_face=blended_face,
                    target_img=reference_preview_img,
                    model=get_swap_models()[0],
                    swapping_options=InswappperOptions(
                        face_restorer_name="CodeFormer",
                        restorer_visibility=1,
                        upscaler_name="Lanczos",
                        codeformer_weight=1,
                        improved_mask=True,
                        color_corrections=False,
                        sharpen=True,
                    ),
                )
                preview_image = result.image

                if path:
                    file_path = path
                else:
                    file_path = os.path.join(
                        get_checkpoint_path(), f"{name}.safetensors"
                    )
                    if not overwrite:
                        file_number = 1
                        while os.path.exists(file_path):
                            file_path = os.path.join(
                                get_checkpoint_path(),
                                f"{name}_{file_number}.safetensors",
                            )
                            file_number += 1
                save_face(filename=file_path, face=blended_face)
                preview_image.save(file_path + ".png")
                try:
                    data = load_face(file_path)
                    logger.debug(data)
                except Exception as e:
                    logger.error("Error loading checkpoint, after creation %s", e)
                    traceback.print_exc()

                return preview_image

        else:
            logger.error("No face found")
            return None  # type: ignore
    except Exception as e:
        logger.error("Failed to build checkpoint %s", e)
        traceback.print_exc()
        return None


def save_face(face: Face, filename: str) -> None:
    try:
        tensors = {
            "embedding": torch.tensor(face["embedding"]),
            "gender": torch.tensor(face["gender"]),
            "age": torch.tensor(face["age"]),
        }
        save_file(tensors, filename)
    except Exception as e:
        traceback.print_exc
        logger.error("Failed to save checkpoint %s", e)
        raise e


def load_face(name: str) -> Optional[Face]:
    if name.startswith("data:application/face;base64,"):
        with tempfile.NamedTemporaryFile(delete=True) as temp_file:
            api_utils.base64_to_safetensors(name, temp_file.name)
            face = {}
            with safe_open(temp_file.name, framework="pt", device="cpu") as f:
                for k in f.keys():
                    logger.debug("load key %s", k)
                    face[k] = f.get_tensor(k).numpy()
            return Face(face)

    filename = matching_checkpoint(name)
    if filename is None:
        return None

    if filename.endswith(".pkl"):
        logger.warning(
            "Pkl files for faces are deprecated to enhance safety, they will be unsupported in future versions."
        )
        logger.warning("The file will be converted to .safetensors")
        logger.warning(
            "You can also use this script https://gist.github.com/glucauze/4a3c458541f2278ad801f6625e5b9d3d"
        )
        with open(filename, "rb") as file:
            logger.info("Load pkl")
            face = Face(pickle.load(file))
            logger.warning(
                "Convert to safetensors, you can remove the pkl version once you have ensured that the safetensor is working"
            )
            save_face(face, filename.replace(".pkl", ".safetensors"))
        return face

    elif filename.endswith(".safetensors"):
        face = {}
        with safe_open(filename, framework="pt", device="cpu") as f:
            for k in f.keys():
                logger.debug("load key %s", k)
                face[k] = f.get_tensor(k).numpy()
        return Face(face)

    raise NotImplementedError("Unknown file type, face extraction not implemented")


def get_checkpoint_path() -> str:
    checkpoint_path = os.path.join(scripts.basedir(), "models", "faceswaplab", "faces")
    os.makedirs(checkpoint_path, exist_ok=True)
    return checkpoint_path


def matching_checkpoint(name: str) -> Optional[str]:
    """
    Retrieve the full path of a checkpoint file matching the given name.

    If the name already includes a path separator, it is returned as-is. Otherwise, the function looks for a matching
    file with the extensions ".safetensors" or ".pkl" in the checkpoint directory.

    Args:
        name (str): The name or path of the checkpoint file.

    Returns:
        Optional[str]: The full path of the matching checkpoint file, or None if no match is found.
    """

    # If the name already includes a path separator, return it as is
    if os.path.sep in name:
        return name

    # If the name doesn't end with the specified extensions, look for a matching file
    if not (name.endswith(".safetensors") or name.endswith(".pkl")):
        # Try appending each extension and check if the file exists in the checkpoint path
        for ext in [".safetensors", ".pkl"]:
            full_path = os.path.join(get_checkpoint_path(), name + ext)
            if os.path.exists(full_path):
                return full_path
        # If no matching file is found, return None
        return None

    # If the name already ends with the specified extensions, simply complete the path
    return os.path.join(get_checkpoint_path(), name)


def get_face_checkpoints() -> List[str]:
    """
    Retrieve a list of face checkpoint paths.

    This function searches for face files with the extension ".safetensors" in the specified directory and returns a list
    containing the paths of those files.

    Returns:
        list: A list of face paths, including the string "None" as the first element.
    """
    faces_path = os.path.join(get_checkpoint_path(), "*.safetensors")
    faces = glob.glob(faces_path)

    faces_path = os.path.join(get_checkpoint_path(), "*.pkl")
    faces += glob.glob(faces_path)

    return ["None"] + [os.path.basename(face) for face in sorted(faces)]
