# 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 inspect

import pytest
import torch

from diffusers.hooks import HookRegistry, ModelHook
from diffusers.training_utils import free_memory
from diffusers.utils.logging import get_logger

from ..testing_utils import CaptureLogger, torch_device


logger = get_logger(__name__)  # pylint: disable=invalid-name


class DummyBlock(torch.nn.Module):
    def __init__(self, in_features: int, hidden_features: int, out_features: int) -> None:
        super().__init__()

        self.proj_in = torch.nn.Linear(in_features, hidden_features)
        self.activation = torch.nn.ReLU()
        self.proj_out = torch.nn.Linear(hidden_features, out_features)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.proj_in(x)
        x = self.activation(x)
        x = self.proj_out(x)
        return x


class DummyModel(torch.nn.Module):
    def __init__(self, in_features: int, hidden_features: int, out_features: int, num_layers: int) -> None:
        super().__init__()

        self.linear_1 = torch.nn.Linear(in_features, hidden_features)
        self.activation = torch.nn.ReLU()
        self.blocks = torch.nn.ModuleList(
            [DummyBlock(hidden_features, hidden_features, hidden_features) for _ in range(num_layers)]
        )
        self.linear_2 = torch.nn.Linear(hidden_features, out_features)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.linear_1(x)
        x = self.activation(x)
        for block in self.blocks:
            x = block(x)
        x = self.linear_2(x)
        return x


class SignatureModel(torch.nn.Module):
    def forward(self, hidden_states, timestep, encoder_hidden_states=None):
        return hidden_states + timestep.sum()


class AddHook(ModelHook):
    def __init__(self, value: int):
        super().__init__()
        self.value = value

    def pre_forward(self, module: torch.nn.Module, *args, **kwargs):
        logger.debug("AddHook pre_forward")
        args = ((x + self.value) if torch.is_tensor(x) else x for x in args)
        return args, kwargs

    def post_forward(self, module, output):
        logger.debug("AddHook post_forward")
        return output


class MultiplyHook(ModelHook):
    def __init__(self, value: int):
        super().__init__()
        self.value = value

    def pre_forward(self, module, *args, **kwargs):
        logger.debug("MultiplyHook pre_forward")
        args = ((x * self.value) if torch.is_tensor(x) else x for x in args)
        return args, kwargs

    def post_forward(self, module, output):
        logger.debug("MultiplyHook post_forward")
        return output

    def __repr__(self):
        return f"MultiplyHook(value={self.value})"


class StatefulAddHook(ModelHook):
    _is_stateful = True

    def __init__(self, value: int):
        super().__init__()
        self.value = value
        self.increment = 0

    def pre_forward(self, module, *args, **kwargs):
        logger.debug("StatefulAddHook pre_forward")
        add_value = self.value + self.increment
        self.increment += 1
        args = ((x + add_value) if torch.is_tensor(x) else x for x in args)
        return args, kwargs

    def reset_state(self, module):
        self.increment = 0


class SkipLayerHook(ModelHook):
    def __init__(self, skip_layer: bool):
        super().__init__()
        self.skip_layer = skip_layer

    def pre_forward(self, module, *args, **kwargs):
        logger.debug("SkipLayerHook pre_forward")
        return args, kwargs

    def new_forward(self, module, *args, **kwargs):
        logger.debug("SkipLayerHook new_forward")
        if self.skip_layer:
            return args[0]
        return self.fn_ref.original_forward(*args, **kwargs)

    def post_forward(self, module, output):
        logger.debug("SkipLayerHook post_forward")
        return output


class TestHooks:
    in_features = 4
    hidden_features = 8
    out_features = 4
    num_layers = 2

    def setup_method(self):
        params = self.get_module_parameters()
        self.model = DummyModel(**params)
        self.model.to(torch_device)

    def teardown_method(self):
        del self.model
        gc.collect()
        free_memory()

    def get_module_parameters(self):
        return {
            "in_features": self.in_features,
            "hidden_features": self.hidden_features,
            "out_features": self.out_features,
            "num_layers": self.num_layers,
        }

    def get_generator(self):
        return torch.Generator(device=torch_device).manual_seed(0)

    def test_hook_registry(self):
        registry = HookRegistry.check_if_exists_or_initialize(self.model)
        registry.register_hook(AddHook(1), "add_hook")
        registry.register_hook(MultiplyHook(2), "multiply_hook")

        registry_repr = repr(registry)
        expected_repr = "HookRegistry(\n  (0) add_hook - AddHook\n  (1) multiply_hook - MultiplyHook(value=2)\n)"

        assert len(registry.hooks) == 2
        assert registry._hook_order == ["add_hook", "multiply_hook"]
        assert registry_repr == expected_repr

        registry.remove_hook("add_hook")

        assert len(registry.hooks) == 1
        assert registry._hook_order == ["multiply_hook"]

    def test_stateful_hook(self):
        registry = HookRegistry.check_if_exists_or_initialize(self.model)
        registry.register_hook(StatefulAddHook(1), "stateful_add_hook")

        assert registry.hooks["stateful_add_hook"].increment == 0

        input = torch.randn(1, 4, device=torch_device, generator=self.get_generator())
        num_repeats = 3

        for i in range(num_repeats):
            result = self.model(input)
            if i == 0:
                output1 = result

        assert registry.get_hook("stateful_add_hook").increment == num_repeats

        registry.reset_stateful_hooks()
        output2 = self.model(input)

        assert registry.get_hook("stateful_add_hook").increment == 1
        assert torch.allclose(output1, output2)

    def test_child_registries_cache_invalidation(self):
        # Unit-level part of the regression test for
        # https://github.com/huggingface/diffusers/issues/14037: the parent registry caches its
        # child registries, so a hook registered on a child block after the cache was built is
        # invisible to the parent until the cache is invalidated.
        parent = HookRegistry.check_if_exists_or_initialize(self.model)

        # Build the parent's child-registry cache while no block carries a hook yet.
        assert parent._get_child_registries() == []

        # Register a hook on a child block. The parent's cached (empty) list is now stale.
        block = self.model.blocks[0]
        child = HookRegistry.check_if_exists_or_initialize(block)
        child.register_hook(AddHook(1), "add_hook")
        assert parent._get_child_registries() == []  # still stale before invalidation

        # Invalidating across the tree makes the new child registry reachable from the parent.
        parent.invalidate_child_registries_cache()
        assert child in parent._get_child_registries()

    def test_inference(self):
        registry = HookRegistry.check_if_exists_or_initialize(self.model)
        registry.register_hook(AddHook(1), "add_hook")
        registry.register_hook(MultiplyHook(2), "multiply_hook")

        input = torch.randn(1, 4, device=torch_device, generator=self.get_generator())
        output1 = self.model(input).mean().detach().cpu().item()

        registry.remove_hook("multiply_hook")
        new_input = input * 2
        output2 = self.model(new_input).mean().detach().cpu().item()

        registry.remove_hook("add_hook")
        new_input = input * 2 + 1
        output3 = self.model(new_input).mean().detach().cpu().item()

        assert output1 == pytest.approx(output2, abs=5e-6)
        assert output1 == pytest.approx(output3, abs=5e-6)
        assert output2 == pytest.approx(output3, abs=5e-6)

    def test_skip_layer_hook(self):
        registry = HookRegistry.check_if_exists_or_initialize(self.model)
        registry.register_hook(SkipLayerHook(skip_layer=True), "skip_layer_hook")

        input = torch.zeros(1, 4, device=torch_device)
        output = self.model(input).mean().detach().cpu().item()
        assert output == 0.0

        registry.remove_hook("skip_layer_hook")
        registry.register_hook(SkipLayerHook(skip_layer=False), "skip_layer_hook")
        output = self.model(input).mean().detach().cpu().item()
        assert output != 0.0

    def test_skip_layer_internal_block(self):
        registry = HookRegistry.check_if_exists_or_initialize(self.model.linear_1)
        input = torch.zeros(1, 4, device=torch_device)

        registry.register_hook(SkipLayerHook(skip_layer=True), "skip_layer_hook")
        with pytest.raises(RuntimeError, match="mat1 and mat2 shapes cannot be multiplied"):
            self.model(input).mean().detach().cpu().item()

        registry.remove_hook("skip_layer_hook")
        output = self.model(input).mean().detach().cpu().item()
        assert output != 0.0

        registry = HookRegistry.check_if_exists_or_initialize(self.model.blocks[1])
        registry.register_hook(SkipLayerHook(skip_layer=True), "skip_layer_hook")
        output = self.model(input).mean().detach().cpu().item()
        assert output != 0.0

    def test_invocation_order_stateful_first(self):
        registry = HookRegistry.check_if_exists_or_initialize(self.model)
        registry.register_hook(StatefulAddHook(1), "add_hook")
        registry.register_hook(AddHook(2), "add_hook_2")
        registry.register_hook(MultiplyHook(3), "multiply_hook")

        input = torch.randn(1, 4, device=torch_device, generator=self.get_generator())

        logger = get_logger(__name__)
        logger.setLevel("DEBUG")

        with CaptureLogger(logger) as cap_logger:
            self.model(input)
        output = cap_logger.out.replace(" ", "").replace("\n", "")
        expected_invocation_order_log = (
            (
                "MultiplyHook pre_forward\n"
                "AddHook pre_forward\n"
                "StatefulAddHook pre_forward\n"
                "AddHook post_forward\n"
                "MultiplyHook post_forward\n"
            )
            .replace(" ", "")
            .replace("\n", "")
        )
        assert output == expected_invocation_order_log

        registry.remove_hook("add_hook")
        with CaptureLogger(logger) as cap_logger:
            self.model(input)
        output = cap_logger.out.replace(" ", "").replace("\n", "")
        expected_invocation_order_log = (
            ("MultiplyHook pre_forward\nAddHook pre_forward\nAddHook post_forward\nMultiplyHook post_forward\n")
            .replace(" ", "")
            .replace("\n", "")
        )
        assert output == expected_invocation_order_log

    def test_invocation_order_stateful_middle(self):
        registry = HookRegistry.check_if_exists_or_initialize(self.model)
        registry.register_hook(AddHook(2), "add_hook")
        registry.register_hook(StatefulAddHook(1), "add_hook_2")
        registry.register_hook(MultiplyHook(3), "multiply_hook")

        input = torch.randn(1, 4, device=torch_device, generator=self.get_generator())

        logger = get_logger(__name__)
        logger.setLevel("DEBUG")

        with CaptureLogger(logger) as cap_logger:
            self.model(input)
        output = cap_logger.out.replace(" ", "").replace("\n", "")
        expected_invocation_order_log = (
            (
                "MultiplyHook pre_forward\n"
                "StatefulAddHook pre_forward\n"
                "AddHook pre_forward\n"
                "AddHook post_forward\n"
                "MultiplyHook post_forward\n"
            )
            .replace(" ", "")
            .replace("\n", "")
        )
        assert output == expected_invocation_order_log

        registry.remove_hook("add_hook")
        with CaptureLogger(logger) as cap_logger:
            self.model(input)
        output = cap_logger.out.replace(" ", "").replace("\n", "")
        expected_invocation_order_log = (
            ("MultiplyHook pre_forward\nStatefulAddHook pre_forward\nMultiplyHook post_forward\n")
            .replace(" ", "")
            .replace("\n", "")
        )
        assert output == expected_invocation_order_log

        registry.remove_hook("add_hook_2")
        with CaptureLogger(logger) as cap_logger:
            self.model(input)
        output = cap_logger.out.replace(" ", "").replace("\n", "")
        expected_invocation_order_log = (
            ("MultiplyHook pre_forward\nMultiplyHook post_forward\n").replace(" ", "").replace("\n", "")
        )
        assert output == expected_invocation_order_log

    def test_invocation_order_stateful_last(self):
        registry = HookRegistry.check_if_exists_or_initialize(self.model)
        registry.register_hook(AddHook(1), "add_hook")
        registry.register_hook(MultiplyHook(2), "multiply_hook")
        registry.register_hook(StatefulAddHook(3), "add_hook_2")

        input = torch.randn(1, 4, device=torch_device, generator=self.get_generator())

        logger = get_logger(__name__)
        logger.setLevel("DEBUG")

        with CaptureLogger(logger) as cap_logger:
            self.model(input)
        output = cap_logger.out.replace(" ", "").replace("\n", "")
        expected_invocation_order_log = (
            (
                "StatefulAddHook pre_forward\n"
                "MultiplyHook pre_forward\n"
                "AddHook pre_forward\n"
                "AddHook post_forward\n"
                "MultiplyHook post_forward\n"
            )
            .replace(" ", "")
            .replace("\n", "")
        )
        assert output == expected_invocation_order_log

        registry.remove_hook("add_hook")
        with CaptureLogger(logger) as cap_logger:
            self.model(input)
        output = cap_logger.out.replace(" ", "").replace("\n", "")
        expected_invocation_order_log = (
            ("StatefulAddHook pre_forward\nMultiplyHook pre_forward\nMultiplyHook post_forward\n")
            .replace(" ", "")
            .replace("\n", "")
        )
        assert output == expected_invocation_order_log

    def test_register_hook_preserves_forward_signature(self):
        model = SignatureModel().to(torch_device)
        sig_before = inspect.signature(model.forward)
        assert list(sig_before.parameters) == ["hidden_states", "timestep", "encoder_hidden_states"]

        registry = HookRegistry.check_if_exists_or_initialize(model)
        registry.register_hook(ModelHook(), "noop")

        assert inspect.signature(model.forward) == sig_before
        assert model.forward.__name__ == "forward"

        registry.register_hook(ModelHook(), "noop_2")
        assert inspect.signature(model.forward) == sig_before
        assert model.forward.__name__ == "forward"

        registry.remove_hook("noop_2")
        assert inspect.signature(model.forward) == sig_before
        registry.remove_hook("noop")
        assert inspect.signature(model.forward) == sig_before

    def test_register_hook_preserves_forward_signature_torch_export(self):
        model = SignatureModel().to(torch_device)
        registry = HookRegistry.check_if_exists_or_initialize(model)
        registry.register_hook(ModelHook(), "noop")
        registry.register_hook(ModelHook(), "noop_2")

        hidden_states = torch.randn(1, 4, 8, device=torch_device)
        timestep = torch.tensor([1.0], device=torch_device)

        eager_output = model(hidden_states, timestep)
        assert tuple(eager_output.shape) == (1, 4, 8)

        exported = torch.export.export(model, args=(), kwargs={"hidden_states": hidden_states, "timestep": timestep})
        assert exported is not None

        exported_output = exported.module()(hidden_states=hidden_states, timestep=timestep)
        assert torch.allclose(exported_output, eager_output)