"""Chained block hooks / attention overrides, on a REAL H3 forward."""

import pytest
import torch


@pytest.fixture(autouse=True)
def _needs_comfy(comfyui):
    """Imports comfy."""


def _clone(tiny):
    return tiny.patcher.clone()


def test_two_block_hooks_on_one_block_both_run_and_the_output_is_the_block(tiny_h3):
    from core import dit_hooks
    base = tiny_h3.run()
    p = _clone(tiny_h3)
    seen = []

    def hook(name):
        def h(args, extra):
            seen.append(name)
            return extra["original_block"](args)
        return h

    dit_hooks.add_block_hook(p, "funui.a", 1, hook("a"))
    dit_hooks.add_block_hook(p, "funui.b", 1, hook("b"))
    out = tiny_h3.run(p)
    assert seen == ["b", "a"], "the second hook replaced the first instead of chaining"
    assert torch.allclose(out[0], base[0]) and torch.allclose(out[1], base[1])


def test_two_attention_overrides_both_see_every_call(tiny_h3):
    from core import dit_hooks
    p = _clone(tiny_h3)
    calls = {"a": 0, "b": 0}

    def ov(name):
        def o(func, q, k, v, *a, **kw):
            calls[name] += 1
            return func(q, k, v, *a, **kw)
        return o

    dit_hooks.add_attention_override(p, "funui.a", ov("a"))
    dit_hooks.add_attention_override(p, "funui.b", ov("b"))
    tiny_h3.run(p)
    # 4 DiT blocks + 1 token-refiner block
    assert calls == {"a": 5, "b": 5}


def test_strip_removes_ours_and_puts_back_a_foreign_hook_underneath(tiny_h3):
    from core import dit_hooks, patching
    p = _clone(tiny_h3)
    foreign = []

    def theirs(args, extra):
        foreign.append(1)
        return extra["original_block"](args)

    dit = p.model_options["transformer_options"].setdefault("patches_replace", {}).setdefault("dit", {})
    dit[("double_block", 0)] = theirs
    dit_hooks.add_block_hook(p, "funui.x", 0, lambda a, e: e["original_block"](a))
    dit_hooks.add_attention_override(p, "funui.x", lambda f, *a, **k: f(*a, **k))

    assert patching.strip(p, "funui.") == 2
    to = p.model_options["transformer_options"]
    assert to["patches_replace"]["dit"][("double_block", 0)] is theirs
    assert "optimized_attention_override" not in to
    tiny_h3.run(p)
    assert foreign == [1]


def test_a_failing_block_hook_is_dropped_and_the_run_matches_an_unhooked_one(tiny_h3):
    from core import dit_hooks, patching
    base = tiny_h3.run()
    p = _clone(tiny_h3)
    dropped = patching.Dropped()
    guarded = patching.GuardedPatcher(p, "funui.bad", dropped)

    def bad(args, extra):
        raise RuntimeError("boom")

    def bad_attn(func, *a, **k):
        raise RuntimeError("boom")

    dit_hooks.add_block_hook(guarded, "funui.bad", 2, bad)
    dit_hooks.add_attention_override(guarded, "funui.bad", bad_attn)
    out = tiny_h3.run(p)
    assert "funui.bad" in dropped
    assert torch.allclose(out[0], base[0])


@pytest.mark.parametrize("spec,expect,problems", [
    ("1", [1], 0), ("0-2", [0, 1, 2], 0), ("3,1", [1, 3], 0),
    ("2-9", [2, 3], 1), ("x", [], 1), ("", [], 0), ("-1", [], 1),
])
def test_parse_blocks_reports_what_it_ignores(spec, expect, problems):
    from core import dit_hooks
    got, said = dit_hooks.parse_blocks(spec, 4)
    assert got == expect and len(said) == problems


def test_without_compiler_switches_it_off_for_the_sample_and_back(tiny_h3):
    import comfy.cli_args
    from comfy.patcher_extension import WrappersMP
    from core import dit_hooks
    p = tiny_h3.patcher.clone()
    dit_hooks.without_compiler(p, "funui.x")
    (outer,), = p.wrappers[WrappersMP.OUTER_SAMPLE].values()
    had = hasattr(comfy.cli_args.args, "disable_comfy_compiler")
    prior = getattr(comfy.cli_args.args, "disable_comfy_compiler", None)
    seen = outer(lambda: comfy.cli_args.args.disable_comfy_compiler)
    assert seen is True
    assert hasattr(comfy.cli_args.args, "disable_comfy_compiler") == had
    assert getattr(comfy.cli_args.args, "disable_comfy_compiler", None) == prior


@pytest.mark.parametrize("module,settings", [
    ("block_repeat", {"enabled": True, "blocks": "1"}),
    ("shadow_negative", {"enabled": True}),
])
def test_block_replacing_modules_turn_the_compiler_off(tiny_h3, module, settings):
    from comfy.patcher_extension import WrappersMP
    from modules.sampling.modifiers.nodes import FunUILoadModifiers
    patched, _ = FunUILoadModifiers.execute(tiny_h3.patcher, {module: settings}).result
    assert patched.wrappers.get(WrappersMP.OUTER_SAMPLE), f"{module} left the compiler on"


def test_a_huge_range_is_clamped_before_it_is_built_and_still_reported():
    import time
    from core import dit_hooks
    t = time.time()
    got, problems = dit_hooks.parse_blocks("1-99999999999", 48)
    assert got == list(range(1, 48)) and len(problems) == 1 and "past" in problems[0]
    assert time.time() - t < 1.0