"""Unit tests for inteprelens.core.hooks using the shared llama fixture."""
import pytest
import torch

from inteprelens.core.hooks import (
    autograd_node_hook,
    capture_backward_pre,
    capture_gradients,
    capture_kwargs_hook,
    tensor_grad_hook,
)

# ---------------------------------------------------------------------------
# helpers
# ---------------------------------------------------------------------------

def _simple_input(shared_model, llama_tokenizer, device="cuda"):
    """Return inputs_embeds with requires_grad=True so backward flows through frozen model."""
    enc = llama_tokenizer("Paris is the capital of France", return_tensors="pt")
    input_ids = enc["input_ids"].to(device)
    with torch.no_grad():
        embeds = shared_model.get_input_embeddings()(input_ids).detach().requires_grad_(True)
    return {"inputs_embeds": embeds}


# ---------------------------------------------------------------------------
# 1. capture_gradients
# ---------------------------------------------------------------------------

def test_capture_gradients_output(shared_model, llama_adapter, llama_tokenizer):
    inputs = _simple_input(shared_model, llama_tokenizer, shared_model.device)
    target_id = llama_tokenizer(" Paris", add_special_tokens=False)["input_ids"][0]

    attn = llama_adapter.get_attention_module(layer=0)
    with capture_gradients(attn, site="output") as grads:
        out = shared_model(**inputs)
        loss = -out.logits[0, -1, target_id]
        loss.backward()

    assert len(grads) == 1
    grad_t = grads[0][0]
    assert grad_t is not None
    assert grad_t.shape[-1] == shared_model.config.hidden_size


def test_capture_gradients_input(shared_model, llama_adapter, llama_tokenizer):
    inputs = _simple_input(shared_model, llama_tokenizer, shared_model.device)
    target_id = llama_tokenizer(" France", add_special_tokens=False)["input_ids"][0]

    attn = llama_adapter.get_attention_module(layer=0)
    with capture_gradients(attn, site="input") as grads:
        out = shared_model(**inputs)
        loss = -out.logits[0, -1, target_id]
        loss.backward()

    assert len(grads) == 1


def test_capture_gradients_removed_after_context(shared_model, llama_adapter):
    attn = llama_adapter.get_attention_module(layer=0)
    before = len(attn._backward_hooks)
    with capture_gradients(attn) as _:
        pass
    assert len(attn._backward_hooks) == before


# ---------------------------------------------------------------------------
# 2. capture_kwargs_hook
# ---------------------------------------------------------------------------

def test_capture_kwargs_hook(shared_model, llama_adapter, llama_tokenizer):
    inputs = _simple_input(shared_model, llama_tokenizer, shared_model.device)
    block0 = llama_adapter.get_layers()[0]

    with capture_kwargs_hook(block0) as all_kwargs:
        with torch.no_grad():
            shared_model(**inputs)

    assert len(all_kwargs) >= 1
    # attention_mask may be present depending on HF version
    assert isinstance(all_kwargs[0], dict)


def test_capture_kwargs_hook_removed(shared_model, llama_adapter):
    block0 = llama_adapter.get_layers()[0]
    before = len(block0._forward_pre_hooks)
    with capture_kwargs_hook(block0) as _:
        pass
    assert len(block0._forward_pre_hooks) == before


# ---------------------------------------------------------------------------
# 3. tensor_grad_hook
# ---------------------------------------------------------------------------

def test_tensor_grad_hook(shared_model, llama_adapter, llama_tokenizer):
    inputs = _simple_input(shared_model, llama_tokenizer, shared_model.device)
    target_id = llama_tokenizer(" France", add_special_tokens=False)["input_ids"][0]

    captured_act: list = []
    fwd_handle = llama_adapter.get_mlp(layer=0).register_forward_hook(
        lambda m, i, o: captured_act.append(o)
    )
    out = shared_model(**inputs)
    fwd_handle.remove()

    act = captured_act[0]
    with tensor_grad_hook(lambda: act) as grads:
        loss = -out.logits[0, -1, target_id]
        loss.backward()

    assert len(grads) == 1
    assert grads[0].shape == act.shape


# ---------------------------------------------------------------------------
# 4. autograd_node_hook
# ---------------------------------------------------------------------------

def test_autograd_node_hook(shared_model, llama_adapter, llama_tokenizer):
    inputs = _simple_input(shared_model, llama_tokenizer, shared_model.device)
    target_id = llama_tokenizer(" France", add_special_tokens=False)["input_ids"][0]

    captured_act: list = []
    fwd_handle = llama_adapter.get_mlp(layer=0).register_forward_hook(
        lambda m, i, o: captured_act.append(o)
    )
    out = shared_model(**inputs)
    fwd_handle.remove()

    act = captured_act[0]
    assert act.grad_fn is not None

    with autograd_node_hook(act) as node_grads:
        loss = -out.logits[0, -1, target_id]
        loss.backward()

    assert len(node_grads) >= 1


def test_autograd_node_hook_requires_grad_fn():
    leaf = torch.tensor([1.0])
    with pytest.raises(ValueError, match="grad_fn"):
        with autograd_node_hook(leaf):
            pass


# ---------------------------------------------------------------------------
# 5. capture_backward_pre
# ---------------------------------------------------------------------------

def test_capture_backward_pre(shared_model, llama_adapter, llama_tokenizer):
    inputs = _simple_input(shared_model, llama_tokenizer, shared_model.device)
    target_id = llama_tokenizer(" France", add_special_tokens=False)["input_ids"][0]

    attn = llama_adapter.get_attention_module(layer=0)
    with capture_backward_pre(attn) as pre_grads:
        out = shared_model(**inputs)
        loss = -out.logits[0, -1, target_id]
        loss.backward()

    assert len(pre_grads) == 1
    # grad_output is a tuple; first element should be a tensor
    assert pre_grads[0][0] is not None
