"""
Personality Steering via Activation Intervention

Tools for constructing steering vectors and applying causal interventions
to modify LLM personality during generation.
"""

from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable

import torch
import torch.nn as nn

from .collector import ActivationDataset, CollectedActivations


DIMENSIONS = [
    "openness",
    "conscientiousness",
    "extraversion",
    "agreeableness",
    "neuroticism",
    "assertiveness",
    "ambition",
    "resilience",
    "integrity",
    "curiosity",
]


@dataclass
class SteeringVector:
    """
    A direction in activation space that corresponds to a personality dimension.

    Constructed from the difference between high-trait and low-trait activations.
    Can be added to activations during generation to steer model personality.
    """

    dimension: str
    layer: str
    vector: torch.Tensor
    metadata: dict[str, Any]

    @classmethod
    def from_activations(
        cls,
        dimension: str,
        layer: str,
        high_activations: list[torch.Tensor],
        low_activations: list[torch.Tensor],
    ) -> SteeringVector:
        """
        Construct steering vector from high/low trait activation sets.

        Args:
            dimension: Personality dimension name
            layer: Layer identifier (e.g., "layer_16_residual")
            high_activations: Activations from high-trait responses
            low_activations: Activations from low-trait responses

        Returns:
            SteeringVector pointing from low to high trait
        """
        high_mean = torch.stack(high_activations).mean(dim=0)
        low_mean = torch.stack(low_activations).mean(dim=0)

        vector = high_mean - low_mean

        return cls(
            dimension=dimension,
            layer=layer,
            vector=vector,
            metadata={
                "num_high_samples": len(high_activations),
                "num_low_samples": len(low_activations),
                "vector_norm": vector.norm().item(),
            },
        )

    @classmethod
    def from_dataset(
        cls,
        dataset: ActivationDataset,
        dimension: str,
        layer: str,
        high_threshold: float = 0.7,
        low_threshold: float = 0.3,
    ) -> SteeringVector:
        """
        Construct steering vector from activation dataset.

        Args:
            dataset: Dataset with activations and personality labels
            dimension: Dimension to extract
            layer: Layer to use
            high_threshold: Score threshold for high trait
            low_threshold: Score threshold for low trait
        """
        dim_idx = DIMENSIONS.index(dimension)
        high_activations = []
        low_activations = []

        for i in range(len(dataset)):
            sample = dataset[i]
            score = sample["labels"][dim_idx].item()

            if score >= high_threshold:
                high_activations.append(sample["activation"])
            elif score <= low_threshold:
                low_activations.append(sample["activation"])

        if len(high_activations) < 10 or len(low_activations) < 10:
            raise ValueError(
                f"Insufficient samples: {len(high_activations)} high, {len(low_activations)} low. "
                "Consider adjusting thresholds."
            )

        return cls.from_activations(
            dimension=dimension,
            layer=layer,
            high_activations=high_activations,
            low_activations=low_activations,
        )

    def normalize(self) -> SteeringVector:
        """Return unit-normalized steering vector."""
        return SteeringVector(
            dimension=self.dimension,
            layer=self.layer,
            vector=self.vector / self.vector.norm(),
            metadata={**self.metadata, "normalized": True},
        )

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

        torch.save(self.vector, path / "vector.pt")
        with open(path / "metadata.json", "w") as f:
            json.dump(
                {
                    "dimension": self.dimension,
                    "layer": self.layer,
                    "metadata": self.metadata,
                },
                f,
                indent=2,
            )

    @classmethod
    def load(cls, path: Path) -> SteeringVector:
        """Load steering vector from disk."""
        path = Path(path)
        vector = torch.load(path / "vector.pt", weights_only=True)

        with open(path / "metadata.json") as f:
            config = json.load(f)

        return cls(
            dimension=config["dimension"],
            layer=config["layer"],
            vector=vector,
            metadata=config["metadata"],
        )


@dataclass
class SteeringResult:
    """Result of a steering intervention."""

    original_response: str
    steered_response: str
    original_scores: dict[str, float]
    steered_scores: dict[str, float]
    intervention_config: dict[str, Any]

    @property
    def dimension_shift(self) -> dict[str, float]:
        """Compute per-dimension change from intervention."""
        return {
            dim: self.steered_scores.get(dim, 0) - self.original_scores.get(dim, 0)
            for dim in DIMENSIONS
        }

    @property
    def target_shift(self) -> float:
        """Shift in the targeted dimension."""
        target_dim = self.intervention_config.get("dimension")
        if target_dim:
            return self.dimension_shift.get(target_dim, 0)
        return 0.0


class PersonalitySteerer:
    """
    Applies steering vectors to modify model personality during generation.

    Usage:
        steerer = PersonalitySteerer(model, tokenizer)
        steerer.load_vector(steering_vector)
        result = steerer.generate_steered(prompt, strength=1.0)
    """

    def __init__(
        self,
        model: nn.Module,
        tokenizer: Any,
        device: str = "cuda",
    ):
        self.model = model
        self.tokenizer = tokenizer
        self.device = device

        self.steering_vectors: dict[str, SteeringVector] = {}
        self._hooks: list[Any] = []

    def load_vector(self, vector: SteeringVector) -> None:
        """Load a steering vector for use during generation."""
        self.steering_vectors[vector.dimension] = vector

    def load_vectors_from_dir(self, directory: Path) -> None:
        """Load all steering vectors from a directory."""
        directory = Path(directory)
        for dim_dir in directory.iterdir():
            if dim_dir.is_dir() and (dim_dir / "vector.pt").exists():
                vector = SteeringVector.load(dim_dir)
                self.load_vector(vector)

    def _get_layer_module(self, layer_name: str) -> nn.Module:
        """Get layer module from layer name (e.g., 'layer_16_residual')."""
        parts = layer_name.split("_")
        layer_idx = int(parts[1])

        # Navigate to layer
        for path in ["model.layers", "transformer.h", "encoder.layer"]:
            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 for {layer_name}")

    def _create_steering_hook(
        self,
        vector: torch.Tensor,
        strength: float,
        position: str = "all",
    ) -> Callable:
        """Create hook that adds steering vector to activations."""

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

            # Reshape vector for broadcasting
            steering = vector.to(hidden_states.device, hidden_states.dtype)
            if steering.dim() == 1:
                steering = steering.unsqueeze(0).unsqueeze(0)  # [1, 1, hidden]

            # Apply to specified positions
            if position == "all":
                hidden_states = hidden_states + strength * steering
            elif position == "last":
                hidden_states[:, -1, :] = hidden_states[:, -1, :] + strength * steering.squeeze()
            elif position == "first":
                hidden_states[:, 0, :] = hidden_states[:, 0, :] + strength * steering.squeeze()

            if rest is not None:
                return (hidden_states,) + rest
            return hidden_states

        return hook

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

    @torch.no_grad()
    def generate_steered(
        self,
        prompt: str,
        dimension: str,
        strength: float = 1.0,
        max_new_tokens: int = 256,
        system_prompt: str | None = None,
        position: str = "all",
        **generate_kwargs: Any,
    ) -> str:
        """
        Generate response with personality steering.

        Args:
            prompt: User prompt
            dimension: Which personality dimension to steer
            strength: Steering strength (negative = opposite direction)
            max_new_tokens: Maximum tokens to generate
            system_prompt: Optional system prompt
            position: Where to apply steering ("all", "last", "first")

        Returns:
            Generated response text
        """
        if dimension not in self.steering_vectors:
            raise ValueError(f"No steering vector loaded for {dimension}")

        vector = self.steering_vectors[dimension]

        # Build 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.device)

        # Register steering hook
        layer_module = self._get_layer_module(vector.layer)
        hook = layer_module.register_forward_hook(
            self._create_steering_hook(vector.vector, strength, position)
        )
        self._hooks.append(hook)

        try:
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                pad_token_id=self.tokenizer.eos_token_id,
                **generate_kwargs,
            )

            response_tokens = outputs[0, inputs.input_ids.shape[1] :]
            return self.tokenizer.decode(response_tokens, skip_special_tokens=True)

        finally:
            self._clear_hooks()

    @torch.no_grad()
    def generate_comparison(
        self,
        prompt: str,
        dimension: str,
        strengths: list[float] = [-1.0, 0.0, 1.0],
        max_new_tokens: int = 256,
        system_prompt: str | None = None,
        **generate_kwargs: Any,
    ) -> dict[float, str]:
        """
        Generate responses at multiple steering strengths for comparison.

        Args:
            prompt: User prompt
            dimension: Dimension to steer
            strengths: List of steering strengths to test
            max_new_tokens: Maximum tokens per generation
            system_prompt: Optional system prompt

        Returns:
            Dict mapping strength to generated response
        """
        results = {}

        for strength in strengths:
            if strength == 0.0:
                # Baseline without steering
                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

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

                outputs = self.model.generate(
                    **inputs,
                    max_new_tokens=max_new_tokens,
                    pad_token_id=self.tokenizer.eos_token_id,
                    **generate_kwargs,
                )

                response_tokens = outputs[0, inputs.input_ids.shape[1] :]
                results[strength] = self.tokenizer.decode(
                    response_tokens, skip_special_tokens=True
                )
            else:
                results[strength] = self.generate_steered(
                    prompt=prompt,
                    dimension=dimension,
                    strength=strength,
                    max_new_tokens=max_new_tokens,
                    system_prompt=system_prompt,
                    **generate_kwargs,
                )

        return results


class ActivationPatcher:
    """
    Patches activations from one generation into another.

    Used for causal intervention experiments: if patching source activations
    causes target to behave like source, those activations causally control
    the behavior difference.
    """

    def __init__(self, model: nn.Module, tokenizer: Any, device: str = "cuda"):
        self.model = model
        self.tokenizer = tokenizer
        self.device = device
        self._hooks: list[Any] = []

    def _get_layer_module(self, layer_idx: int) -> nn.Module:
        """Get layer module by index."""
        for path in ["model.layers", "transformer.h", "encoder.layer"]:
            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 _clear_hooks(self) -> None:
        """Remove all hooks."""
        for hook in self._hooks:
            hook.remove()
        self._hooks = []

    @torch.no_grad()
    def patch_and_generate(
        self,
        prompt: str,
        source_activations: CollectedActivations,
        target_layers: list[int],
        max_new_tokens: int = 256,
        patch_mode: str = "replace",  # "replace" or "add"
        **generate_kwargs: Any,
    ) -> str:
        """
        Generate with patched activations from source.

        Args:
            prompt: Prompt to generate response for
            source_activations: Activations to patch from
            target_layers: Which layers to patch
            max_new_tokens: Max tokens to generate
            patch_mode: "replace" overwrites, "add" adds difference

        Returns:
            Generated response with patched activations
        """
        inputs = self.tokenizer(
            prompt,
            return_tensors="pt",
            truncation=True,
            max_length=4096,
        ).to(self.device)

        # Register patching hooks
        for layer_idx in target_layers:
            layer_key = f"layer_{layer_idx}_residual"
            if layer_key not in source_activations.activations:
                continue

            source_act = source_activations.activations[layer_key]
            layer_module = self._get_layer_module(layer_idx)

            def make_hook(src_act: torch.Tensor, mode: str):
                def hook(module: nn.Module, input: Any, output: Any) -> Any:
                    if isinstance(output, tuple):
                        hidden = output[0]
                        rest = output[1:]
                    else:
                        hidden = output
                        rest = None

                    src = src_act.to(hidden.device, hidden.dtype)

                    # Handle sequence length mismatch
                    if src.dim() == 1:
                        src = src.unsqueeze(0).unsqueeze(0)
                    if src.shape[1] != hidden.shape[1]:
                        # Pad or truncate
                        if src.shape[1] < hidden.shape[1]:
                            src = src.expand(-1, hidden.shape[1], -1)
                        else:
                            src = src[:, : hidden.shape[1], :]

                    if mode == "replace":
                        hidden = src.expand_as(hidden)
                    else:  # add
                        hidden = hidden + src

                    if rest is not None:
                        return (hidden,) + rest
                    return hidden

                return hook

            hook = layer_module.register_forward_hook(make_hook(source_act, patch_mode))
            self._hooks.append(hook)

        try:
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                pad_token_id=self.tokenizer.eos_token_id,
                **generate_kwargs,
            )

            response_tokens = outputs[0, inputs.input_ids.shape[1] :]
            return self.tokenizer.decode(response_tokens, skip_special_tokens=True)

        finally:
            self._clear_hooks()
