"""Generic transformer tracing built on standard PyTorch hooks."""

from __future__ import annotations

from collections.abc import Iterable
from contextlib import ExitStack
from typing import Any

from torch import Tensor, nn

from ..models.adapters import ModelAdapter, get_adapter
from .logits import project_layer_logits
from .residuals import resolve_post_attention_anchor
from .sites import ALL_SITES, SiteName
from .types import TraceResult


def _extract_tensor(value: Any) -> Tensor:
    """Extract the main tensor from hook inputs or outputs."""

    if isinstance(value, Tensor):
        return value
    if isinstance(value, (list, tuple)):
        for item in value:
            if isinstance(item, Tensor):
                return item
    raise TypeError(f"Could not extract a tensor from hook value of type {type(value)!r}.")


def _clone_tensor(tensor: Tensor, *, detach: bool, clone: bool) -> Tensor:
    if detach:
        tensor = tensor.detach()
    if clone:
        tensor = tensor.clone()
    return tensor


def _normalize_layers(adapter: ModelAdapter, layers: Iterable[int] | None) -> tuple[int, ...]:
    if layers is None:
        return tuple(range(adapter.num_layers))
    normalized = tuple(sorted(set(layers)))
    for layer_idx in normalized:
        if layer_idx < 0 or layer_idx >= adapter.num_layers:
            raise IndexError(f"Requested layer index {layer_idx} outside valid range.")
    return normalized


def _normalize_sites(sites: Iterable[SiteName] | None) -> tuple[SiteName, ...]:
    if sites is None:
        return ("block_output", "final_norm", "logits")
    normalized = tuple(sites)
    invalid = [site for site in normalized if site not in ALL_SITES]
    if invalid:
        raise ValueError(f"Unsupported sites requested: {invalid}.")
    return normalized


def _register_capture(
    module: nn.Module,
    *,
    on_input: bool = False,
    site: SiteName | None = None,
    layer: int | None = None,
    trace: TraceResult | None = None,
    store: dict[int, Tensor] | None = None,
    detach: bool,
    clone: bool,
) -> Any:
    """Register a forward hook that captures either input or output tensors."""
    if on_input:
        def hook(_module: nn.Module, inputs: tuple[Any, ...]) -> None:
            t = _clone_tensor(_extract_tensor(inputs), detach=detach, clone=clone)
            if store is not None:
                store[layer] = t  # type: ignore[index]
            else:
                trace.add(site, t, layer=layer)  # type: ignore[union-attr]

        return module.register_forward_pre_hook(hook)
    else:
        def hook(_module: nn.Module, _inputs: tuple[Any, ...], output: Any) -> None:  # type: ignore[misc]
            t = _clone_tensor(_extract_tensor(output), detach=detach, clone=clone)
            if store is not None:
                store[layer] = t  # type: ignore[index]
            else:
                trace.add(site, t, layer=layer)  # type: ignore[union-attr]

        return module.register_forward_hook(hook)


class TransformerTracer:
    """Trace hidden states and logits from named transformer sites."""

    def __init__(
        self,
        model: nn.Module,
        adapter: ModelAdapter | None = None,
    ):
        self.model = model
        self.adapter = adapter if adapter is not None else get_adapter(model)

    def trace(
        self,
        input_ids: Tensor,
        *,
        attention_mask: Tensor | None = None,
        layers: Iterable[int] | None = None,
        sites: Iterable[SiteName] | None = None,
        detach: bool = True,
        clone: bool = True,
        **model_kwargs: Any,
    ) -> TraceResult:
        requested_sites = _normalize_sites(sites)
        layer_indices = _normalize_layers(self.adapter, layers)
        trace = TraceResult(
            metadata={
                "requested_layers": layer_indices,
                "requested_sites": requested_sites,
            }
        )

        layers_for_logits = set(layer_indices) if "logits" in requested_sites else set()
        hidden_for_logits: dict[int, Tensor] = {}
        model_layers = self.adapter.get_layers()

        with ExitStack() as stack:
            for layer_idx in layer_indices:
                layer = model_layers[layer_idx]
                if "block_input" in requested_sites:
                    stack.enter_context(
                        _HookHandleContext(
                            _register_capture(
                                layer,
                                on_input=True,
                                site="block_input",
                                layer=layer_idx,
                                trace=trace,
                                detach=detach,
                                clone=clone,
                            )
                        )
                    )
                if "block_output" in requested_sites:
                    handle = _register_capture(
                        layer,
                        site="block_output",
                        layer=layer_idx,
                        trace=trace,
                        detach=detach,
                        clone=clone,
                    )
                    stack.enter_context(_HookHandleContext(handle))
                elif layer_idx in layers_for_logits:
                    handle = _register_capture(
                        layer,
                        store=hidden_for_logits,
                        layer=layer_idx,
                        detach=detach,
                        clone=clone,
                    )
                    stack.enter_context(_HookHandleContext(handle))
                if "attn_o_proj_pre" in requested_sites:
                    stack.enter_context(
                        _HookHandleContext(
                            _register_capture(
                                self.adapter.get_attention_output_proj(layer),
                                on_input=True,
                                site="attn_o_proj_pre",
                                layer=layer_idx,
                                trace=trace,
                                detach=detach,
                                clone=clone,
                            )
                        )
                    )
                if "attn_output" in requested_sites:
                    stack.enter_context(
                        _HookHandleContext(
                            _register_capture(
                                self.adapter.get_attention_module(layer),
                                site="attn_output",
                                layer=layer_idx,
                                trace=trace,
                                detach=detach,
                                clone=clone,
                            )
                        )
                    )
                if "post_attn_residual" in requested_sites:
                    stack.enter_context(
                        _HookHandleContext(
                            _register_capture(
                                resolve_post_attention_anchor(self.adapter, layer),
                                on_input=True,
                                site="post_attn_residual",
                                layer=layer_idx,
                                trace=trace,
                                detach=detach,
                                clone=clone,
                            )
                        )
                    )
                if "mlp_input" in requested_sites:
                    stack.enter_context(
                        _HookHandleContext(
                            _register_capture(
                                self.adapter.get_mlp(layer),
                                on_input=True,
                                site="mlp_input",
                                layer=layer_idx,
                                trace=trace,
                                detach=detach,
                                clone=clone,
                            )
                        )
                    )
                if "mlp_output" in requested_sites:
                    stack.enter_context(
                        _HookHandleContext(
                            _register_capture(
                                self.adapter.get_mlp(layer),
                                site="mlp_output",
                                layer=layer_idx,
                                trace=trace,
                                detach=detach,
                                clone=clone,
                            )
                        )
                    )

            if "final_norm" in requested_sites:
                final_norm = self.adapter.get_final_norm()
                if final_norm is None:
                    raise ValueError("This model does not expose a final norm module.")
                stack.enter_context(
                    _HookHandleContext(
                        _register_capture(
                            final_norm,
                            site="final_norm",
                            layer=None,
                            trace=trace,
                            detach=detach,
                            clone=clone,
                        )
                    )
                )

            if "lm_head_input" in requested_sites:
                stack.enter_context(
                    _HookHandleContext(
                        _register_capture(
                            self.adapter.get_lm_head(),
                            on_input=True,
                            site="lm_head_input",
                            layer=None,
                            trace=trace,
                            detach=detach,
                            clone=clone,
                        )
                    )
                )

            outputs = self.model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                **model_kwargs,
            )

        trace.final_logits = outputs.logits.detach().clone() if hasattr(outputs, "logits") else None

        if "logits" in requested_sites:
            if "block_output" in trace.activations:
                hidden_for_logits = {
                    layer_idx: trace.get("block_output", layer_idx)
                    for layer_idx in layer_indices
                    if trace.has("block_output", layer_idx)
                }
            if hidden_for_logits:
                projected = project_layer_logits(
                    self.adapter,
                    hidden_for_logits,
                    source_site="block_output",
                )
                for layer_idx, logits in projected.items():
                    trace.add("logits", logits, layer_idx)
            else:
                raise ValueError("The tracer could not capture block outputs for logit projection.")

        return trace


class _HookHandleContext:
    """Light wrapper to use removable hook handles with ExitStack."""

    def __init__(self, handle: Any):
        self.handle = handle

    def __enter__(self) -> Any:
        return self.handle

    def __exit__(self, *_exc_info: Any) -> None:
        self.handle.remove()
