from types import SimpleNamespace

import torch

from h3_sparse_attention import H3AccelerationConfig, install_h3_acceleration


class _Attention:
    def forward(self, x):
        return x


class _Block:
    def __init__(self):
        self.attn = _Attention()

    def forward(self, x):
        return self.attn.forward(x)


def test_compile_default_installs_and_restores_per_block(monkeypatch):
    blocks = [_Block(), _Block()]
    pipe = SimpleNamespace(transformer=SimpleNamespace(transformer_blocks=blocks))
    compile_calls = []
    disable_calls = []

    def fake_compile(fn, **kwargs):
        compile_calls.append(kwargs)
        return lambda *args, **kw: fn(*args, **kw)

    def fake_disable(fn):
        disable_calls.append(fn)
        return fn

    monkeypatch.setattr(torch, "compile", fake_compile)
    monkeypatch.setattr(torch.compiler, "disable", fake_disable)
    monkeypatch.setattr(torch._inductor, "list_mode_options", lambda mode: {})

    config = H3AccelerationConfig()
    assert config.torch_compile is True
    assert config.compile_reproducibility_fix is True
    assert config.vae_fp16 is False
    with install_h3_acceleration(pipe, config) as plugin:
        assert plugin.config is config
        assert len(compile_calls) == len(disable_calls) == 2
        assert all(call["options"]["deterministic"] for call in compile_calls)
        assert all("forward" in block.__dict__ for block in blocks)
        assert all("forward" in block.attn.__dict__ for block in blocks)
    assert all("forward" not in block.__dict__ for block in blocks)
    assert all("forward" not in block.attn.__dict__ for block in blocks)


def test_reproducibility_fix_can_restore_legacy_whole_block_compile(monkeypatch):
    blocks = [_Block(), _Block()]
    pipe = SimpleNamespace(transformer=SimpleNamespace(transformer_blocks=blocks))
    compile_calls = []

    def fake_compile(fn, **kwargs):
        compile_calls.append(kwargs)
        return lambda *args, **kw: fn(*args, **kw)

    monkeypatch.setattr(torch, "compile", fake_compile)
    monkeypatch.setattr(
        torch.compiler,
        "disable",
        lambda fn: (_ for _ in ()).throw(
            AssertionError("legacy mode must not install the eager attention boundary")
        ),
    )
    monkeypatch.setattr(
        torch._inductor,
        "list_mode_options",
        lambda mode: (_ for _ in ()).throw(
            AssertionError("legacy mode must not install deterministic options")
        ),
    )

    config = H3AccelerationConfig(compile_reproducibility_fix=False)
    with install_h3_acceleration(pipe, config):
        assert all("forward" in block.__dict__ for block in blocks)
        assert all("forward" not in block.attn.__dict__ for block in blocks)
    assert compile_calls == [
        {"mode": None, "dynamic": None},
        {"mode": None, "dynamic": None},
    ]


def test_explicit_no_compile_is_not_a_noop_fallback(monkeypatch):
    block = _Block()
    pipe = SimpleNamespace(transformer=SimpleNamespace(transformer_blocks=[block]))
    monkeypatch.setattr(torch, "compile", lambda *args, **kwargs: (_ for _ in ()).throw(
        AssertionError("compile should not be called")
    ))
    with install_h3_acceleration(pipe, H3AccelerationConfig(torch_compile=False)):
        assert "forward" not in block.__dict__