"""
Activation Collector for Transformer Models

Collects intermediate activations during forward passes for interpretability analysis.
Supports HuggingFace Transformers models with configurable collection points.
"""

from __future__ import annotations

import json
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Literal

import torch
import torch.nn as nn
from torch.utils.data import Dataset

# Type aliases
CollectionPoint = Literal["attention", "mlp", "residual", "all"]
LayerSelection = list[int] | Literal["all", "quarter", "eighth"]


@dataclass
class ActivationConfig:
    """Configuration for activation collection."""

    collection_points: list[CollectionPoint] = field(
        default_factory=lambda: ["residual"]
    )
    layers: LayerSelection = "quarter"  # Collect every 4th layer by default
    token_aggregation: Literal["last", "mean", "all"] = "last"
    dtype: torch.dtype = torch.float16
    device: str = "cuda"


@dataclass
class CollectedActivations:
    """Container for collected activations from a single forward pass."""

    prompt: str
    response: str
    activations: dict[str, torch.Tensor]  # "layer_N_type" -> tensor
    metadata: dict[str, Any] = field(default_factory=dict)

    def save(self, path: Path) -> None:
        """Save activations to disk."""
        path.mkdir(parents=True, exist_ok=True)

        # Save tensors
        tensors = {k: v.cpu() for k, v in self.activations.items()}
        torch.save(tensors, path / "activations.pt")

        # Save metadata
        meta = {
            "prompt": self.prompt,
            "response": self.response,
            "metadata": self.metadata,
            "activation_keys": list(self.activations.keys()),
        }
        with open(path / "metadata.json", "w") as f:
            json.dump(meta, f, indent=2)

    @classmethod
    def load(cls, path: Path) -> CollectedActivations:
        """Load activations from disk."""
        tensors = torch.load(path / "activations.pt", weights_only=True)
        with open(path / "metadata.json") as f:
            meta = json.load(f)

        return cls(
            prompt=meta["prompt"],
            response=meta["response"],
            activations=tensors,
            metadata=meta["metadata"],
        )


class ActivationCollector:
    """
    Collects activations from transformer models during generation.

    Usage:
        collector = ActivationCollector(model, tokenizer, config)
        result = collector.collect(prompt, max_new_tokens=100)
        # result.activations["layer_16_residual"] -> tensor of shape [seq_len, hidden_dim]
    """

    def __init__(
        self,
        model: nn.Module,
        tokenizer: Any,
        config: ActivationConfig | None = None,
    ):
        self.model = model
        self.tokenizer = tokenizer
        self.config = config or ActivationConfig()

        # Determine model architecture
        self.num_layers = self._get_num_layers()
        self.hidden_dim = self._get_hidden_dim()

        # Resolve layer selection
        self.target_layers = self._resolve_layers()

        # Storage for hooks
        self._hooks: list[Any] = []
        self._activations: dict[str, torch.Tensor] = {}

    def _get_num_layers(self) -> int:
        """Detect number of transformer layers."""
        # Common attribute names across architectures
        for attr in ["num_hidden_layers", "n_layer", "num_layers"]:
            if hasattr(self.model.config, attr):
                return getattr(self.model.config, attr)

        # Fallback: count layers directly
        if hasattr(self.model, "model") and hasattr(self.model.model, "layers"):
            return len(self.model.model.layers)

        raise ValueError("Could not determine number of layers")

    def _get_hidden_dim(self) -> int:
        """Detect hidden dimension."""
        for attr in ["hidden_size", "d_model", "n_embd"]:
            if hasattr(self.model.config, attr):
                return getattr(self.model.config, attr)
        raise ValueError("Could not determine hidden dimension")

    def _resolve_layers(self) -> list[int]:
        """Convert layer selection to explicit list of layer indices."""
        layers = self.config.layers

        if layers == "all":
            return list(range(self.num_layers))
        elif layers == "quarter":
            return list(range(0, self.num_layers, 4))
        elif layers == "eighth":
            return list(range(0, self.num_layers, 8))
        elif isinstance(layers, list):
            return [l if l >= 0 else self.num_layers + l for l in layers]
        else:
            raise ValueError(f"Unknown layer selection: {layers}")

    def _get_layer_module(self, layer_idx: int) -> nn.Module:
        """Get the transformer layer module by index."""
        # Try common attribute paths
        for path in ["model.layers", "transformer.h", "encoder.layer", "decoder.layers"]:
            parts = path.split(".")
            module = self.model
            try:
                for part in parts:
                    module = getattr(module, part)
                return module[layer_idx]
            except (AttributeError, IndexError, TypeError):
                continue

        raise ValueError(f"Could not find layer {layer_idx}")

    def _create_hook(self, name: str):
        """Create a forward hook that stores activations."""

        def hook(module: nn.Module, input: Any, output: Any) -> None:
            # Handle different output formats
            if isinstance(output, tuple):
                hidden_states = output[0]
            else:
                hidden_states = output

            # Apply token aggregation
            if self.config.token_aggregation == "last":
                activation = hidden_states[:, -1, :]  # [batch, hidden]
            elif self.config.token_aggregation == "mean":
                activation = hidden_states.mean(dim=1)  # [batch, hidden]
            else:  # "all"
                activation = hidden_states  # [batch, seq, hidden]

            # Store with configured dtype
            self._activations[name] = activation.detach().to(self.config.dtype)

        return hook

    def _register_hooks(self) -> None:
        """Register forward hooks on target layers."""
        self._clear_hooks()

        for layer_idx in self.target_layers:
            layer = self._get_layer_module(layer_idx)

            if "residual" in self.config.collection_points or "all" in self.config.collection_points:
                # Hook the full layer output (residual stream)
                hook = layer.register_forward_hook(
                    self._create_hook(f"layer_{layer_idx}_residual")
                )
                self._hooks.append(hook)

            if "attention" in self.config.collection_points or "all" in self.config.collection_points:
                # Hook attention output
                attn_module = self._get_attention_module(layer)
                if attn_module is not None:
                    hook = attn_module.register_forward_hook(
                        self._create_hook(f"layer_{layer_idx}_attention")
                    )
                    self._hooks.append(hook)

            if "mlp" in self.config.collection_points or "all" in self.config.collection_points:
                # Hook MLP output
                mlp_module = self._get_mlp_module(layer)
                if mlp_module is not None:
                    hook = mlp_module.register_forward_hook(
                        self._create_hook(f"layer_{layer_idx}_mlp")
                    )
                    self._hooks.append(hook)

    def _get_attention_module(self, layer: nn.Module) -> nn.Module | None:
        """Find the attention submodule within a layer."""
        for attr in ["self_attn", "attention", "attn"]:
            if hasattr(layer, attr):
                return getattr(layer, attr)
        return None

    def _get_mlp_module(self, layer: nn.Module) -> nn.Module | None:
        """Find the MLP submodule within a layer."""
        for attr in ["mlp", "feed_forward", "ffn", "ff"]:
            if hasattr(layer, attr):
                return getattr(layer, attr)
        return None

    def _clear_hooks(self) -> None:
        """Remove all registered hooks."""
        for hook in self._hooks:
            hook.remove()
        self._hooks = []
        self._activations = {}

    @torch.no_grad()
    def collect(
        self,
        prompt: str,
        max_new_tokens: int = 256,
        system_prompt: str | None = None,
        **generate_kwargs: Any,
    ) -> CollectedActivations:
        """
        Generate a response and collect activations.

        Args:
            prompt: User prompt
            max_new_tokens: Maximum tokens to generate
            system_prompt: Optional system prompt for persona steering
            **generate_kwargs: Additional arguments for model.generate()

        Returns:
            CollectedActivations with response text and activation tensors
        """
        # Build full prompt
        if system_prompt:
            messages = [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": prompt},
            ]
            if hasattr(self.tokenizer, "apply_chat_template"):
                full_prompt = self.tokenizer.apply_chat_template(
                    messages, tokenize=False, add_generation_prompt=True
                )
            else:
                full_prompt = f"{system_prompt}\n\nUser: {prompt}\n\nAssistant:"
        else:
            full_prompt = prompt

        # Tokenize
        inputs = self.tokenizer(
            full_prompt,
            return_tensors="pt",
            truncation=True,
            max_length=4096,
        ).to(self.config.device)

        # Register hooks
        self._register_hooks()

        try:
            # Generate with activation collection
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                pad_token_id=self.tokenizer.eos_token_id,
                **generate_kwargs,
            )

            # Decode response (only new tokens)
            response_tokens = outputs[0, inputs.input_ids.shape[1] :]
            response = self.tokenizer.decode(response_tokens, skip_special_tokens=True)

            return CollectedActivations(
                prompt=prompt,
                response=response,
                activations=dict(self._activations),
                metadata={
                    "system_prompt": system_prompt,
                    "target_layers": self.target_layers,
                    "collection_points": self.config.collection_points,
                    "token_aggregation": self.config.token_aggregation,
                },
            )

        finally:
            self._clear_hooks()

    def collect_on_text(self, text: str) -> CollectedActivations:
        """
        Collect activations on existing text (no generation).

        Useful for analyzing responses you already have.
        """
        inputs = self.tokenizer(
            text,
            return_tensors="pt",
            truncation=True,
            max_length=4096,
        ).to(self.config.device)

        self._register_hooks()

        try:
            # Single forward pass
            with torch.no_grad():
                self.model(**inputs)

            return CollectedActivations(
                prompt=text,
                response="",
                activations=dict(self._activations),
                metadata={
                    "mode": "text_only",
                    "target_layers": self.target_layers,
                    "collection_points": self.config.collection_points,
                },
            )

        finally:
            self._clear_hooks()


class ActivationDataset(Dataset):
    """
    PyTorch Dataset for activation data with personality labels.

    Loads activations and Lindr personality scores for probe training.
    """

    def __init__(
        self,
        data_dir: Path,
        target_layer: str = "layer_16_residual",
        dimensions: list[str] | None = None,
    ):
        """
        Args:
            data_dir: Directory containing collected activations
            target_layer: Which layer's activations to use
            dimensions: Which personality dimensions to include (default: all 10)
        """
        self.data_dir = Path(data_dir)
        self.target_layer = target_layer
        self.dimensions = dimensions or [
            "openness",
            "conscientiousness",
            "extraversion",
            "agreeableness",
            "neuroticism",
            "assertiveness",
            "ambition",
            "resilience",
            "integrity",
            "curiosity",
        ]

        # Find all sample directories
        self.samples = sorted(
            [d for d in self.data_dir.iterdir() if d.is_dir() and (d / "metadata.json").exists()]
        )

        if len(self.samples) == 0:
            raise ValueError(f"No samples found in {data_dir}")

    def __len__(self) -> int:
        return len(self.samples)

    def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
        sample_dir = self.samples[idx]

        # Load activations
        activations = torch.load(sample_dir / "activations.pt", weights_only=True)

        # Load metadata with personality scores
        with open(sample_dir / "metadata.json") as f:
            meta = json.load(f)

        # Extract target activation
        if self.target_layer not in activations:
            available = list(activations.keys())
            raise KeyError(f"Layer {self.target_layer} not found. Available: {available}")

        activation = activations[self.target_layer]

        # Extract personality labels
        scores = meta["metadata"].get("personality_scores", {})
        labels = torch.tensor(
            [scores.get(dim, 0.5) for dim in self.dimensions],
            dtype=torch.float32,
        )

        # Extract confidence (for weighted training)
        confidence = meta["metadata"].get("confidence", 1.0)

        return {
            "activation": activation.squeeze(0),  # [hidden_dim]
            "labels": labels,  # [num_dimensions]
            "confidence": torch.tensor(confidence, dtype=torch.float32),
        }

    def get_layer_keys(self) -> list[str]:
        """Get available layer keys from first sample."""
        if len(self.samples) == 0:
            return []
        activations = torch.load(self.samples[0] / "activations.pt", weights_only=True)
        return list(activations.keys())
