import math
import sys
import types
from pathlib import Path

import torch


sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from nodes import (  # noqa: E402
    ANIMA_COEFFICIENTS,
    _STEP_INFO_KEY,
    AnimaTeaCache,
    AnimaTeaCacheRuntime,
    CacheDecision,
    poly1d,
)


def step_options(index, total=5, labels=None, uuids=None):
    options = {
        _STEP_INFO_KEY: {
            "index": index,
            "total": total,
            "percent": index / total,
            "schedule_key": (total,),
        }
    }
    if labels is not None:
        options["cond_or_uncond"] = labels
    if uuids is not None:
        options["uuids"] = uuids
    return options


def runtime(threshold=0.05, start=0.0, end=1.0):
    return AnimaTeaCacheRuntime(threshold, start, end, torch.device("cpu"))


def seed_residual(cache_runtime, decision, tensor):
    cache_runtime.commit_residual(decision, torch.full_like(tensor, 3.0))


def test_poly1d_is_quartic_with_five_coefficients():
    value = 0.125
    expected = sum(
        coefficient * value ** (len(ANIMA_COEFFICIENTS) - index - 1)
        for index, coefficient in enumerate(ANIMA_COEFFICIENTS)
    )
    assert math.isclose(poly1d(ANIMA_COEFFICIENTS, value), expected)


def test_first_and_last_steps_force_full():
    cache_runtime = runtime(threshold=100.0)
    tensor = torch.ones(1, 1, 1, 1, 2)

    first = cache_runtime.decide(tensor, step_options(0))
    assert first.should_calc
    seed_residual(cache_runtime, first, tensor)

    middle = cache_runtime.decide(tensor * 1.001, step_options(1))
    assert not middle.should_calc

    last = cache_runtime.decide(tensor * 1.002, step_options(4))
    assert last.should_calc


def test_threshold_refreshes_residual():
    cache_runtime = runtime(threshold=0.01)
    tensor = torch.ones(1, 1, 1, 1, 2)
    first = cache_runtime.decide(tensor, step_options(0))
    seed_residual(cache_runtime, first, tensor)

    decision = cache_runtime.decide(tensor * 2, step_options(1))
    assert decision.should_calc
    assert cache_runtime.entries[0].accumulated_distance == 0.0


def test_cfg_entries_are_separate_and_reused_in_batch_order():
    cache_runtime = runtime(threshold=100.0)
    tensor = torch.stack((torch.ones(2), torch.ones(2) * 2)).reshape(2, 1, 1, 1, 2)
    options = step_options(0, labels=[1, 0], uuids=["negative", "positive"])
    first = cache_runtime.decide(tensor, options)
    assert first.keys == [(1, "negative", 0), (0, "positive", 0)]
    cache_runtime.commit_residual(
        first,
        torch.stack((torch.ones(2) * 3, torch.ones(2) * 4)).reshape(2, 1, 1, 1, 2),
    )

    options = step_options(1, labels=[1, 0], uuids=["negative", "positive"])
    second = cache_runtime.decide(tensor * 1.001, options)
    assert not second.should_calc
    reused = cache_runtime.reuse_residual(second, torch.zeros_like(tensor))
    assert torch.equal(reused[0], torch.full_like(reused[0], 3.0))
    assert torch.equal(reused[1], torch.full_like(reused[1], 4.0))


def test_reentering_step_zero_resets_interrupted_run():
    cache_runtime = runtime(threshold=100.0)
    tensor = torch.ones(1, 1, 1, 1, 2)
    first = cache_runtime.decide(tensor, step_options(0))
    seed_residual(cache_runtime, first, tensor)
    cache_runtime.decide(tensor * 1.001, step_options(1))

    restarted = cache_runtime.decide(tensor, step_options(0))
    assert restarted.should_calc
    assert cache_runtime.entries[0].previous_residual is None


def test_new_schedule_and_shape_change_reset_cache():
    cache_runtime = runtime(threshold=100.0)
    tensor = torch.ones(1, 1, 1, 1, 2)
    first = cache_runtime.decide(tensor, step_options(0))
    seed_residual(cache_runtime, first, tensor)

    new_schedule = step_options(1)
    new_schedule[_STEP_INFO_KEY]["schedule_key"] = (99,)
    assert cache_runtime.decide(tensor, new_schedule).should_calc

    cache_runtime.commit_residual(
        CacheDecision(True, True, [0], [slice(0, 1)]),
        tensor,
    )
    larger = torch.ones(1, 1, 1, 2, 2)
    assert cache_runtime.decide(larger, step_options(2)).should_calc


def test_outside_percent_range_forces_full():
    cache_runtime = runtime(threshold=100.0, start=0.25, end=0.75)
    tensor = torch.ones(1, 1, 1, 1, 2)
    first = cache_runtime.decide(tensor, step_options(0))
    seed_residual(cache_runtime, first, tensor)

    assert not cache_runtime.decide(tensor * 1.001, step_options(2)).should_calc
    assert cache_runtime.decide(tensor * 1.002, step_options(4)).should_calc


def test_missing_step_metadata_disables_cache():
    cache_runtime = runtime()
    decision = cache_runtime.decide(torch.ones(1, 1, 1, 1, 2), {})
    assert decision.should_calc
    assert not decision.cacheable
    assert cache_runtime.entries == {}


def test_build_step_info_matches_current_sigma():
    cache_runtime = runtime()
    sigmas = torch.tensor([1.0, 0.7, 0.4, 0.1, 0.0])
    info = cache_runtime._build_step_info(torch.tensor([0.69]), sigmas)
    assert info["index"] == 1
    assert info["total"] == 4
    assert math.isclose(info["percent"], 1 / 4)


def test_first_block_signal_matches_norm1_formula():
    class Block:
        use_adaln_lora = True
        layer_norm_self_attn = torch.nn.Identity()

        @staticmethod
        def adaln_modulation_self_attn(embedding):
            return torch.cat((embedding, embedding * 2, embedding * 3), dim=-1)

    class Model:
        blocks = [Block()]

    cache_runtime = runtime()
    hidden = torch.ones(1, 1, 1, 1, 2)
    embedding = torch.tensor([[[1.0, 2.0]]])
    adaln_lora = torch.ones(1, 1, 6)
    extra = torch.full_like(hidden, 2.0)

    result = cache_runtime._first_block_modulated_input(
        Model(), hidden, embedding, adaln_lora, extra
    )
    shift = embedding + 1
    scale = embedding * 2 + 1
    expected = (hidden + extra) * (1 + scale[:, :, None, None, :]) + shift[
        :, :, None, None, :
    ]
    assert torch.equal(result, expected)


def test_diffusion_wrapper_preserves_later_wrappers_and_replaces_original():
    cache_runtime = runtime()
    calls = []
    cache_runtime.forward = lambda model, *args, **kwargs: calls.append(model) or "teacache"

    class Executor:
        class_obj = object()
        original = lambda: "original"

        def __call__(self, *args, **kwargs):
            return self.original(*args, **kwargs)

    executor = Executor()
    result = cache_runtime.diffusion_model_wrapper(
        executor,
        transformer_options={
            "sigmas": torch.tensor([1.0]),
            "sample_sigmas": torch.tensor([1.0, 0.0]),
        },
    )
    assert result == "teacache"
    assert calls == [executor.class_obj]


def test_node_registers_diffusion_wrapper_in_model_options():
    class FakeAnima:
        pass

    class FakeModelPatcher:
        def __init__(self):
            self.diffusion_model = FakeAnima()
            self.model_options = {}

        def get_model_object(self, _name):
            return self.diffusion_model

        def clone(self):
            clone = FakeModelPatcher()
            clone.diffusion_model = self.diffusion_model
            clone.model_options = {
                key: value.copy() if isinstance(value, dict) else value
                for key, value in self.model_options.items()
            }
            return clone

    def add_wrapper_with_key(wrapper_type, key, wrapper, options, is_model_options=False):
        assert is_model_options
        transformer_options = options.setdefault("transformer_options", {})
        transformer_options.setdefault("wrappers", {}).setdefault(wrapper_type, {}).setdefault(
            key, []
        ).append(wrapper)

    modules = {
        "comfy": types.ModuleType("comfy"),
        "comfy.ldm": types.ModuleType("comfy.ldm"),
        "comfy.ldm.anima": types.ModuleType("comfy.ldm.anima"),
        "comfy.ldm.anima.model": types.ModuleType("comfy.ldm.anima.model"),
        "comfy.model_management": types.ModuleType("comfy.model_management"),
        "comfy.patcher_extension": types.ModuleType("comfy.patcher_extension"),
    }
    modules["comfy.ldm.anima.model"].Anima = FakeAnima
    modules["comfy.model_management"].get_torch_device = lambda: torch.device("cpu")
    modules["comfy.patcher_extension"].WrappersMP = types.SimpleNamespace(
        DIFFUSION_MODEL="diffusion_model"
    )
    modules["comfy.patcher_extension"].add_wrapper_with_key = add_wrapper_with_key
    modules["comfy"].model_management = modules["comfy.model_management"]
    modules["comfy"].patcher_extension = modules["comfy.patcher_extension"]
    previous = {name: sys.modules.get(name) for name in modules}
    sys.modules.update(modules)
    try:
        patched = AnimaTeaCache().apply_teacache(
            FakeModelPatcher(), 0.05, 0.0, 1.0, "cpu"
        )[0]
        wrappers = patched.model_options["transformer_options"]["wrappers"]
        assert len(wrappers["diffusion_model"]["anima_teacache"]) == 1
    finally:
        for name, module in previous.items():
            if module is None:
                sys.modules.pop(name, None)
            else:
                sys.modules[name] = module


def test_non_anima_model_is_rejected():
    class FakeModelPatcher:
        @staticmethod
        def get_model_object(_name):
            return object()

    class FakeAnima:
        pass

    comfy_module = types.ModuleType("comfy")
    ldm_module = types.ModuleType("comfy.ldm")
    anima_package = types.ModuleType("comfy.ldm.anima")
    anima_model_module = types.ModuleType("comfy.ldm.anima.model")
    anima_model_module.Anima = FakeAnima
    previous = {
        name: sys.modules.get(name)
        for name in ("comfy", "comfy.ldm", "comfy.ldm.anima", "comfy.ldm.anima.model")
    }
    sys.modules.update(
        {
            "comfy": comfy_module,
            "comfy.ldm": ldm_module,
            "comfy.ldm.anima": anima_package,
            "comfy.ldm.anima.model": anima_model_module,
        }
    )
    try:
        try:
            AnimaTeaCache().apply_teacache(FakeModelPatcher(), 0.05, 0.0, 1.0, "cpu")
        except TypeError as error:
            assert "only supports" in str(error)
        else:
            raise AssertionError("Expected non-Anima model rejection")
    finally:
        for name, module in previous.items():
            if module is None:
                sys.modules.pop(name, None)
            else:
                sys.modules[name] = module