from typing import Any, Dict, List, Tuple

import base64
import torch

from diffusers import AutoencoderKLHunyuanVideo
from safetensors.torch import _tobytes


def prepare_tensor(tensor: torch.Tensor) -> Tuple[str, List[int], str]:
    tensor_data = base64.b64encode(_tobytes(tensor, "inputs")).decode("utf-8")
    shape = list(tensor.shape)
    dtype = str(tensor.dtype).split(".")[-1]
    return tensor_data, shape, dtype


def unpack_tensor(tensor_data: str, shape: List[int], dtype: str) -> torch.Tensor:
    tensor = base64.b64decode(tensor_data.encode("utf-8"))
    DTYPE_MAP = {
        "float16": torch.float16,
        "float32": torch.float32,
        "bfloat16": torch.bfloat16,
    }
    torch_dtype = DTYPE_MAP.get(dtype)
    tensor = torch.frombuffer(bytearray(tensor), dtype=torch_dtype).reshape(shape)
    return tensor


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

    @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 = data["inputs"]
        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.")

        shape = parameters.get("shape")
        dtype = parameters.get("dtype")

        tensor = unpack_tensor(tensor_data, shape, dtype)

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

        tensor = tensor / self.vae.config.scaling_factor

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

        tensor_data, shape, dtype = prepare_tensor(frames)

        return {"tensor": tensor_data, "shape": shape, "dtype": dtype}