# coding=utf-8
# Copyright 2025 HuggingFace Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import io
import json
from typing import List, Literal, Optional, Union, cast

import requests

from .deprecation_utils import deprecate
from .import_utils import is_safetensors_available, is_torch_available


if is_torch_available():
    import torch

    from ..image_processor import VaeImageProcessor
    from ..video_processor import VideoProcessor

    if is_safetensors_available():
        import safetensors.torch

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


from PIL import Image


def detect_image_type(data: bytes) -> str:
    if data.startswith(b"\xff\xd8"):
        return "jpeg"
    elif data.startswith(b"\x89PNG\r\n\x1a\n"):
        return "png"
    elif data.startswith(b"GIF87a") or data.startswith(b"GIF89a"):
        return "gif"
    elif data.startswith(b"BM"):
        return "bmp"
    return "unknown"


def check_inputs_decode(
    endpoint: str,
    tensor: "torch.Tensor",
    processor: Optional[Union["VaeImageProcessor", "VideoProcessor"]] = None,
    do_scaling: bool = True,
    scaling_factor: Optional[float] = None,
    shift_factor: Optional[float] = None,
    output_type: Literal["mp4", "pil", "pt"] = "pil",
    return_type: Literal["mp4", "pil", "pt"] = "pil",
    image_format: Literal["png", "jpg"] = "jpg",
    partial_postprocess: bool = False,
    input_tensor_type: Literal["binary"] = "binary",
    output_tensor_type: Literal["binary"] = "binary",
    height: Optional[int] = None,
    width: Optional[int] = None,
):
    if tensor.ndim == 3 and height is None and width is None:
        raise ValueError("`height` and `width` required for packed latents.")
    if (
        output_type == "pt"
        and return_type == "pil"
        and not partial_postprocess
        and not isinstance(processor, (VaeImageProcessor, VideoProcessor))
    ):
        raise ValueError("`processor` is required.")
    if do_scaling and scaling_factor is None:
        deprecate(
            "do_scaling",
            "1.0.0",
            "`do_scaling` is deprecated, pass `scaling_factor` and `shift_factor` if required.",
            standard_warn=False,
        )


def postprocess_decode(
    response: requests.Response,
    processor: Optional[Union["VaeImageProcessor", "VideoProcessor"]] = None,
    output_type: Literal["mp4", "pil", "pt"] = "pil",
    return_type: Literal["mp4", "pil", "pt"] = "pil",
    partial_postprocess: bool = False,
):
    if output_type == "pt" or (output_type == "pil" and processor is not None):
        output_tensor = response.content
        parameters = response.headers
        shape = json.loads(parameters["shape"])
        dtype = parameters["dtype"]
        torch_dtype = DTYPE_MAP[dtype]
        output_tensor = torch.frombuffer(bytearray(output_tensor), dtype=torch_dtype).reshape(shape)
    if output_type == "pt":
        if partial_postprocess:
            if return_type == "pil":
                output = [Image.fromarray(image.numpy()) for image in output_tensor]
                if len(output) == 1:
                    output = output[0]
            elif return_type == "pt":
                output = output_tensor
        else:
            if processor is None or return_type == "pt":
                output = output_tensor
            else:
                if isinstance(processor, VideoProcessor):
                    output = cast(
                        List[Image.Image],
                        processor.postprocess_video(output_tensor, output_type="pil")[0],
                    )
                else:
                    output = cast(
                        Image.Image,
                        processor.postprocess(output_tensor, output_type="pil")[0],
                    )
    elif output_type == "pil" and return_type == "pil" and processor is None:
        output = Image.open(io.BytesIO(response.content)).convert("RGB")
        detected_format = detect_image_type(response.content)
        output.format = detected_format
    elif output_type == "pil" and processor is not None:
        if return_type == "pil":
            output = [
                Image.fromarray(image)
                for image in (output_tensor.permute(0, 2, 3, 1).float().numpy() * 255).round().astype("uint8")
            ]
        elif return_type == "pt":
            output = output_tensor
    elif output_type == "mp4" and return_type == "mp4":
        output = response.content
    return output


def prepare_decode(
    tensor: "torch.Tensor",
    processor: Optional[Union["VaeImageProcessor", "VideoProcessor"]] = None,
    do_scaling: bool = True,
    scaling_factor: Optional[float] = None,
    shift_factor: Optional[float] = None,
    output_type: Literal["mp4", "pil", "pt"] = "pil",
    image_format: Literal["png", "jpg"] = "jpg",
    partial_postprocess: bool = False,
    height: Optional[int] = None,
    width: Optional[int] = None,
):
    headers = {}
    parameters = {
        "image_format": image_format,
        "output_type": output_type,
        "partial_postprocess": partial_postprocess,
        "shape": list(tensor.shape),
        "dtype": str(tensor.dtype).split(".")[-1],
    }
    if do_scaling and scaling_factor is not None:
        parameters["scaling_factor"] = scaling_factor
    if do_scaling and shift_factor is not None:
        parameters["shift_factor"] = shift_factor
    if do_scaling and scaling_factor is None:
        parameters["do_scaling"] = do_scaling
    elif do_scaling and scaling_factor is None and shift_factor is None:
        parameters["do_scaling"] = do_scaling
    if height is not None and width is not None:
        parameters["height"] = height
        parameters["width"] = width
    headers["Content-Type"] = "tensor/binary"
    headers["Accept"] = "tensor/binary"
    if output_type == "pil" and image_format == "jpg" and processor is None:
        headers["Accept"] = "image/jpeg"
    elif output_type == "pil" and image_format == "png" and processor is None:
        headers["Accept"] = "image/png"
    elif output_type == "mp4":
        headers["Accept"] = "text/plain"
    tensor_data = safetensors.torch._tobytes(tensor, "tensor")
    return {"data": tensor_data, "params": parameters, "headers": headers}


def remote_decode(
    endpoint: str,
    tensor: "torch.Tensor",
    processor: Optional[Union["VaeImageProcessor", "VideoProcessor"]] = None,
    do_scaling: bool = True,
    scaling_factor: Optional[float] = None,
    shift_factor: Optional[float] = None,
    output_type: Literal["mp4", "pil", "pt"] = "pil",
    return_type: Literal["mp4", "pil", "pt"] = "pil",
    image_format: Literal["png", "jpg"] = "jpg",
    partial_postprocess: bool = False,
    input_tensor_type: Literal["binary"] = "binary",
    output_tensor_type: Literal["binary"] = "binary",
    height: Optional[int] = None,
    width: Optional[int] = None,
) -> Union[Image.Image, List[Image.Image], bytes, "torch.Tensor"]:
    """
    Hugging Face Hybrid Inference that allow running VAE decode remotely.

    Args:
        endpoint (`str`):
            Endpoint for Remote Decode.
        tensor (`torch.Tensor`):
            Tensor to be decoded.
        processor (`VaeImageProcessor` or `VideoProcessor`, *optional*):
            Used with `return_type="pt"`, and `return_type="pil"` for Video models.
        do_scaling (`bool`, default `True`, *optional*):
            **DEPRECATED**. **pass `scaling_factor`/`shift_factor` instead.** **still set
            do_scaling=None/do_scaling=False for no scaling until option is removed** When `True` scaling e.g. `latents
            / self.vae.config.scaling_factor` is applied remotely. If `False`, input must be passed with scaling
            applied.
        scaling_factor (`float`, *optional*):
            Scaling is applied when passed e.g. [`latents /
            self.vae.config.scaling_factor`](https://github.com/huggingface/diffusers/blob/7007febae5cff000d4df9059d9cf35133e8b2ca9/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py#L1083C37-L1083C77).
            - SD v1: 0.18215
            - SD XL: 0.13025
            - Flux: 0.3611
            If `None`, input must be passed with scaling applied.
        shift_factor (`float`, *optional*):
            Shift is applied when passed e.g. `latents + self.vae.config.shift_factor`.
            - Flux: 0.1159
            If `None`, input must be passed with scaling applied.
        output_type (`"mp4"` or `"pil"` or `"pt", default `"pil"):
            **Endpoint** output type. Subject to change. Report feedback on preferred type.

            `"mp4": Supported by video models. Endpoint returns `bytes` of video. `"pil"`: Supported by image and video
            models.
                Image models: Endpoint returns `bytes` of an image in `image_format`. Video models: Endpoint returns
                `torch.Tensor` with partial `postprocessing` applied.
                    Requires `processor` as a flag (any `None` value will work).
            `"pt"`: Support by image and video models. Endpoint returns `torch.Tensor`.
                With `partial_postprocess=True` the tensor is postprocessed `uint8` image tensor.

            Recommendations:
                `"pt"` with `partial_postprocess=True` is the smallest transfer for full quality. `"pt"` with
                `partial_postprocess=False` is the most compatible with third party code. `"pil"` with
                `image_format="jpg"` is the smallest transfer overall.

        return_type (`"mp4"` or `"pil"` or `"pt", default `"pil"):
            **Function** return type.

            `"mp4": Function returns `bytes` of video. `"pil"`: Function returns `PIL.Image.Image`.
                With `output_type="pil" no further processing is applied. With `output_type="pt" a `PIL.Image.Image` is
                created.
                    `partial_postprocess=False` `processor` is required. `partial_postprocess=True` `processor` is
                    **not** required.
            `"pt"`: Function returns `torch.Tensor`.
                `processor` is **not** required. `partial_postprocess=False` tensor is `float16` or `bfloat16`, without
                denormalization. `partial_postprocess=True` tensor is `uint8`, denormalized.

        image_format (`"png"` or `"jpg"`, default `jpg`):
            Used with `output_type="pil"`. Endpoint returns `jpg` or `png`.

        partial_postprocess (`bool`, default `False`):
            Used with `output_type="pt"`. `partial_postprocess=False` tensor is `float16` or `bfloat16`, without
            denormalization. `partial_postprocess=True` tensor is `uint8`, denormalized.

        input_tensor_type (`"binary"`, default `"binary"`):
            Tensor transfer type.

        output_tensor_type (`"binary"`, default `"binary"`):
            Tensor transfer type.

        height (`int`, **optional**):
            Required for `"packed"` latents.

        width (`int`, **optional**):
            Required for `"packed"` latents.

    Returns:
        output (`Image.Image` or `List[Image.Image]` or `bytes` or `torch.Tensor`).
    """
    if input_tensor_type == "base64":
        deprecate(
            "input_tensor_type='base64'",
            "1.0.0",
            "input_tensor_type='base64' is deprecated. Using `binary`.",
            standard_warn=False,
        )
        input_tensor_type = "binary"
    if output_tensor_type == "base64":
        deprecate(
            "output_tensor_type='base64'",
            "1.0.0",
            "output_tensor_type='base64' is deprecated. Using `binary`.",
            standard_warn=False,
        )
        output_tensor_type = "binary"
    check_inputs_decode(
        endpoint,
        tensor,
        processor,
        do_scaling,
        scaling_factor,
        shift_factor,
        output_type,
        return_type,
        image_format,
        partial_postprocess,
        input_tensor_type,
        output_tensor_type,
        height,
        width,
    )
    kwargs = prepare_decode(
        tensor=tensor,
        processor=processor,
        do_scaling=do_scaling,
        scaling_factor=scaling_factor,
        shift_factor=shift_factor,
        output_type=output_type,
        image_format=image_format,
        partial_postprocess=partial_postprocess,
        height=height,
        width=width,
    )
    response = requests.post(endpoint, **kwargs)
    if not response.ok:
        raise RuntimeError(response.json())
    output = postprocess_decode(
        response=response,
        processor=processor,
        output_type=output_type,
        return_type=return_type,
        partial_postprocess=partial_postprocess,
    )
    return output


def check_inputs_encode(
    endpoint: str,
    image: Union["torch.Tensor", Image.Image],
    scaling_factor: Optional[float] = None,
    shift_factor: Optional[float] = None,
):
    pass


def postprocess_encode(
    response: requests.Response,
):
    output_tensor = response.content
    parameters = response.headers
    shape = json.loads(parameters["shape"])
    dtype = parameters["dtype"]
    torch_dtype = DTYPE_MAP[dtype]
    output_tensor = torch.frombuffer(bytearray(output_tensor), dtype=torch_dtype).reshape(shape)
    return output_tensor


def prepare_encode(
    image: Union["torch.Tensor", Image.Image],
    scaling_factor: Optional[float] = None,
    shift_factor: Optional[float] = None,
):
    headers = {}
    parameters = {}
    if scaling_factor is not None:
        parameters["scaling_factor"] = scaling_factor
    if shift_factor is not None:
        parameters["shift_factor"] = shift_factor
    if isinstance(image, torch.Tensor):
        data = safetensors.torch._tobytes(image.contiguous(), "tensor")
        parameters["shape"] = list(image.shape)
        parameters["dtype"] = str(image.dtype).split(".")[-1]
    else:
        buffer = io.BytesIO()
        image.save(buffer, format="PNG")
        data = buffer.getvalue()
    return {"data": data, "params": parameters, "headers": headers}


def remote_encode(
    endpoint: str,
    image: Union["torch.Tensor", Image.Image],
    scaling_factor: Optional[float] = None,
    shift_factor: Optional[float] = None,
) -> "torch.Tensor":
    """
    Hugging Face Hybrid Inference that allow running VAE encode remotely.

    Args:
        endpoint (`str`):
            Endpoint for Remote Decode.
        image (`torch.Tensor` or `PIL.Image.Image`):
            Image to be encoded.
        scaling_factor (`float`, *optional*):
            Scaling is applied when passed e.g. [`latents * self.vae.config.scaling_factor`].
            - SD v1: 0.18215
            - SD XL: 0.13025
            - Flux: 0.3611
            If `None`, input must be passed with scaling applied.
        shift_factor (`float`, *optional*):
            Shift is applied when passed e.g. `latents - self.vae.config.shift_factor`.
            - Flux: 0.1159
            If `None`, input must be passed with scaling applied.

    Returns:
        output (`torch.Tensor`).
    """
    check_inputs_encode(
        endpoint,
        image,
        scaling_factor,
        shift_factor,
    )
    kwargs = prepare_encode(
        image=image,
        scaling_factor=scaling_factor,
        shift_factor=shift_factor,
    )
    response = requests.post(endpoint, **kwargs)
    if not response.ok:
        raise RuntimeError(response.json())
    output = postprocess_encode(
        response=response,
    )
    return output