"""
Unified Wisent interface for multi-modal contrastive steering.

This module provides a high-level API that works across all modalities
while maintaining backward compatibility with the text-only API.
"""
from __future__ import annotations

from typing import Any, Dict, List, Optional, Union, TypeVar
from dataclasses import dataclass, field
import logging

import torch

from wisent.core.adapters.base import BaseAdapter, SteeringConfig, AdapterError
from wisent.core.constants import DEFAULT_SCALE
from wisent.core.adapters.text import TextAdapter
from wisent.core.adapters import AudioAdapter
from wisent.core.adapters import VideoAdapter
from wisent.core.adapters import RoboticsAdapter
from wisent.core.adapters import MultimodalAdapter
from wisent.core.errors import UnknownTypeError, NoTrainedVectorsError
from wisent.core.modalities import (
    Modality,
    ModalityContent,
    ContentType,
    TextContent,
    AudioContent,
    VideoContent,
    ImageContent,
    RobotState,
    RobotTrajectory,
    MultimodalContent,
    wrap_content,
)
from wisent.core.activations.core.atoms import LayerActivations
from wisent.core.steering_methods.core.atoms import BaseSteeringMethod

__all__ = [
    "Wisent",
    "TraitConfig",
]

logger = logging.getLogger(__name__)

InputT = TypeVar("InputT", bound=ModalityContent)
OutputT = TypeVar("OutputT")

from wisent.core._wisent_pairs import WisentPairsMixin
from wisent.core._wisent_training import WisentTrainingMixin
from wisent.core._wisent_io import WisentIOMixin

@dataclass
class TraitConfig:
    """
    Configuration for a steering trait.

    Attributes:
        name: Unique identifier for the trait
        description: Human-readable description
        steering_vectors: Per-layer steering vectors
        default_scale: Default steering strength
        layers: Which layers to apply to (None = all recommended)
    """
    name: str
    description: str = ""
    steering_vectors: LayerActivations | None = None
    default_scale: float = DEFAULT_SCALE
    layers: List[str] | None = None


class Wisent(WisentPairsMixin, WisentTrainingMixin, WisentIOMixin):
    """
    Unified interface for multi-modal contrastive steering.

    Wisent provides a simple, consistent API for steering AI models across
    different modalities (text, audio, video, robotics). The core workflow is:

    1. Create a Wisent instance with an appropriate adapter
    2. Add contrastive pairs (positive/negative examples)
    3. Train steering vectors from the pairs
    4. Generate outputs with steering applied

    Example (Text):
        >>> wisent = Wisent.for_text("meta-llama/Llama-3-8B-Instruct")
        >>> wisent.add_pair(
        ...     positive="I'd be happy to help you with that.",
        ...     negative="I refuse to help with that.",
        ...     trait="helpfulness"
        ... )
        >>> wisent.train()
        >>> response = wisent.generate("How do I cook pasta?", steer={"helpfulness": 1.5})

    Example (Audio):
        >>> wisent = Wisent.for_audio("openai/whisper-large-v3")
        >>> wisent.add_pair(
        ...     positive=AudioContent.from_file("calm_speech.wav"),
        ...     negative=AudioContent.from_file("angry_speech.wav"),
        ...     trait="calmness"
        ... )
        >>> wisent.train()
        >>> transcript = wisent.generate(audio_input, steer={"calmness": 1.2})

    Example (Video):
        >>> wisent = Wisent.for_video("MCG-NJU/videomae-base")
        >>> wisent.add_pair(
        ...     positive=VideoContent.from_file("safe_action.mp4"),
        ...     negative=VideoContent.from_file("unsafe_action.mp4"),
        ...     trait="safety"
        ... )
        >>> wisent.train()
        >>> embedding = wisent.generate(video_input, steer={"safety": 2.0})

    Example (Robotics):
        >>> wisent = Wisent.for_robotics(model=my_policy_network)
        >>> wisent.add_pair(
        ...     positive=gentle_trajectory,
        ...     negative=forceful_trajectory,
        ...     trait="gentleness"
        ... )
        >>> wisent.train()
        >>> action = wisent.act(robot_state, steer={"gentleness": 1.5})
    """

    def __init__(
        self,
        adapter: BaseAdapter,
        steering_method: str = "caa",
        device: str | None = None,
    ):
        """
        Initialize Wisent with an adapter.

        Args:
            adapter: The modality-specific adapter to use
            steering_method: Method for computing steering vectors ("caa", etc.)
            device: Override device (uses adapter's device by default)
        """
        self.adapter = adapter
        self.steering_method_name = steering_method
        self.device = device or adapter.device

        # Storage for pairs and trained vectors
        self._pairs: Dict[str, List[tuple[ModalityContent, ModalityContent]]] = {}
        self._traits: Dict[str, TraitConfig] = {}
        self._trained = False

    # ==================== Factory Methods ====================

    @classmethod
    def for_text(
        cls,
        model_name: str,
        device: str | None = None,
        **kwargs: Any,
    ) -> "Wisent":
        """
        Create a Wisent instance for text/LLM steering.

        Args:
            model_name: HuggingFace model identifier
            device: Target device
            **kwargs: Additional adapter arguments

        Returns:
            Configured Wisent instance
        """
        adapter = TextAdapter(model_name=model_name, device=device, **kwargs)
        return cls(adapter)

    @classmethod
    def for_audio(
        cls,
        model_name: str,
        device: str | None = None,
        **kwargs: Any,
    ) -> "Wisent":
        """
        Create a Wisent instance for audio steering.

        Args:
            model_name: HuggingFace model identifier
            device: Target device
            **kwargs: Additional adapter arguments

        Returns:
            Configured Wisent instance
        """
        adapter = AudioAdapter(model_name=model_name, device=device, **kwargs)
        return cls(adapter)

    @classmethod
    def for_video(
        cls,
        model_name: str,
        device: str | None = None,
        **kwargs: Any,
    ) -> "Wisent":
        """
        Create a Wisent instance for video steering.

        Args:
            model_name: HuggingFace model identifier
            device: Target device
            **kwargs: Additional adapter arguments

        Returns:
            Configured Wisent instance
        """
        adapter = VideoAdapter(model_name=model_name, device=device, **kwargs)
        return cls(adapter)

    @classmethod
    def for_robotics(
        cls,
        model: torch.nn.Module,
        device: str | None = None,
        **kwargs: Any,
    ) -> "Wisent":
        """
        Create a Wisent instance for robotics policy steering.

        Args:
            model: Policy network
            device: Target device
            **kwargs: Additional adapter arguments

        Returns:
            Configured Wisent instance
        """
        adapter = RoboticsAdapter(model=model, device=device, **kwargs)
        return cls(adapter)

    @classmethod
    def for_multimodal(
        cls,
        model_name: str,
        device: str | None = None,
        **kwargs: Any,
    ) -> "Wisent":
        """
        Create a Wisent instance for multimodal (VLM, etc.) steering.

        Args:
            model_name: HuggingFace model identifier
            device: Target device
            **kwargs: Additional adapter arguments

        Returns:
            Configured Wisent instance
        """
        adapter = MultimodalAdapter(model_name=model_name, device=device, **kwargs)
        return cls(adapter)

    # ==================== Pair Management ====================

