# 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 gc
import glob
import inspect
from functools import wraps

import pytest
import torch
from accelerate.utils.modeling import compute_module_sizes

from diffusers.utils.torch_utils import get_torch_cuda_device_capability

from ...testing_utils import (
    _check_safetensors_serialization,
    assert_tensors_close,
    backend_empty_cache,
    backend_max_memory_allocated,
    backend_reset_peak_memory_stats,
    backend_synchronize,
    is_cpu_offload,
    is_group_offload,
    is_memory,
    require_accelerator,
    torch_device,
)
from .common import BaseModelOutputMixin, cast_inputs_to_dtype, check_device_map_is_respected


def reseed_generator_input(model, inputs_dict):
    """Re-seed a `generator` entry in `inputs_dict`, in place, if the model's forward takes one.

    A `torch.Generator` is stateful, so a model that draws noise in its forward (e.g. a diffusion decoder)
    returns something different on every call. Tests below reuse one `inputs_dict` across several forwards
    and compare the outputs, which only holds if each pass starts from the same generator state.

    The signature is read off the class, not the instance: offloading hooks replace `model.forward` with a
    `(*args, **kwargs)` wrapper, so an instance lookup finds no `generator` on exactly the models these
    tests offload.
    """
    if "generator" in inspect.signature(type(model).forward).parameters:
        inputs_dict["generator"] = torch.manual_seed(0)


def require_offload_support(func):
    """
    Decorator to skip tests if model doesn't support offloading (requires _no_split_modules).
    """

    @wraps(func)
    def wrapper(self, *args, **kwargs):
        if self.model_class._no_split_modules is None:
            pytest.skip("Test not supported for this model as `_no_split_modules` is not set.")
        return func(self, *args, **kwargs)

    return wrapper


def require_group_offload_support(func):
    """
    Decorator to skip tests if model doesn't support group offloading.
    """

    @wraps(func)
    def wrapper(self, *args, **kwargs):
        if not self.model_class._supports_group_offloading:
            pytest.skip("Model does not support group offloading.")
        return func(self, *args, **kwargs)

    return wrapper


@is_cpu_offload
class CPUOffloadTesterMixin(BaseModelOutputMixin):
    """
    Mixin class for testing CPU offloading functionality.

    Expected from config mixin:
        - model_class: The model class to test

    Optional properties:
        - model_split_percents: List of percentages for splitting model across devices (default: [0.5, 0.7])

    Expected methods from config mixin:
        - get_init_dict(): Returns dict of arguments to initialize the model
        - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass

    Pytest mark: cpu_offload
        Use `pytest -m "not cpu_offload"` to skip these tests
    """

    @property
    def model_split_percents(self) -> list[float]:
        """List of percentages for splitting model across devices during offloading tests."""
        return [0.5, 0.7]

    @require_offload_support
    @torch.no_grad()
    def test_cpu_offload(self, base_model_output, tmp_path, atol=1e-5, rtol=0):
        torch.manual_seed(0)
        config = self.get_init_dict()
        inputs_dict = self.get_dummy_inputs()
        model = self.model_class(**config).eval()

        model = model.to(torch_device)

        model_size = compute_module_sizes(model)[""]
        # We test several splits of sizes to make sure it works
        max_gpu_sizes = [int(p * model_size) for p in self.model_split_percents]
        model.cpu().save_pretrained(str(tmp_path))

        for max_size in max_gpu_sizes:
            max_memory = {0: max_size, "cpu": model_size * 2}
            new_model = self.model_class.from_pretrained(str(tmp_path), device_map="auto", max_memory=max_memory)
            # Making sure part of the model will actually end up offloaded
            assert set(new_model.hf_device_map.values()) == {0, "cpu"}, "Model should be split between GPU and CPU"

            check_device_map_is_respected(new_model, new_model.hf_device_map)
            torch.manual_seed(0)
            new_output = new_model(**inputs_dict)

            assert_tensors_close(
                base_model_output, new_output[0], atol=atol, rtol=rtol, msg="Output should match with CPU offloading"
            )

    @require_offload_support
    @torch.no_grad()
    def test_disk_offload_without_safetensors(self, base_model_output, tmp_path, atol=1e-5, rtol=0):
        torch.manual_seed(0)
        config = self.get_init_dict()
        inputs_dict = self.get_dummy_inputs()
        model = self.model_class(**config).eval()

        model = model.to(torch_device)

        model_size = compute_module_sizes(model)[""]
        max_size = int(self.model_split_percents[0] * model_size)
        # Force disk offload by setting very small CPU memory
        max_memory = {0: max_size, "cpu": int(0.1 * max_size)}

        model.cpu().save_pretrained(str(tmp_path), safe_serialization=False)
        # This errors out because it's missing an offload folder
        with pytest.raises(ValueError):
            new_model = self.model_class.from_pretrained(str(tmp_path), device_map="auto", max_memory=max_memory)

        new_model = self.model_class.from_pretrained(
            str(tmp_path), device_map="auto", max_memory=max_memory, offload_folder=str(tmp_path)
        )

        check_device_map_is_respected(new_model, new_model.hf_device_map)
        torch.manual_seed(0)
        new_output = new_model(**inputs_dict)

        assert_tensors_close(
            base_model_output, new_output[0], atol=atol, rtol=rtol, msg="Output should match with disk offloading"
        )

    @require_offload_support
    @torch.no_grad()
    def test_disk_offload_with_safetensors(self, base_model_output, tmp_path, atol=1e-5, rtol=0):
        torch.manual_seed(0)
        config = self.get_init_dict()
        inputs_dict = self.get_dummy_inputs()
        model = self.model_class(**config).eval()

        model = model.to(torch_device)

        model_size = compute_module_sizes(model)[""]
        model.cpu().save_pretrained(str(tmp_path))

        max_size = int(self.model_split_percents[0] * model_size)
        max_memory = {0: max_size, "cpu": max_size}
        new_model = self.model_class.from_pretrained(
            str(tmp_path), device_map="auto", offload_folder=str(tmp_path), max_memory=max_memory
        )

        check_device_map_is_respected(new_model, new_model.hf_device_map)
        torch.manual_seed(0)
        new_output = new_model(**inputs_dict)

        assert_tensors_close(
            base_model_output,
            new_output[0],
            atol=atol,
            rtol=rtol,
            msg="Output should match with disk offloading (safetensors)",
        )


@is_group_offload
class GroupOffloadTesterMixin(BaseModelOutputMixin):
    """
    Mixin class for testing group offloading functionality.

    Expected from config mixin:
        - model_class: The model class to test

    Expected methods from config mixin:
        - get_init_dict(): Returns dict of arguments to initialize the model
        - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass

    Pytest mark: group_offload
        Use `pytest -m "not group_offload"` to skip these tests
    """

    @require_group_offload_support
    @pytest.mark.parametrize("record_stream", [False, True])
    def test_group_offloading(self, base_model_output, record_stream, atol=1e-5, rtol=0):
        init_dict = self.get_init_dict()
        inputs_dict = self.get_dummy_inputs()

        @torch.no_grad()
        def run_forward(model):
            assert all(
                module._diffusers_hook.get_hook("group_offloading") is not None
                for module in model.modules()
                if hasattr(module, "_diffusers_hook")
            ), "Group offloading hook should be set"
            model.eval()
            reseed_generator_input(model, inputs_dict)
            return model(**inputs_dict)[0]

        output_without_group_offloading = base_model_output

        torch.manual_seed(0)
        model = self.model_class(**init_dict)
        model.enable_group_offload(torch_device, offload_type="block_level", num_blocks_per_group=1)
        output_with_group_offloading1 = run_forward(model)

        torch.manual_seed(0)
        model = self.model_class(**init_dict)
        model.enable_group_offload(torch_device, offload_type="block_level", num_blocks_per_group=1, non_blocking=True)
        output_with_group_offloading2 = run_forward(model)

        torch.manual_seed(0)
        model = self.model_class(**init_dict)
        model.enable_group_offload(torch_device, offload_type="leaf_level")
        output_with_group_offloading3 = run_forward(model)

        torch.manual_seed(0)
        model = self.model_class(**init_dict)
        model.enable_group_offload(
            torch_device, offload_type="leaf_level", use_stream=True, record_stream=record_stream
        )
        output_with_group_offloading4 = run_forward(model)

        assert_tensors_close(
            output_without_group_offloading,
            output_with_group_offloading1,
            atol=atol,
            rtol=rtol,
            msg="Output should match with block-level offloading",
        )
        assert_tensors_close(
            output_without_group_offloading,
            output_with_group_offloading2,
            atol=atol,
            rtol=rtol,
            msg="Output should match with non-blocking block-level offloading",
        )
        assert_tensors_close(
            output_without_group_offloading,
            output_with_group_offloading3,
            atol=atol,
            rtol=rtol,
            msg="Output should match with leaf-level offloading",
        )
        assert_tensors_close(
            output_without_group_offloading,
            output_with_group_offloading4,
            atol=atol,
            rtol=rtol,
            msg="Output should match with leaf-level offloading with stream",
        )

    @require_group_offload_support
    @pytest.mark.parametrize("record_stream", [False, True])
    @pytest.mark.parametrize("offload_type", ["block_level", "leaf_level"])
    @torch.no_grad()
    def test_group_offloading_with_layerwise_casting(self, record_stream, offload_type):
        torch.manual_seed(0)
        init_dict = self.get_init_dict()
        inputs_dict = self.get_dummy_inputs()
        model = self.model_class(**init_dict)

        model.to(torch_device)
        model.eval()
        _ = model(**inputs_dict)[0]

        torch.manual_seed(0)
        init_dict = self.get_init_dict()
        inputs_dict = self.get_dummy_inputs()
        storage_dtype, compute_dtype = torch.float16, torch.float32
        inputs_dict = cast_inputs_to_dtype(inputs_dict, torch.float32, compute_dtype)
        model = self.model_class(**init_dict)
        model.eval()
        additional_kwargs = {} if offload_type == "leaf_level" else {"num_blocks_per_group": 1}
        model.enable_group_offload(
            torch_device, offload_type=offload_type, use_stream=True, record_stream=record_stream, **additional_kwargs
        )
        model.enable_layerwise_casting(storage_dtype=storage_dtype, compute_dtype=compute_dtype)
        _ = model(**inputs_dict)[0]

    @require_group_offload_support
    @pytest.mark.parametrize("record_stream", [False, True])
    @pytest.mark.parametrize("offload_type", ["block_level", "leaf_level"])
    @torch.no_grad()
    @torch.inference_mode()
    def test_group_offloading_with_disk(self, tmp_path, record_stream, offload_type, atol=1e-5, rtol=0):
        def _run_forward(model, inputs_dict):
            reseed_generator_input(model, inputs_dict)
            torch.manual_seed(0)
            return model(**inputs_dict)[0]

        init_dict = self.get_init_dict()
        inputs_dict = self.get_dummy_inputs()
        torch.manual_seed(0)
        model = self.model_class(**init_dict)

        model.eval()
        model.to(torch_device)
        output_without_group_offloading = _run_forward(model, inputs_dict)

        torch.manual_seed(0)
        model = self.model_class(**init_dict)
        model.eval()

        num_blocks_per_group = None if offload_type == "leaf_level" else 1
        additional_kwargs = {} if offload_type == "leaf_level" else {"num_blocks_per_group": num_blocks_per_group}
        tmpdir = str(tmp_path)
        model.enable_group_offload(
            torch_device,
            offload_type=offload_type,
            offload_to_disk_path=tmpdir,
            use_stream=True,
            record_stream=record_stream,
            **additional_kwargs,
        )
        has_safetensors = glob.glob(f"{tmpdir}/*.safetensors")
        assert has_safetensors, "No safetensors found in the directory."

        # For "leaf-level", there is a prefetching hook which makes this check a bit non-deterministic
        # in nature. So, skip it.
        if offload_type != "leaf_level":
            is_correct, extra_files, missing_files = _check_safetensors_serialization(
                module=model,
                offload_to_disk_path=tmpdir,
                offload_type=offload_type,
                num_blocks_per_group=num_blocks_per_group,
                block_modules=model._group_offload_block_modules
                if hasattr(model, "_group_offload_block_modules")
                else None,
            )
            if not is_correct:
                if extra_files:
                    raise ValueError(f"Found extra files: {', '.join(extra_files)}")
                elif missing_files:
                    raise ValueError(f"Following files are missing: {', '.join(missing_files)}")

        output_with_group_offloading = _run_forward(model, inputs_dict)
        assert_tensors_close(
            output_without_group_offloading,
            output_with_group_offloading,
            atol=atol,
            rtol=rtol,
            msg="Output should match with disk-based group offloading",
        )


class LayerwiseCastingTesterMixin:
    """
    Mixin class for testing layerwise dtype casting for memory optimization.

    Expected from config mixin:
        - model_class: The model class to test

    Expected methods from config mixin:
        - get_init_dict(): Returns dict of arguments to initialize the model
        - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass
    """

    @pytest.mark.xfail(
        condition=torch_device == "mps",
        reason="MPS does not support float8 casting.",
        strict=True,
    )
    @torch.no_grad()
    def test_layerwise_casting_memory(self):
        MB_TOLERANCE = 0.2
        LEAST_COMPUTE_CAPABILITY = 8.0

        def reset_memory_stats():
            gc.collect()
            backend_synchronize(torch_device)
            backend_empty_cache(torch_device)
            backend_reset_peak_memory_stats(torch_device)

        def get_memory_usage(storage_dtype, compute_dtype):
            torch.manual_seed(0)
            config = self.get_init_dict()
            inputs_dict = self.get_dummy_inputs()
            inputs_dict = cast_inputs_to_dtype(inputs_dict, torch.float32, compute_dtype)
            model = self.model_class(**config).eval()
            model = model.to(torch_device, dtype=compute_dtype)
            model.enable_layerwise_casting(storage_dtype=storage_dtype, compute_dtype=compute_dtype)

            reset_memory_stats()
            model(**inputs_dict)
            model_memory_footprint = model.get_memory_footprint()
            peak_inference_memory_allocated_mb = backend_max_memory_allocated(torch_device) / 1024**2

            return model_memory_footprint, peak_inference_memory_allocated_mb

        fp32_memory_footprint, fp32_max_memory = get_memory_usage(torch.float32, torch.float32)
        fp8_e4m3_fp32_memory_footprint, fp8_e4m3_fp32_max_memory = get_memory_usage(torch.float8_e4m3fn, torch.float32)
        fp8_e4m3_bf16_memory_footprint, fp8_e4m3_bf16_max_memory = get_memory_usage(
            torch.float8_e4m3fn, torch.bfloat16
        )

        compute_capability = get_torch_cuda_device_capability() if torch_device == "cuda" else None
        assert fp8_e4m3_bf16_memory_footprint < fp8_e4m3_fp32_memory_footprint < fp32_memory_footprint, (
            "Memory footprint should decrease with lower precision storage"
        )

        # NOTE: the following assertion would fail on our CI (running Tesla T4) due to bf16 using more memory than fp32.
        # On other devices, such as DGX (Ampere) and Audace (Ada), the test passes. So, we conditionally check it.
        if compute_capability and compute_capability >= LEAST_COMPUTE_CAPABILITY:
            assert fp8_e4m3_bf16_max_memory < fp8_e4m3_fp32_max_memory, (
                "Peak memory should be lower with bf16 compute on newer GPUs"
            )

        # On this dummy test case with a small model, sometimes fp8_e4m3_fp32 max memory usage is higher than fp32 by a few
        # bytes. This only happens for some models, so we allow a small tolerance.
        # For any real model being tested, the order would be fp8_e4m3_bf16 < fp8_e4m3_fp32 < fp32.
        assert (
            fp8_e4m3_fp32_max_memory < fp32_max_memory
            or abs(fp8_e4m3_fp32_max_memory - fp32_max_memory) < MB_TOLERANCE
        ), "Peak memory should be lower or within tolerance with fp8 storage"

    @pytest.mark.xfail(
        condition=torch_device == "mps",
        reason="MPS does not support float8 casting.",
        strict=True,
    )
    def test_layerwise_casting_training(self):
        def test_fn(storage_dtype, compute_dtype):
            if torch.device(torch_device).type == "cpu" and compute_dtype == torch.bfloat16:
                pytest.skip("Skipping test because CPU doesn't go well with bfloat16.")

            model = self.model_class(**self.get_init_dict())
            model = model.to(torch_device, dtype=compute_dtype)
            model.enable_layerwise_casting(storage_dtype=storage_dtype, compute_dtype=compute_dtype)
            model.train()

            inputs_dict = self.get_dummy_inputs()
            inputs_dict = cast_inputs_to_dtype(inputs_dict, torch.float32, compute_dtype)
            with torch.amp.autocast(device_type=torch.device(torch_device).type):
                output = model(**inputs_dict, return_dict=False)[0]

                input_tensor = inputs_dict[self.main_input_name]
                noise = torch.randn((input_tensor.shape[0],) + self.output_shape).to(torch_device)
                noise = cast_inputs_to_dtype(noise, torch.float32, compute_dtype)
                loss = torch.nn.functional.mse_loss(output, noise)

            loss.backward()

        test_fn(torch.float16, torch.float32)
        test_fn(torch.float8_e4m3fn, torch.float32)
        test_fn(torch.float8_e5m2, torch.float32)
        test_fn(torch.float8_e4m3fn, torch.bfloat16)


@is_memory
@require_accelerator
class MemoryTesterMixin(CPUOffloadTesterMixin, GroupOffloadTesterMixin, LayerwiseCastingTesterMixin):
    """
    Combined mixin class for all memory optimization tests including CPU/disk offloading,
    group offloading, and layerwise dtype casting.

    This mixin inherits from:
        - CPUOffloadTesterMixin: CPU and disk offloading tests
        - GroupOffloadTesterMixin: Group offloading tests (block-level and leaf-level)
        - LayerwiseCastingTesterMixin: Layerwise dtype casting tests

    Expected from config mixin:
        - model_class: The model class to test

    Optional properties:
        - model_split_percents: List of percentages for splitting model across devices (default: [0.5, 0.7])

    Expected methods from config mixin:
        - get_init_dict(): Returns dict of arguments to initialize the model
        - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass

    Pytest mark: memory
        Use `pytest -m "not memory"` to skip these tests
    """