# coding=utf-8
# Copyright 2026 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 PIL.Image
import pytest
import torch

from diffusers.utils import load_image
from diffusers.utils.constants import (
    DECODE_ENDPOINT_FLUX,
    DECODE_ENDPOINT_SD_V1,
    DECODE_ENDPOINT_SD_XL,
    ENCODE_ENDPOINT_FLUX,
    ENCODE_ENDPOINT_SD_V1,
    ENCODE_ENDPOINT_SD_XL,
)
from diffusers.utils.remote_utils import (
    remote_decode,
    remote_encode,
)

from ..testing_utils import (
    enable_full_determinism,
    slow,
)


enable_full_determinism()

IMAGE = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg?download=true"

# The `ENCODE_ENDPOINT_*` Inference Endpoints these tests call have been undeployed and now answer every
# request with `404 NOT_FOUND`, which `remote_encode` surfaces as a `RuntimeError`. Nothing in the repo can
# bring them back, so the tests are marked `xfail` rather than deleted: if the endpoints are ever redeployed
# the suite reports XPASS and the marker can be dropped.
ENDPOINTS_GONE = pytest.mark.xfail(
    reason="ENCODE_ENDPOINT_* are no longer deployed and return 404 NOT_FOUND.",
    raises=RuntimeError,
)


class RemoteAutoencoderKLEncodeMixin:
    channels: int = None
    endpoint: str = None
    decode_endpoint: str = None
    dtype: torch.dtype = None
    scaling_factor: float = None
    shift_factor: float = None
    image: PIL.Image.Image = None

    def get_dummy_inputs(self):
        if self.image is None:
            self.image = load_image(IMAGE)
        inputs = {
            "endpoint": self.endpoint,
            "image": self.image,
            "scaling_factor": self.scaling_factor,
            "shift_factor": self.shift_factor,
        }
        return inputs

    @ENDPOINTS_GONE
    def test_image_input(self):
        inputs = self.get_dummy_inputs()
        height, width = inputs["image"].height, inputs["image"].width
        output = remote_encode(**inputs)
        assert list(output.shape) == [1, self.channels, height // 8, width // 8]
        decoded = remote_decode(
            tensor=output,
            endpoint=self.decode_endpoint,
            scaling_factor=self.scaling_factor,
            shift_factor=self.shift_factor,
            image_format="png",
        )
        assert decoded.height == height
        assert decoded.width == width
        # image_slice = torch.from_numpy(np.array(inputs["image"])[0, -3:, -3:].flatten())
        # decoded_slice = torch.from_numpy(np.array(decoded)[0, -3:, -3:].flatten())
        # TODO: how to test this? encode->decode is lossy. expected slice of encoded latent?


# The tests below hit live HF Inference Endpoints, which are not part of the fast CI contract, so they are
# gated behind `RUN_SLOW` instead of running on every push.
@slow
class TestRemoteAutoencoderKLSDv1(RemoteAutoencoderKLEncodeMixin):
    channels = 4
    endpoint = ENCODE_ENDPOINT_SD_V1
    decode_endpoint = DECODE_ENDPOINT_SD_V1
    dtype = torch.float16
    scaling_factor = 0.18215
    shift_factor = None


@slow
class TestRemoteAutoencoderKLSDXL(RemoteAutoencoderKLEncodeMixin):
    channels = 4
    endpoint = ENCODE_ENDPOINT_SD_XL
    decode_endpoint = DECODE_ENDPOINT_SD_XL
    dtype = torch.float16
    scaling_factor = 0.13025
    shift_factor = None


@slow
class TestRemoteAutoencoderKLFlux(RemoteAutoencoderKLEncodeMixin):
    channels = 16
    endpoint = ENCODE_ENDPOINT_FLUX
    decode_endpoint = DECODE_ENDPOINT_FLUX
    dtype = torch.bfloat16
    scaling_factor = 0.3611
    shift_factor = 0.1159


class RemoteAutoencoderKLEncodeSlowTestMixin:
    channels: int = 4
    endpoint: str = None
    decode_endpoint: str = None
    dtype: torch.dtype = None
    scaling_factor: float = None
    shift_factor: float = None
    image: PIL.Image.Image = None

    def get_dummy_inputs(self):
        if self.image is None:
            self.image = load_image(IMAGE)
        inputs = {
            "endpoint": self.endpoint,
            "image": self.image,
            "scaling_factor": self.scaling_factor,
            "shift_factor": self.shift_factor,
        }
        return inputs

    @ENDPOINTS_GONE
    def test_multi_res(self, tmp_path):
        inputs = self.get_dummy_inputs()
        for height in {
            320,
            512,
            640,
            704,
            896,
            1024,
            1208,
            1384,
            1536,
            1608,
            1864,
            2048,
        }:
            for width in {
                320,
                512,
                640,
                704,
                896,
                1024,
                1208,
                1384,
                1536,
                1608,
                1864,
                2048,
            }:
                inputs["image"] = inputs["image"].resize(
                    (
                        width,
                        height,
                    )
                )
                output = remote_encode(**inputs)
                assert list(output.shape) == [1, self.channels, height // 8, width // 8]
                decoded = remote_decode(
                    tensor=output,
                    endpoint=self.decode_endpoint,
                    scaling_factor=self.scaling_factor,
                    shift_factor=self.shift_factor,
                    image_format="png",
                )
                assert decoded.height == height
                assert decoded.width == width
                decoded.save(tmp_path / f"test_multi_res_{height}_{width}.png")


@slow
class TestRemoteAutoencoderKLSDv1Slow(RemoteAutoencoderKLEncodeSlowTestMixin):
    endpoint = ENCODE_ENDPOINT_SD_V1
    decode_endpoint = DECODE_ENDPOINT_SD_V1
    dtype = torch.float16
    scaling_factor = 0.18215
    shift_factor = None


@slow
class TestRemoteAutoencoderKLSDXLSlow(RemoteAutoencoderKLEncodeSlowTestMixin):
    endpoint = ENCODE_ENDPOINT_SD_XL
    decode_endpoint = DECODE_ENDPOINT_SD_XL
    dtype = torch.float16
    scaling_factor = 0.13025
    shift_factor = None


@slow
class TestRemoteAutoencoderKLFluxSlow(RemoteAutoencoderKLEncodeSlowTestMixin):
    channels = 16
    endpoint = ENCODE_ENDPOINT_FLUX
    decode_endpoint = DECODE_ENDPOINT_FLUX
    dtype = torch.bfloat16
    scaling_factor = 0.3611
    shift_factor = 0.1159