import cv2
import torch
import numpy as np
import PIL
from PIL import Image
from typing import Tuple
import diffusers
from diffusers.utils import load_image
from diffusers.models import ControlNetModel
from diffusers.pipelines.controlnet.multicontrolnet import MultiControlNetModel
from insightface.app import FaceAnalysis
from style_template import styles
from pipeline_stable_diffusion_xl_instantid_full import StableDiffusionXLInstantIDPipeline, draw_kps
from controlnet_aux import OpenposeDetector
import torch.nn.functional as F
from torchvision.transforms import Compose
import os
from huggingface_hub import hf_hub_download
import base64
import io
import json

# global variable
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if str(device).__contains__("cuda") else torch.float32
STYLE_NAMES = list(styles.keys())
DEFAULT_STYLE_NAME = "Spring Festival"

class EndpointHandler:
    def __init__(self, model_dir):
        # Ensure the necessary files are downloaded
        controlnet_config = hf_hub_download(repo_id="InstantX/InstantID", filename="ControlNetModel/config.json", local_dir=os.path.join(model_dir, "checkpoints"))
        controlnet_model = hf_hub_download(repo_id="InstantX/InstantID", filename="ControlNetModel/diffusion_pytorch_model.safetensors", local_dir=os.path.join(model_dir, "checkpoints"))
        face_adapter = hf_hub_download(repo_id="InstantX/InstantID", filename="ip-adapter.bin", local_dir=os.path.join(model_dir, "checkpoints"))

        dir_path = os.path.join(model_dir, "models", "antelopev2")
        if not os.path.exists(dir_path):
            print(f"Model path {dir_path} does not exist. Attempting to download.")
            self.app = FaceAnalysis(name='antelopev2', root=model_dir, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
        else:
            print(f"Model path {dir_path} exists. Skipping download.")
            self.app = FaceAnalysis(name='antelopev2', root=model_dir, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])

        self.app.prepare(ctx_id=0, det_size=(640, 640))
        openpose = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")

        # Path to InstantID models
        controlnet_path = os.path.join(model_dir, "checkpoints", "ControlNetModel")

        # Load pipeline face ControlNetModel
        self.controlnet_identitynet = ControlNetModel.from_pretrained(
            controlnet_path, torch_dtype=dtype
        )

        # controlnet-pose
        controlnet_pose_model = "thibaud/controlnet-openpose-sdxl-1.0"
        controlnet_canny_model = "diffusers/controlnet-canny-sdxl-1.0"

        controlnet_pose = ControlNetModel.from_pretrained(
            controlnet_pose_model, torch_dtype=dtype
        ).to(device)
        controlnet_canny = ControlNetModel.from_pretrained(
            controlnet_canny_model, torch_dtype=dtype
        ).to(device)

        def get_canny_image(image, t1=100, t2=200):
            image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
            edges = cv2.Canny(image, t1, t2)
            return Image.fromarray(edges, "L")

        self.controlnet_map = {
            "pose": controlnet_pose,
            "canny": controlnet_canny
        }

        self.controlnet_map_fn = {
            "pose": openpose,
            "canny": get_canny_image
        }

        pretrained_model_name_or_path = "wangqixun/YamerMIX_v8"

        self.pipe = StableDiffusionXLInstantIDPipeline.from_pretrained(
            pretrained_model_name_or_path,
            controlnet=[self.controlnet_identitynet],
            torch_dtype=dtype,
            safety_checker=None,
            feature_extractor=None,
        ).to(device)

        self.pipe.scheduler = diffusers.EulerDiscreteScheduler.from_config(
            self.pipe.scheduler.config
        )

        # load and disable LCM
        self.pipe.load_lora_weights("latent-consistency/lcm-lora-sdxl")
        self.pipe.disable_lora()

        self.pipe.cuda()
        self.pipe.load_ip_adapter_instantid(face_adapter)
        self.pipe.image_proj_model.to("cuda")
        self.pipe.unet.to("cuda")

    def __call__(self, data):

        def convert_from_cv2_to_image(img: np.ndarray) -> Image:
            return Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))

        def convert_from_image_to_cv2(img: Image) -> np.ndarray:
            return cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)

        def resize_img(
            input_image,
            max_side=1280,
            min_side=1024,
            size=None,
            pad_to_max_side=False,
            mode=PIL.Image.BILINEAR,
            base_pixel_number=64,
        ):
            w, h = input_image.size
            if size is not None:
                w_resize_new, h_resize_new = size
            else:
                ratio = min_side / min(h, w)
                w, h = round(ratio * w), round(ratio * h)
                ratio = max_side / max(h, w)
                input_image = input_image.resize([round(ratio * w), round(ratio * h)], mode)
                w_resize_new = (round(ratio * w) // base_pixel_number) * base_pixel_number
                h_resize_new = (round(ratio * h) // base_pixel_number) * base_pixel_number
            input_image = input_image.resize([w_resize_new, h_resize_new], mode)

            if pad_to_max_side:
                res = np.ones([max_side, max_side, 3], dtype=np.uint8) * 255
                offset_x = (max_side - w_resize_new) // 2
                offset_y = (max_side - h_resize_new) // 2
                res[
                    offset_y : offset_y + h_resize_new, offset_x : offset_x + w_resize_new
                ] = np.array(input_image)
                input_image = Image.fromarray(res)
            return input_image

        def apply_style(
            style_name: str, positive: str, negative: str = ""
        ) -> Tuple[str, str]:
            p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
            return p.replace("{prompt}", positive), n + " " + negative

        face_image_path = data.pop("face_image_path", "https://i.ibb.co/GQzm527/examples-musk-resize.jpg")
        pose_image_path = data.pop("pose_image_path", "https://i.ibb.co/TRCK4MS/examples-poses-pose2.jpg")
        style_name = data.pop("style_name", DEFAULT_STYLE_NAME)
        prompt = data.pop("inputs", "a man flying in the sky in Mars")
        negative_prompt = data.pop("negative_prompt", "(lowres, low quality, worst quality:1.2), (text:1.2), watermark, (frame:1.2), deformed, ugly, deformed eyes, blur, out of focus, blurry, deformed cat, deformed, photo, anthropomorphic cat, monochrome, pet collar, gun, weapon, blue, 3d, drones, drone, buildings in background, green")

        identitynet_strength_ratio = data.pop("identitynet_strength_ratio", 0.8)
        adapter_strength_ratio = data.pop("adapter_strength_ratio", 0.8)
        pose_strength = data.pop("pose_strength", 0.5)
        canny_strength = data.pop("canny_strength", 0.3)
        num_steps = data.pop("num_steps", 20)
        guidance_scale = data.pop("guidance_scale", 5.0)
        controlnet_selection = data.pop("controlnet_selection", ["pose", "canny"])
        scheduler = data.pop("scheduler", "EulerDiscreteScheduler")
        enable_fast_inference = data.pop("enable_fast_inference", False)
        enhance_non_face_region = data.pop("enhance_non_face_region", False)
        seed = data.pop("seed", 42)

        # Ensure required fields are present
        data.setdefault("prompt", prompt)
        data.setdefault("style", style_name)
        data.setdefault("num_steps", num_steps)
        data.setdefault("enable_LCM", enable_fast_inference)
        data.setdefault("enhance_face_region", enhance_non_face_region)

        # Enable LCM if fast inference is enabled
        if enable_fast_inference:
            self.pipe.enable_lora()
        else:
            self.pipe.disable_lora()

        scheduler_class_name = scheduler.split("-")[0]

        add_kwargs = {}
        if len(scheduler.split("-")) > 1:
            add_kwargs["use_karras_sigmas"] = True
        if len(scheduler.split("-")) > 2:
            add_kwargs["algorithm_type"] = "sde-dpmsolver++"
        scheduler = getattr(diffusers, scheduler_class_name)
        self.pipe.scheduler = scheduler.from_config(self.pipe.scheduler.config, **add_kwargs)

        # apply the style template
        prompt, negative_prompt = apply_style(style_name, prompt, negative_prompt)

        face_image = load_image(face_image_path)
        face_image = resize_img(face_image, max_side=1024)
        face_image_cv2 = convert_from_image_to_cv2(face_image)
        height, width, _ = face_image_cv2.shape

        # Extract face features
        face_info = self.app.get(face_image_cv2)

        face_info = sorted(
            face_info,
            key=lambda x: (x["bbox"][2] - x["bbox"][0]) * x["bbox"][3] - x["bbox"][1],
        )[
            -1
        ]  # only use the maximum face
        face_emb = face_info["embedding"]
        face_kps = draw_kps(convert_from_cv2_to_image(face_image_cv2), face_info["kps"])
        img_controlnet = face_image
        if pose_image_path is not None:
            pose_image = load_image(pose_image_path)
            pose_image = resize_img(pose_image, max_side=1024)
            img_controlnet = pose_image
            pose_image_cv2 = convert_from_image_to_cv2(pose_image)

            face_info = self.app.get(pose_image_cv2)

            face_info = face_info[-1]
            face_kps = draw_kps(pose_image, face_info["kps"])

            width, height = face_kps.size

        control_mask = np.zeros([height, width, 3])
        x1, y1, x2, y2 = face_info["bbox"]
        x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
        control_mask[y1:y2, x1:x2] = 255
        control_mask = Image.fromarray(control_mask.astype(np.uint8))

        controlnet_scales = {
            "pose": pose_strength,
            "canny": canny_strength
        }
        self.pipe.controlnet = MultiControlNetModel(
            [self.controlnet_identitynet]
            + [self.controlnet_map[s] for s in controlnet_selection]
        )
        control_scales = [float(identitynet_strength_ratio)] + [
            controlnet_scales[s] for s in controlnet_selection
        ]
        control_images = [face_kps] + [
            self.controlnet_map_fn[s](img_controlnet).resize((width, height))
            for s in controlnet_selection
        ]

        generator = torch.Generator(device=device).manual_seed(seed)

        print("Start inference...")
        print(f"[Debug] Prompt: {prompt}, \n[Debug] Neg Prompt: {negative_prompt}")
        print(f"[Debug] Number of Inference Steps: {num_steps}")

        self.pipe.set_ip_adapter_scale(adapter_strength_ratio)
        images = self.pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            image_embeds=face_emb,
            image=control_images,
            control_mask=control_mask,
            controlnet_conditioning_scale=control_scales,
            num_inference_steps=num_steps,
            guidance_scale=guidance_scale,
            height=height,
            width=width,
            generator=generator,
        ).images

        # Convert the output image to base64
        buffered = io.BytesIO()
        images[0].save(buffered, format="JPEG")
        img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")

        return {"generated_image_base64": img_str}