"""Minimal hook-based interventions for module-backed transformer sites."""

from __future__ import annotations

import contextlib
from collections.abc import Callable, Iterable
from dataclasses import dataclass
from typing import Any

from torch import Tensor, nn

from ..models.adapters import ModelAdapter, get_adapter
from .sites import SiteName
from .tracer import _extract_tensor

TensorTransform = Callable[[Tensor], Tensor]


@dataclass(frozen=True)
class InterventionSpec:
    """A tensor transform applied at a specific module-backed site."""

    site: SiteName
    transform: TensorTransform
    layer: int | None = None


def _replace_first_tensor(original: Any, updated: Tensor) -> Any:
    if isinstance(original, Tensor):
        return updated
    if isinstance(original, tuple):
        values = list(original)
        for idx, item in enumerate(values):
            if isinstance(item, Tensor):
                values[idx] = updated
                return tuple(values)
    if isinstance(original, list):
        values = list(original)
        for idx, item in enumerate(values):
            if isinstance(item, Tensor):
                values[idx] = updated
                return values
    raise TypeError(f"Unsupported hook payload type for intervention: {type(original)!r}.")


def _resolve_site_module(
    adapter: ModelAdapter,
    site: SiteName,
    layer: int | None,
) -> tuple[nn.Module, str]:
    if site == "final_norm":
        module = adapter.get_final_norm()
        if module is None:
            raise ValueError("This model does not expose a final norm module.")
        return module, "output"
    if site == "lm_head_input":
        return adapter.get_lm_head(), "input"
    if layer is None:
        raise ValueError(f"Layer index is required for site {site!r}.")

    block = adapter.get_layers()[layer]
    if site == "attn_o_proj_pre":
        return adapter.get_attention_output_proj(block), "input"
    if site == "attn_output":
        return adapter.get_attention_module(block), "output"
    if site == "mlp_input":
        return adapter.get_mlp(block), "input"
    if site == "mlp_output":
        return adapter.get_mlp(block), "output"
    raise ValueError(f"Interventions are not supported at site {site!r}.")


@contextlib.contextmanager
def apply_interventions(
    model: nn.Module,
    interventions: Iterable[InterventionSpec],
    *,
    adapter: ModelAdapter | None = None,
):
    """Apply temporary interventions at module-backed transformer sites."""

    resolved_adapter = adapter if adapter is not None else get_adapter(model)
    hooks = []

    try:
        for spec in interventions:
            module, hook_kind = _resolve_site_module(resolved_adapter, spec.site, spec.layer)
            if hook_kind == "input":
                def pre_hook(
                    _module: nn.Module,
                    inputs: tuple[Any, ...],
                    transform: TensorTransform = spec.transform,
                ) -> tuple[Any, ...]:
                    tensor = _extract_tensor(inputs)
                    return (transform(tensor),)

                hooks.append(module.register_forward_pre_hook(pre_hook))
            else:
                def post_hook(
                    _module: nn.Module,
                    _inputs: tuple[Any, ...],
                    output: Any,
                    transform: TensorTransform = spec.transform,
                ) -> Any:
                    tensor = _extract_tensor(output)
                    return _replace_first_tensor(output, transform(tensor))

                hooks.append(module.register_forward_hook(post_hook))
        yield
    finally:
        for hook in hooks:
            hook.remove()
