from typing import Dict, List, Any
import torch
from base64 import b64decode
from diffusers import AutoencoderKL
from diffusers.image_processor import VaeImageProcessor


class EndpointHandler:
    def __init__(self, path=""):
        self.device = "cuda"
        self.dtype = torch.bfloat16
        self.vae = (
            AutoencoderKL.from_pretrained(path, torch_dtype=self.dtype)
            .to(self.device, self.dtype)
            .eval()
        )

        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)

    @staticmethod
    def _unpack_latents(latents, height, width, vae_scale_factor):
        batch_size, num_patches, channels = latents.shape

        # VAE applies 8x compression on images but we must also account for packing which requires
        # latent height and width to be divisible by 2.
        height = 2 * (int(height) // (vae_scale_factor * 2))
        width = 2 * (int(width) // (vae_scale_factor * 2))

        latents = latents.view(batch_size, height // 2, width // 2, channels // 4, 2, 2)
        latents = latents.permute(0, 3, 1, 4, 2, 5)

        latents = latents.reshape(batch_size, channels // (2 * 2), height, width)

        return latents

    @torch.no_grad()
    def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
        """
        Args:
            data (:obj:):
                includes the input data and the parameters for the inference.
        """
        tensor = data["inputs"]
        tensor = b64decode(tensor.encode("utf-8"))
        parameters = data.get("parameters", {})
        if "shape" not in parameters:
            raise ValueError("Expected `shape` in parameters.")
        if "dtype" not in parameters:
            raise ValueError("Expected `dtype` in parameters.")
        if "height" not in parameters:
            raise ValueError("Expected `height` in parameters.")
        if "width" not in parameters:
            raise ValueError("Expected `width` in parameters.")

        DTYPE_MAP = {
            "float16": torch.float16,
            "float32": torch.float32,
            "bfloat16": torch.bfloat16,
        }

        shape = parameters.get("shape")
        dtype = DTYPE_MAP.get(parameters.get("dtype"))
        height = parameters.get("height")
        width = parameters.get("width")

        tensor = torch.frombuffer(bytearray(tensor), dtype=dtype).reshape(shape)

        tensor = tensor.to(self.device, self.dtype)

        tensor = self._unpack_latents(tensor, height, width, self.vae_scale_factor)
        tensor = (
            tensor / self.vae.config.scaling_factor
        ) + self.vae.config.shift_factor

        with torch.no_grad():
            image = self.vae.decode(tensor, return_dict=False)[0]

        image = self.image_processor.postprocess(image, output_type="pil")

        return image[0]