"""Core intercept decorator for activation caching and patching."""

from functools import wraps
from typing import Any, Callable, Dict, Set, Union

InterventionFn = Callable[[Any], Any]
Intervention = Union[Any, InterventionFn]


def intercept(fn: Callable) -> Callable:
    """
    Decorator that adds activation caching and patching to a JAX function.

    The decorated function must accept a `_save` parameter, which is a
    function with signature: _save(name: str, value: Array) -> Array

    The `_save` function:
    - Caches the value if the name is in the `cache` set
    - Applies an intervention if the name is in the `patch` dict
    - Returns the (possibly modified) value

    Args:
        fn: Function with signature fn(*args, _save, **kwargs) -> output

    Returns:
        Wrapped function with signature:
        fn(*args, cache: Set[str] = set(), patch: Dict[str, Intervention] = {},
           **kwargs) -> (output, cached_activations)

    Example:
        @intercept
        def mlp(params, x, _save):
            x = jnp.dot(x, params['w1']) + params['b1']
            x = _save('pre_relu', x)
            x = jax.nn.relu(x)
            x = _save('post_relu', x)
            return x

        # Basic usage - no caching
        output, cache = mlp(params, x)

        # Cache specific points
        output, cache = mlp(params, x, cache={'pre_relu', 'post_relu'})

        # Patch - replace activation
        output, cache = mlp(params, x, patch={'post_relu': zeros})

        # Patch - modify activation
        output, cache = mlp(params, x, patch={'post_relu': lambda act: act * 0.5})
    """

    @wraps(fn)
    def wrapper(
        *args,
        cache: Set[str] = frozenset(),
        patch: Dict[str, Intervention] | None = None,
        **kwargs,
    ):
        if patch is None:
            patch = {}

        cached: Dict[str, Any] = {}

        def _save(name: str, val: Any) -> Any:
            # Cache if requested
            if name in cache:
                cached[name] = val

            # Apply intervention if provided
            if name in patch:
                intervention = patch[name]
                if callable(intervention):
                    return intervention(val)
                else:
                    return intervention

            return val

        output = fn(*args, _save=_save, **kwargs)
        return output, cached

    return wrapper
