"""
neuroflow.models -- Pydantic data models for the NeuroFlow SDK.

Provides typed, validated representations of all entities in the
NeuroFlow platform: devices, sessions, streaming data, and extracted
features. These models are used both for API responses and for
local data access.

All models use Pydantic v2 for runtime validation, serialization,
and JSON schema generation.
"""

from __future__ import annotations

from datetime import datetime
from enum import Enum
from typing import Any, Optional

from pydantic import BaseModel, Field


# ──────────────────────────────────────────────────────────────
# Enums
# ──────────────────────────────────────────────────────────────

class DeviceStatus(str, Enum):
    """Lifecycle state of a device connection."""

    DISCONNECTED = "disconnected"
    CONNECTING = "connecting"
    CONNECTED = "connected"
    STREAMING = "streaming"
    RECONNECTING = "reconnecting"
    ERROR = "error"


class SessionStatus(str, Enum):
    """Lifecycle state of a recording session."""

    CREATED = "created"
    RECORDING = "recording"
    STOPPED = "stopped"
    PROCESSING = "processing"
    COMPLETED = "completed"
    FAILED = "failed"


class StreamMessageType(str, Enum):
    """Type of message received over a WebSocket stream."""

    EEG_DATA = "eeg_data"
    PROCESSED_DATA = "processed_data"
    FEATURES = "features"
    QUALITY = "quality"
    EVENT = "event"
    ERROR = "error"
    PING = "ping"
    PONG = "pong"


# ──────────────────────────────────────────────────────────────
# Device Models
# ──────────────────────────────────────────────────────────────

class Device(BaseModel):
    """
    Represents a neural data acquisition device.

    A device can be a physical EEG headset (Muse, OpenBCI, etc.)
    or a synthetic board used for development and testing.

    Attributes:
        device_id: Unique identifier assigned by the NeuroFlow server.
        name: Human-readable device name (e.g. "Muse S").
        manufacturer: Device manufacturer name.
        board_id: BrainFlow board identifier (-1 for synthetic).
        status: Current connection state.
        sampling_rate: Data sampling rate in Hz.
        num_channels: Number of EEG channels.
        connection_type: How the device connects (ble, serial, wifi, synthetic).
        is_streaming: Whether the device is currently streaming data.
        total_samples: Total number of samples received since connection.
        signal_quality: Overall signal quality score (0.0 to 1.0).
        uptime_seconds: Time since the device was connected.
    """

    device_id: str
    name: str = ""
    manufacturer: str = ""
    board_id: int = -1
    status: DeviceStatus = DeviceStatus.DISCONNECTED
    sampling_rate: int = 0
    num_channels: int = 0
    connection_type: str = ""
    is_streaming: bool = False
    total_samples: int = 0
    signal_quality: float = 0.0
    uptime_seconds: float = 0.0

    def __repr__(self) -> str:
        streaming = " [streaming]" if self.is_streaming else ""
        return (
            f"Device(id={self.device_id!r}, name={self.name!r}, "
            f"status={self.status.value}{streaming})"
        )


# ──────────────────────────────────────────────────────────────
# Session Models
# ──────────────────────────────────────────────────────────────

class Session(BaseModel):
    """
    Represents a recording session on a connected device.

    A session captures EEG data over a period of time, optionally
    processing it through a signal pipeline for feature extraction.

    Attributes:
        session_id: Unique session identifier.
        device_id: Identifier of the device being recorded.
        name: Human-readable session name.
        status: Current session lifecycle state.
        subject_id: Optional identifier for the test subject.
        pipeline_name: Name of the processing pipeline applied.
        started_at: When the session started recording.
        stopped_at: When the session stopped (None if still active).
        duration_seconds: Total duration of the session.
        total_samples: Total number of EEG samples captured.
        notes: Free-form notes about the session.
        tags: Categorization tags for the session.
    """

    session_id: str
    device_id: str = ""
    name: str = ""
    status: SessionStatus = SessionStatus.CREATED
    subject_id: str = ""
    pipeline_name: str = ""
    started_at: Optional[datetime] = None
    stopped_at: Optional[datetime] = None
    duration_seconds: float = 0.0
    total_samples: int = 0
    notes: str = ""
    tags: list[str] = Field(default_factory=list)

    @property
    def is_active(self) -> bool:
        """Whether the session is currently recording."""
        return self.status == SessionStatus.RECORDING

    def __repr__(self) -> str:
        return (
            f"Session(id={self.session_id!r}, device={self.device_id!r}, "
            f"status={self.status.value}, duration={self.duration_seconds:.1f}s)"
        )


# ──────────────────────────────────────────────────────────────
# Stream Data Models
# ──────────────────────────────────────────────────────────────

class StreamData(BaseModel):
    """
    A single chunk of streaming data received from a device.

    This represents one window of real-time data, typically
    containing 0.1-1.0 seconds of EEG samples.

    Attributes:
        timestamp: Unix timestamp when this data was captured.
        message_type: Type of stream message (eeg_data, features, quality, etc.).
        eeg: Raw EEG data as a list of channels, each a list of float samples.
        channel_names: Names of the EEG channels.
        num_samples: Number of samples in this data chunk.
        sampling_rate: Sampling rate of the data in Hz.
        features: Extracted features (if pipeline is active).
        quality: Per-channel signal quality scores.
        overall_quality: Average quality score across all channels.
        metadata: Additional message-specific data.
    """

    timestamp: float = 0.0
    message_type: StreamMessageType = StreamMessageType.EEG_DATA
    eeg: list[list[float]] = Field(default_factory=list)
    channel_names: list[str] = Field(default_factory=list)
    num_samples: int = 0
    sampling_rate: int = 0
    features: dict[str, Any] = Field(default_factory=dict)
    quality: dict[str, float] = Field(default_factory=dict)
    overall_quality: float = 0.0
    metadata: dict[str, Any] = Field(default_factory=dict)

    def __repr__(self) -> str:
        return (
            f"StreamData(type={self.message_type.value}, "
            f"samples={self.num_samples}, channels={len(self.eeg)})"
        )


# ──────────────────────────────────────────────────────────────
# Features Model
# ──────────────────────────────────────────────────────────────

class Features(BaseModel):
    """
    Extracted neuroscience features from EEG data.

    Features are computed by NeuroFlow's processing pipelines and
    include band powers (delta, theta, alpha, beta, gamma), spectral
    metrics, temporal statistics, and derived cognitive indicators.

    Attributes:
        session_id: Session the features were extracted from.
        timestamp: When the features were computed.
        band_powers: Power in standard EEG frequency bands (keyed by band name).
        spectral: Spectral features (dominant frequency, entropy, etc.).
        temporal: Temporal features (Hjorth parameters, zero crossings, etc.).
        ratios: Computed ratios (theta/beta for attention, alpha/theta for relaxation).
        raw: The complete raw feature dictionary as returned by the pipeline.
    """

    session_id: str = ""
    timestamp: float = 0.0
    band_powers: dict[str, float] = Field(default_factory=dict)
    spectral: dict[str, float] = Field(default_factory=dict)
    temporal: dict[str, float] = Field(default_factory=dict)
    ratios: dict[str, float] = Field(default_factory=dict)
    raw: dict[str, Any] = Field(default_factory=dict)

    @classmethod
    def from_raw_features(cls, raw_features: dict[str, Any], session_id: str = "") -> "Features":
        """
        Parse a raw feature dictionary into categorized feature groups.

        This method takes the flat dictionary returned by NeuroFlow's
        feature extractors and organizes it into band_powers, spectral,
        temporal, and ratio categories.

        Args:
            raw_features: Flat dictionary of feature_name -> value.
            session_id: Optional session identifier.

        Returns:
            A Features instance with categorized feature groups.
        """
        import time as _time

        band_powers: dict[str, float] = {}
        spectral: dict[str, float] = {}
        temporal: dict[str, float] = {}
        ratios: dict[str, float] = {}

        band_keywords = ("delta", "theta", "alpha", "beta", "gamma", "power")
        spectral_keywords = ("freq", "entropy", "spread", "edge", "psd")
        temporal_keywords = ("mean", "std", "rms", "ptp", "hjorth", "zero_crossing", "line_length", "skewness", "kurtosis")
        ratio_keywords = ("ratio",)

        for key, value in raw_features.items():
            if not isinstance(value, (int, float)):
                continue
            key_lower = key.lower()
            if any(kw in key_lower for kw in ratio_keywords):
                ratios[key] = float(value)
            elif any(kw in key_lower for kw in band_keywords):
                band_powers[key] = float(value)
            elif any(kw in key_lower for kw in spectral_keywords):
                spectral[key] = float(value)
            elif any(kw in key_lower for kw in temporal_keywords):
                temporal[key] = float(value)
            else:
                # Uncategorized features go into band_powers as a catch-all
                band_powers[key] = float(value)

        return cls(
            session_id=session_id,
            timestamp=_time.time(),
            band_powers=band_powers,
            spectral=spectral,
            temporal=temporal,
            ratios=ratios,
            raw=raw_features,
        )

    def __repr__(self) -> str:
        total = len(self.band_powers) + len(self.spectral) + len(self.temporal) + len(self.ratios)
        return f"Features(total={total}, bands={len(self.band_powers)}, ratios={len(self.ratios)})"


# ──────────────────────────────────────────────────────────────
# Session Data Model
# ──────────────────────────────────────────────────────────────

class SessionData(BaseModel):
    """
    Complete data from a finished recording session.

    Returned when you request session data after stopping a session.

    Attributes:
        session_id: Session identifier.
        sampling_rate: Sampling rate of the recorded data.
        num_channels: Number of EEG channels.
        num_samples: Total samples recorded.
        channel_names: Names of each EEG channel.
        duration_seconds: Total recording duration.
        eeg_data: EEG data as [channels][samples].
        timestamps: Sample timestamps.
        features: Extracted features (empty if no pipeline was used).
        quality_scores: Per-channel quality scores.
    """

    session_id: str
    sampling_rate: int = 250
    num_channels: int = 0
    num_samples: int = 0
    channel_names: list[str] = Field(default_factory=list)
    duration_seconds: float = 0.0
    eeg_data: list[list[float]] = Field(default_factory=list)
    timestamps: list[float] = Field(default_factory=list)
    features: dict[str, Any] = Field(default_factory=dict)
    quality_scores: list[float] = Field(default_factory=list)

    def __repr__(self) -> str:
        return (
            f"SessionData(id={self.session_id!r}, "
            f"channels={self.num_channels}, samples={self.num_samples}, "
            f"duration={self.duration_seconds:.1f}s)"
        )


# ──────────────────────────────────────────────────────────────
# Pipeline Model
# ──────────────────────────────────────────────────────────────

class PipelineInfo(BaseModel):
    """
    Information about a signal processing pipeline.

    Attributes:
        name: Pipeline identifier (e.g. "meditation", "focus_detection").
        steps: Ordered list of processing step names.
        feature_extractors: List of feature extraction step names.
        is_preset: Whether this is a built-in preset pipeline.
        description: Human-readable pipeline description.
    """

    name: str
    steps: list[str] = Field(default_factory=list)
    feature_extractors: list[str] = Field(default_factory=list)
    is_preset: bool = False
    description: str = ""


# ──────────────────────────────────────────────────────────────
# System Status Model
# ──────────────────────────────────────────────────────────────

class SystemStatus(BaseModel):
    """
    NeuroFlow server system status.

    Attributes:
        version: Server version string.
        uptime_seconds: Time since server start.
        active_devices: Number of currently connected devices.
        active_sessions: Number of sessions currently recording.
        active_streams: Number of active WebSocket streams.
        total_sessions: Total sessions recorded since server start.
        total_samples_processed: Total samples processed through pipelines.
        registered_pipelines: List of available pipeline names.
        available_presets: List of built-in preset pipeline names.
    """

    version: str = ""
    uptime_seconds: float = 0.0
    active_devices: int = 0
    active_sessions: int = 0
    active_streams: int = 0
    total_sessions: int = 0
    total_samples_processed: int = 0
    registered_pipelines: list[str] = Field(default_factory=list)
    available_presets: list[str] = Field(default_factory=list)
