import base64
import json
import sys
from collections import defaultdict
from io import BytesIO
from pprint import pprint
from typing import Any, Dict, List
import os
import re
from pathlib import Path
from typing import Union
from concurrent.futures import ThreadPoolExecutor
import numpy as np
from PIL import ImageFilter


import torch
from diffusers import (
    DiffusionPipeline,
    DPMSolverMultistepScheduler,
    DPMSolverSinglestepScheduler,
    EulerAncestralDiscreteScheduler,
    utils,
)
from safetensors.torch import load_file
from torch import autocast, tensor
import torchvision.transforms
from PIL import Image

REPO_DIR = Path(__file__).resolve().parent

# if local avoid repo url
# print(os.getcwd())

# set device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

if device.type != "cuda":
    raise ValueError("need to run on GPU")


class EndpointHandler:
    LORA_PATHS = {
        "hairdetailer": str(REPO_DIR / "lora/hairdetailer.safetensors"),
        "lora_leica": str(REPO_DIR / "lora/lora_leica.safetensors"),
        "epiNoiseoffset_v2": str(REPO_DIR / "lora/epiNoiseoffset_v2.safetensors"),
        "MBHU-TT2FRS": str(REPO_DIR / "lora/MBHU-TT2FRS.safetensors"),
        "ShinyOiledSkin_v20": str(
            REPO_DIR / "lora/ShinyOiledSkin_v20-LoRA.safetensors"
        ),
        "polyhedron_new_skin_v1.1": str(
            REPO_DIR / "lora/polyhedron_new_skin_v1.1.safetensors"
        ),
        "detailed_eye-10": str(REPO_DIR / "lora/detailed_eye-10.safetensors"),
        "add_detail": str(REPO_DIR / "lora/add_detail.safetensors"),
        "MuscleGirl_v1": str(REPO_DIR / "lora/MuscleGirl_v1.safetensors"),
        "flat2": str(REPO_DIR / "lora/flat2.safetensors"),
    }

    TEXTUAL_INVERSION = [
        {
            "weight_name": str(REPO_DIR / "embeddings/EasyNegative.safetensors"),
            "token": "easynegative",
        },
        {
            "weight_name": str(REPO_DIR / "embeddings/badhandv4.pt"),
            "token": "badhandv4",
        },
        {
            "weight_name": str(REPO_DIR / "embeddings/bad-artist-anime.pt"),
            "token": "bad-artist-anime",
        },
        {
            "weight_name": str(REPO_DIR / "embeddings/NegfeetV2.pt"),
            "token": "negfeetv2",
        },
        {
            "weight_name": str(REPO_DIR / "embeddings/ng_deepnegative_v1_75t.pt"),
            "token": "ng_deepnegative_v1_75t",
        },
        {
            "weight_name": str(REPO_DIR / "embeddings/bad-hands-5.pt"),
            "token": "bad-hands-5",
        },
    ]

    def __init__(self, path="."):
        self.inference_progress = {}  # Dictionary to store progress of each request
        self.inference_images = {}  # Dictionary to store latest image of each request
        self.total_steps = {}
        self.active_request_ids = list()
        self.inference_in_progress = False

        self.executor = ThreadPoolExecutor(
            max_workers=1
        )  # Vous pouvez ajuster max_workers en fonction de vos besoins

        # load the optimized model
        self.pipe = DiffusionPipeline.from_pretrained(
            path,
            custom_pipeline="lpw_stable_diffusion",  # avoid 77 token limit
            torch_dtype=torch.float16,  # accelerate render
        )
        self.pipe = self.pipe.to(device)

        # https://stablediffusionapi.com/docs/a1111schedulers/

        # DPM++ 2M SDE Karras
        # increase step to avoid high contrast num_inference_steps=30
        # self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(
        #     self.pipe.scheduler.config,
        #     use_karras_sigmas=True,
        #     algorithm_type="sde-dpmsolver++",
        # )

        # DPM++ 2M Karras
        self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(
            self.pipe.scheduler.config,
            use_karras_sigmas=True,
        )

        # Mode boulardus
        self.pipe.safety_checker = None

        # Disable progress bar
        self.pipe.set_progress_bar_config(disable=True)

        # Load negative embeddings to avoid bad hands, etc
        self.load_embeddings()

        # boosts performance by another 20%
        self.pipe.enable_xformers_memory_efficient_attention()
        self.pipe.enable_attention_slicing()
        # may need a requirement in the root with xformer

        # Load loras one time only
        # Must be replaced once we will know how to hot load/unload
        # it use the own made load_lora function
        self.load_selected_loras(
            [
                ["polyhedron_new_skin_v1.1", 0.2],
                ["detailed_eye-10", 0.2],
                ["add_detail", 0.3],
                ["MuscleGirl_v1", 0.2],
            ]
        )

    def load_lora(self, pipeline, lora_path, lora_weight=0.5):
        state_dict = load_file(lora_path)
        LORA_PREFIX_UNET = "lora_unet"
        LORA_PREFIX_TEXT_ENCODER = "lora_te"

        alpha = lora_weight
        visited = []

        for key in state_dict:
            state_dict[key] = state_dict[key].to(device)

        # directly update weight in diffusers model
        for key in state_dict:
            # as we have set the alpha beforehand, so just skip
            if ".alpha" in key or key in visited:
                continue

            if "text" in key:
                layer_infos = (
                    key.split(".")[0]
                    .split(LORA_PREFIX_TEXT_ENCODER + "_")[-1]
                    .split("_")
                )
                curr_layer = pipeline.text_encoder
            else:
                layer_infos = (
                    key.split(".")[0].split(LORA_PREFIX_UNET + "_")[-1].split("_")
                )
                curr_layer = pipeline.unet

            # find the target layer
            temp_name = layer_infos.pop(0)
            while len(layer_infos) > -1:
                try:
                    curr_layer = curr_layer.__getattr__(temp_name)
                    if len(layer_infos) > 0:
                        temp_name = layer_infos.pop(0)
                    elif len(layer_infos) == 0:
                        break
                except Exception:
                    if len(temp_name) > 0:
                        temp_name += "_" + layer_infos.pop(0)
                    else:
                        temp_name = layer_infos.pop(0)

            # org_forward(x) + lora_up(lora_down(x)) * multiplier
            pair_keys = []
            if "lora_down" in key:
                pair_keys.append(key.replace("lora_down", "lora_up"))
                pair_keys.append(key)
            else:
                pair_keys.append(key)
                pair_keys.append(key.replace("lora_up", "lora_down"))

            # update weight
            if len(state_dict[pair_keys[0]].shape) == 4:
                weight_up = (
                    state_dict[pair_keys[0]].squeeze(3).squeeze(2).to(torch.float32)
                )
                weight_down = (
                    state_dict[pair_keys[1]].squeeze(3).squeeze(2).to(torch.float32)
                )
                curr_layer.weight.data += alpha * torch.mm(
                    weight_up, weight_down
                ).unsqueeze(2).unsqueeze(3)
            else:
                weight_up = state_dict[pair_keys[0]].to(torch.float32)
                weight_down = state_dict[pair_keys[1]].to(torch.float32)
                curr_layer.weight.data += alpha * torch.mm(weight_up, weight_down)

            # update visited list
            for item in pair_keys:
                visited.append(item)

        return pipeline

    def load_embeddings(self):
        """Load textual inversions, avoid bad prompts"""
        for model in EndpointHandler.TEXTUAL_INVERSION:
            self.pipe.load_textual_inversion(
                ".", weight_name=model["weight_name"], token=model["token"]
            )

    def load_selected_loras(self, selections):
        """Load Loras models, can lead to marvelous creations"""
        for model_name, weight in selections:
            lora_path = EndpointHandler.LORA_PATHS[model_name]
            # self.pipe.load_lora_weights(lora_path)
            self.load_lora(self.pipe, lora_path, weight)

    def clean_negative_prompt(self, negative_prompt):
        """Clean negative prompt to remove already used negative prompt handlers"""

        # negative_prompt = (
        #     negative_prompt
        #     + """, easynegative, badhandv4, bad-artist-anime, negfeetv2, ng_deepnegative_v1_75t, bad-hands-5, """
        # )

        tokens = [item["token"] for item in self.TEXTUAL_INVERSION]

        # Retirer tous les tokens de negative_prompt s'ils existent déjà
        for token in tokens:
            # Utiliser une expression régulière pour un remplacement insensible à la casse
            negative_prompt = re.sub(
                r"\b" + re.escape(token) + r"\b",
                "",
                negative_prompt,
                flags=re.IGNORECASE,
            ).strip()

        # Ajouter tous les tokens à la fin de negative_prompt
        negative_prompt += " " + " ".join(tokens)

        return negative_prompt

    def clean_request_data(self):
        """Clean up the data related to a specific request ID."""

        # Remove the request ID from the progress dictionary
        self.inference_progress.clear()

        # Remove the request ID from the images dictionary
        self.inference_images.clear()

        # Remove the request ID from the total_steps dictionary
        self.total_steps.clear()

        # Delete request id
        self.active_request_ids.clear()

        # Set inference to False
        self.inference_in_progress = False

    def progress_callback(
        self,
        step: int,
        timestep: int,
        latents: Any,
        request_id: str,
        status: str,
    ):
        try:
            if status == "progress":
                # Latents to numpy
                img_data = self.pipe.decode_latents(latents)
                img_data = (img_data.squeeze() * 255).astype(np.uint8)
                img = Image.fromarray(img_data, "RGB")

                # Apply a blur to the image
                # more intense at the beginning
                if step < int(self.total_steps[self.active_request_ids[0]] / 1.5):
                    img = img.filter(ImageFilter.GaussianBlur(radius=30))
                else:
                    img = img.filter(ImageFilter.GaussianBlur(radius=10))

                # print(img_data)
            else:
                # pil object
                # print(latents)
                img = latents

            buffered = BytesIO()
            img.save(buffered, format="PNG")

            # print(status)
            # Save the image to a file
            # img.save("squirel.png", format="PNG")

            # Encode the image into a base64 string representation
            img_str = base64.b64encode(buffered.getvalue()).decode()

        except Exception as e:
            print(f"Error: {e}")

        # Store progress and image
        progress_percentage = (
            step / self.total_steps[request_id]
        ) * 100  # Assuming self.total_steps is the total number of steps for inference

        self.inference_progress[request_id] = progress_percentage
        self.inference_images[request_id] = img_str

    def check_progress(self, request_id: str) -> Dict[str, Union[str, float]]:
        progress = self.inference_progress.get(request_id, 0)
        latest_image = self.inference_images.get(request_id, None)

        # print(self.inference_progress)

        if progress >= 100:
            status = "complete"
        else:
            status = "in-progress"

        return {
            "flag": "success",
            "status": status,
            "progress": int(progress),
            "image": latest_image,
        }

    def start_inference(self, data: Dict) -> Dict:
        """Start a new inference."""

        global device

        # Which Lora do we load ?
        # selected_models = [
        #     ("ShinyOiledSkin_v20", 0.3),
        #     ("MBHU-TT2FRS", 0.5),
        #     ("hairdetailer", 0.5),
        #     ("lora_leica", 0.5),
        #     ("epiNoiseoffset_v2", 0.5),
        # ]

        # 1. Verify input arguments
        required_fields = [
            "prompt",
            "negative_prompt",
            "width",
            "num_inference_steps",
            "height",
            "guidance_scale",
            "request_id",
        ]

        missing_fields = [field for field in required_fields if field not in data]

        if missing_fields:
            return {
                "flag": "error",
                "message": f"Missing fields: {', '.join(missing_fields)}",
            }

        # Now extract the fields
        prompt = data["prompt"]
        negative_prompt = data["negative_prompt"]
        loras_model = data.get("loras_model", None)
        seed = data.get("seed", None)
        width = data["width"]
        num_inference_steps = data["num_inference_steps"]
        height = data["height"]
        guidance_scale = data["guidance_scale"]
        request_id = data["request_id"]

        # Used for progress checker
        self.total_steps[request_id] = num_inference_steps

        # USe this to add automatically some negative prompts
        forced_negative = self.clean_negative_prompt(negative_prompt)

        # Set the generator seed if provided
        generator = torch.Generator(device="cuda").manual_seed(seed) if seed else None

        # Load the provided Lora models
        # self.pipe.unload_lora_weights()  # Unload models to avoid lora staking
        # if loras_model:
        #     self.load_selected_loras(loras_model)

        # set scale of loras, for now take only first scale of the loaded lora and apply to all until we find the way to apply specified scale
        # scale = {"scale": loras_model[0][1]} if loras_model else None

        try:
            # 2. Process
            with autocast(device.type):
                image = self.pipe.text2img(
                    prompt=prompt,
                    guidance_scale=guidance_scale,
                    num_inference_steps=num_inference_steps,
                    height=height,
                    width=width,
                    negative_prompt=forced_negative,
                    generator=generator,
                    max_embeddings_multiples=5,
                    callback=lambda step, timestep, latents: self.progress_callback(
                        step, timestep, latents, request_id, "progress"
                    ),
                    callback_steps=5,
                    # cross_attention_kwargs={"scale": 0.2},
                ).images[0]

            # print(image)
            self.progress_callback(
                num_inference_steps, 0, image, request_id, "complete"
            )

            self.inference_in_progress = False

            # for debug
            # image.save("squirelb.png", format="PNG")

        except Exception as e:
            # Handle any other exceptions and return an error response
            return {"flag": "error", "message": str(e)}

    def __call__(self, data: Any) -> Dict:
        """Handle incoming requests."""

        action = data.get("action", None)
        request_id = data.get("request_id")

        # Check if the request_id is valid for all actions
        if not request_id:
            return {"flag": "error", "message": "Missing request_id."}

        if action == "check_progress":
            if request_id not in self.active_request_ids:
                return {
                    "flag": "error",
                    "message": "Request id doesn't match any active request.",
                }
            return self.check_progress(request_id)

        elif action == "inference":
            # Check if an inference is already in progress
            if self.inference_in_progress:
                return {
                    "flag": "error",
                    "message": "Another inference is already in progress. Please wait.",
                }

            # Set the inference state to in progress
            self.clean_request_data()
            self.inference_in_progress = True
            self.inference_progress[request_id] = 0
            self.inference_images[request_id] = None
            self.active_request_ids.append(request_id)

            self.executor.submit(self.start_inference, data)

            return {
                "flag": "success",
                "message": "Inference started",
                "request_id": request_id,
            }

        else:
            return {"flag": "error", "message": f"Unsupported action: {action}"}