from __future__ import annotations

from typing import TYPE_CHECKING
import torch

from amber.hooks.detector import Detector
from amber.hooks.hook import HookType, HOOK_FUNCTION_INPUT, HOOK_FUNCTION_OUTPUT
from amber.hooks.utils import extract_tensor_from_output

if TYPE_CHECKING:
    from torch import nn


class LayerActivationDetector(Detector):
    """
    Detector hook that captures and saves activations during inference.
    
    This detector extracts activations from layer outputs and stores them
    for later use (e.g., saving to disk, further analysis).
    """

    def __init__(
            self,
            layer_signature: str | int,
            hook_id: str | None = None
    ):
        """
        Initialize the activation saver detector.
        
        Args:
            layer_signature: Layer to capture activations from
            hook_id: Unique identifier for this hook
            
        Raises:
            ValueError: If layer_signature is None
        """
        if layer_signature is None:
            raise ValueError("layer_signature cannot be None for LayerActivationDetector")

        super().__init__(
            hook_type=HookType.FORWARD,
            hook_id=hook_id,
            store=None,
            layer_signature=layer_signature
        )

    def process_activations(
            self,
            module: torch.nn.Module,
            input: HOOK_FUNCTION_INPUT,
            output: HOOK_FUNCTION_OUTPUT
    ) -> None:
        """
        Extract and store activations from output.
        
        Handles various output types:
        - Plain tensors
        - Tuples/lists of tensors (takes first tensor)
        - Objects with last_hidden_state attribute (e.g., HuggingFace outputs)
        
        Args:
            module: The PyTorch module being hooked
            input: Tuple of input tensors to the module
            output: Output tensor(s) from the module
            
        Raises:
            RuntimeError: If tensor extraction or storage fails
        """
        try:
            tensor = extract_tensor_from_output(output)

            if tensor is not None:
                tensor_cpu = tensor.detach().to("cpu")
                # Store current batch's tensor (overwrites previous)
                self.tensor_metadata['activations'] = tensor_cpu
                # Store activations shape to metadata
                self.metadata['activations_shape'] = tuple(tensor_cpu.shape)
        except Exception as e:
            raise RuntimeError(
                f"Error extracting activations in LayerActivationDetector {self.id}: {e}"
            ) from e

    def get_captured(self) -> torch.Tensor | None:
        """
        Get the captured activations from the current batch.
        
        Returns:
            The captured activation tensor from the current batch or None if no activations captured yet
        """
        return self.tensor_metadata.get('activations')

    def clear_captured(self) -> None:
        """Clear captured activations for current batch."""
        self.tensor_metadata.pop('activations', None)
        self.metadata.pop('activations_shape', None)

