import gc
import os
from abc import ABC, abstractmethod

import PIL.Image
import torch
from controlnet_aux import (
    CannyDetector,
    LineartDetector,
    MidasDetector,
    PidiNetDetector,
    ZoeDetector,
)
from diffusers import (
    AutoencoderKL,
    EulerAncestralDiscreteScheduler,
    StableDiffusionXLAdapterPipeline,
    T2IAdapter,
)

ADAPTER_NAMES = [
    "TencentARC/t2i-adapter-canny-sdxl-1.0",
    "TencentARC/t2i-adapter-sketch-sdxl-1.0",
    "TencentARC/t2i-adapter-lineart-sdxl-1.0",
    "TencentARC/t2i-adapter-depth-midas-sdxl-1.0",
    "TencentARC/t2i-adapter-depth-zoe-sdxl-1.0",
    "TencentARC/t2i-adapter-recolor-sdxl-1.0",
]


class Preprocessor(ABC):
    @abstractmethod
    def to(self, device: torch.device | str) -> "Preprocessor":
        pass

    @abstractmethod
    def __call__(self, image: PIL.Image.Image) -> PIL.Image.Image:
        pass


class CannyPreprocessor(Preprocessor):
    def __init__(self):
        self.model = CannyDetector()

    def to(self, device: torch.device | str) -> Preprocessor:
        return self

    def __call__(self, image: PIL.Image.Image) -> PIL.Image.Image:
        return self.model(image, detect_resolution=384, image_resolution=1024)


class LineartPreprocessor(Preprocessor):
    def __init__(self):
        self.model = LineartDetector.from_pretrained("lllyasviel/Annotators")

    def to(self, device: torch.device | str) -> Preprocessor:
        return self.model.to(device)

    def __call__(self, image: PIL.Image.Image) -> PIL.Image.Image:
        return self.model(image, detect_resolution=384, image_resolution=1024)


class MidasPreprocessor(Preprocessor):
    def __init__(self):
        self.model = MidasDetector.from_pretrained(
            "valhalla/t2iadapter-aux-models", filename="dpt_large_384.pt", model_type="dpt_large"
        )

    def to(self, device: torch.device | str) -> Preprocessor:
        return self.model.to(device)

    def __call__(self, image: PIL.Image.Image) -> PIL.Image.Image:
        return self.model(image, detect_resolution=512, image_resolution=1024)


class PidiNetPreprocessor(Preprocessor):
    def __init__(self):
        self.model = PidiNetDetector.from_pretrained("lllyasviel/Annotators")

    def to(self, device: torch.device | str) -> Preprocessor:
        return self.model.to(device)

    def __call__(self, image: PIL.Image.Image) -> PIL.Image.Image:
        return self.model(image, detect_resolution=512, image_resolution=1024, apply_filter=True)


class RecolorPreprocessor(Preprocessor):
    def to(self, device: torch.device | str) -> Preprocessor:
        return self

    def __call__(self, image: PIL.Image.Image) -> PIL.Image.Image:
        return image.convert("L").convert("RGB")


class ZoePreprocessor(Preprocessor):
    def __init__(self):
        self.model = ZoeDetector.from_pretrained(
            "valhalla/t2iadapter-aux-models", filename="zoed_nk.pth", model_type="zoedepth_nk"
        )

    def to(self, device: torch.device | str) -> Preprocessor:
        return self.model.to(device)

    def __call__(self, image: PIL.Image.Image) -> PIL.Image.Image:
        return self.model(image, gamma_corrected=True, image_resolution=1024)


PRELOAD_PREPROCESSORS_IN_GPU_MEMORY = os.getenv("PRELOAD_PREPROCESSORS_IN_GPU_MEMORY", "1") == "1"
PRELOAD_PREPROCESSORS_IN_CPU_MEMORY = os.getenv("PRELOAD_PREPROCESSORS_IN_CPU_MEMORY", "0") == "1"
if PRELOAD_PREPROCESSORS_IN_GPU_MEMORY:
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    preprocessors_gpu: dict[str, Preprocessor] = {
        "TencentARC/t2i-adapter-canny-sdxl-1.0": CannyPreprocessor().to(device),
        "TencentARC/t2i-adapter-sketch-sdxl-1.0": PidiNetPreprocessor().to(device),
        "TencentARC/t2i-adapter-lineart-sdxl-1.0": LineartPreprocessor().to(device),
        "TencentARC/t2i-adapter-depth-midas-sdxl-1.0": MidasPreprocessor().to(device),
        "TencentARC/t2i-adapter-depth-zoe-sdxl-1.0": ZoePreprocessor().to(device),
        "TencentARC/t2i-adapter-recolor-sdxl-1.0": RecolorPreprocessor().to(device),
    }

    def get_preprocessor(adapter_name: str) -> Preprocessor:
        return preprocessors_gpu[adapter_name]

elif PRELOAD_PREPROCESSORS_IN_CPU_MEMORY:
    preprocessors_cpu: dict[str, Preprocessor] = {
        "TencentARC/t2i-adapter-canny-sdxl-1.0": CannyPreprocessor(),
        "TencentARC/t2i-adapter-sketch-sdxl-1.0": PidiNetPreprocessor(),
        "TencentARC/t2i-adapter-lineart-sdxl-1.0": LineartPreprocessor(),
        "TencentARC/t2i-adapter-depth-midas-sdxl-1.0": MidasPreprocessor(),
        "TencentARC/t2i-adapter-depth-zoe-sdxl-1.0": ZoePreprocessor(),
        "TencentARC/t2i-adapter-recolor-sdxl-1.0": RecolorPreprocessor(),
    }

    def get_preprocessor(adapter_name: str) -> Preprocessor:
        return preprocessors_cpu[adapter_name]

else:

    def get_preprocessor(adapter_name: str) -> Preprocessor:
        if adapter_name == "TencentARC/t2i-adapter-canny-sdxl-1.0":
            return CannyPreprocessor()
        elif adapter_name == "TencentARC/t2i-adapter-sketch-sdxl-1.0":
            return PidiNetPreprocessor()
        elif adapter_name == "TencentARC/t2i-adapter-lineart-sdxl-1.0":
            return LineartPreprocessor()
        elif adapter_name == "TencentARC/t2i-adapter-depth-midas-sdxl-1.0":
            return MidasPreprocessor()
        elif adapter_name == "TencentARC/t2i-adapter-depth-zoe-sdxl-1.0":
            return ZoePreprocessor()
        elif adapter_name == "TencentARC/t2i-adapter-recolor-sdxl-1.0":
            return RecolorPreprocessor()
        else:
            raise ValueError(f"Adapter name must be one of {ADAPTER_NAMES}")

    def download_all_preprocessors():
        for adapter_name in ADAPTER_NAMES:
            get_preprocessor(adapter_name)
        gc.collect()

    download_all_preprocessors()


def download_all_adapters():
    for adapter_name in ADAPTER_NAMES:
        T2IAdapter.from_pretrained(
            adapter_name,
            torch_dtype=torch.float16,
            varient="fp16",
        )
    gc.collect()


download_all_adapters()


class Model:
    MAX_NUM_INFERENCE_STEPS = 50

    def __init__(self, adapter_name: str):
        if adapter_name not in ADAPTER_NAMES:
            raise ValueError(f"Adapter name must be one of {ADAPTER_NAMES}")

        self.preprocessor_name = adapter_name
        self.adapter_name = adapter_name

        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        if torch.cuda.is_available():
            self.preprocessor = get_preprocessor(adapter_name).to(self.device)

            model_id = "stabilityai/stable-diffusion-xl-base-1.0"
            adapter = T2IAdapter.from_pretrained(
                adapter_name,
                torch_dtype=torch.float16,
                varient="fp16",
            ).to(self.device)
            self.pipe = StableDiffusionXLAdapterPipeline.from_pretrained(
                model_id,
                vae=AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16),
                adapter=adapter,
                scheduler=EulerAncestralDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler"),
                torch_dtype=torch.float16,
                variant="fp16",
            ).to(self.device)
            self.pipe.enable_xformers_memory_efficient_attention()
        else:
            self.preprocessor = None  # type: ignore
            self.pipe = None

    def change_preprocessor(self, adapter_name: str) -> None:
        if adapter_name not in ADAPTER_NAMES:
            raise ValueError(f"Adapter name must be one of {ADAPTER_NAMES}")
        if adapter_name == self.preprocessor_name:
            return

        if PRELOAD_PREPROCESSORS_IN_GPU_MEMORY:
            pass
        elif PRELOAD_PREPROCESSORS_IN_CPU_MEMORY:
            self.preprocessor.to("cpu")
        else:
            del self.preprocessor
        self.preprocessor = get_preprocessor(adapter_name).to(self.device)
        self.preprocessor_name = adapter_name
        gc.collect()
        torch.cuda.empty_cache()

    def change_adapter(self, adapter_name: str) -> None:
        if adapter_name not in ADAPTER_NAMES:
            raise ValueError(f"Adapter name must be one of {ADAPTER_NAMES}")
        if adapter_name == self.adapter_name:
            return
        self.pipe.adapter = T2IAdapter.from_pretrained(
            adapter_name,
            torch_dtype=torch.float16,
            varient="fp16",
        ).to(self.device)
        self.adapter_name = adapter_name
        gc.collect()
        torch.cuda.empty_cache()

    def resize_image(self, image: PIL.Image.Image) -> PIL.Image.Image:
        w, h = image.size
        scale = 1024 / max(w, h)
        new_w = int(w * scale)
        new_h = int(h * scale)
        return image.resize((new_w, new_h), PIL.Image.LANCZOS)

    def run(
        self,
        image: PIL.Image.Image,
        prompt: str,
        negative_prompt: str,
        adapter_name: str,
        num_inference_steps: int = 30,
        guidance_scale: float = 5.0,
        adapter_conditioning_scale: float = 1.0,
        cond_tau: float = 1.0,
        seed: int = 0,
        apply_preprocess: bool = True,
    ) -> list[PIL.Image.Image]:
        if not torch.cuda.is_available():
            raise RuntimeError("This demo does not work on CPU.")
        if num_inference_steps > self.MAX_NUM_INFERENCE_STEPS:
            raise ValueError(f"Number of steps must be less than {self.MAX_NUM_INFERENCE_STEPS}")

        # Resize image to avoid OOM
        image = self.resize_image(image)

        self.change_preprocessor(adapter_name)
        self.change_adapter(adapter_name)

        if apply_preprocess:
            image = self.preprocessor(image)

        generator = torch.Generator(device=self.device).manual_seed(seed)
        out = self.pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            image=image,
            num_inference_steps=num_inference_steps,
            adapter_conditioning_scale=adapter_conditioning_scale,
            cond_tau=cond_tau,
            generator=generator,
            guidance_scale=guidance_scale,
        ).images[0]
        return [image, out]