"""User-facing specification types — TimeRangeSpec, SignalSpec, Flex aliases.

These provide a clean, semantic interface for specifying time ranges
and signal selections across all client types.

Flex type aliases (``FlexTimeRange``, ``FlexSignalSpec``) use Pydantic
``BeforeValidator`` to coerce raw Python types (list, tuple, dict) into
canonical specs at the Pydantic model boundary.
"""

from __future__ import annotations

from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Annotated

import numpy as np
import numpy.typing as npt
from pydantic import BaseModel, BeforeValidator, ConfigDict

from radiens_core.models.enums import KeyIndex, SignalType, TrsMode
from radiens_core.models.signals import SigSelect

if TYPE_CHECKING:
    from radiens_core.models.channels import ChannelMetadata
    from radiens_core.models.time_range import TimeRange


class TimeRangeSpec(BaseModel):
    """Semantic time range specification with embedded mode.

    Use factory methods to create specs with clear semantics:

    >>> spec = TimeRangeSpec.subset(0, 10)     # [0, 10) seconds
    >>> spec = TimeRangeSpec.lookback(5.0)      # Last 5 seconds
    >>> spec = TimeRangeSpec.full()              # Entire recording
    """

    model_config = ConfigDict(frozen=True)

    mode: TrsMode
    start: float
    end: float | None = None

    # --- Factory methods ---

    @classmethod
    def subset(cls, start: float, end: float) -> TimeRangeSpec:
        """Absolute time range ``[start, end)`` seconds."""
        if end <= start:
            raise ValueError(f"end ({end}) must be greater than start ({start})")
        return cls(mode=TrsMode.SUBSET, start=start, end=end)

    @classmethod
    def lookback(cls, duration: float) -> TimeRangeSpec:
        """Last ``duration`` seconds from head of cache."""
        if duration <= 0:
            raise ValueError(f"duration must be positive, got {duration}")
        return cls(mode=TrsMode.FROM_HEAD, start=duration, end=None)

    @classmethod
    def to_head(cls, start: float) -> TimeRangeSpec:
        """From ``start`` seconds to current head."""
        return cls(mode=TrsMode.TO_HEAD, start=start, end=None)

    @classmethod
    def full(cls) -> TimeRangeSpec:
        """Full dataset/cache range."""
        return cls(mode=TrsMode.SUBSET, start=0, end=float("inf"))

    @classmethod
    def from_list(cls, time_range: list[float], mode: TrsMode | None = None) -> TimeRangeSpec:
        """Create from ``[start, end]`` or ``[duration, nan]`` list."""
        arr = np.asarray(time_range, dtype=np.float64)
        if len(arr) != 2:
            raise ValueError("time_range must have exactly 2 elements")
        if mode is None:
            mode = TrsMode.FROM_HEAD if np.isnan(arr[1]) else TrsMode.SUBSET
        if mode == TrsMode.FROM_HEAD:
            return cls(mode=mode, start=float(arr[0]), end=None)
        return cls(mode=mode, start=float(arr[0]), end=float(arr[1]))

    # --- Resolution methods ---

    def is_full(self) -> bool:
        """Check if this is a full-range specification."""
        return self.mode == TrsMode.SUBSET and self.end == float("inf")

    def to_array(self) -> npt.NDArray[np.float64]:
        """Convert to ``[start, end]`` or ``[duration, nan]`` array."""
        if self.is_full():
            raise ValueError("Cannot convert full() spec to array without dataset.")
        if self.mode in (TrsMode.FROM_HEAD, TrsMode.TO_HEAD):
            return np.array([self.start, np.nan], dtype=np.float64)
        return np.array([self.start, self.end], dtype=np.float64)

    def resolve(self, dataset_metadata: object | None = None, fs: float | None = None) -> TimeRange | None:
        """Resolve to a concrete TimeRange.

        Parameters
        ----------
        dataset_metadata
            Dataset metadata (must have ``.time_range`` and ``.make_time_range()``).
        fs
            Sampling rate override.

        Returns
        -------
        TimeRange or None
            None for FROM_HEAD mode (caller should use ``to_array()``).
        """
        from radiens_core.models.dataset import DatasetMetadata

        if self.is_full():
            if not isinstance(dataset_metadata, DatasetMetadata):
                raise ValueError("dataset_metadata required to resolve full() specification")
            return dataset_metadata.time_range

        if self.mode == TrsMode.FROM_HEAD:
            return None

        if isinstance(dataset_metadata, DatasetMetadata):
            return dataset_metadata.time_range

        raise ValueError("dataset_metadata required to fully resolve TimeRange")

    @property
    def duration(self) -> float | None:
        """Duration in seconds, if known."""
        if self.mode == TrsMode.FROM_HEAD:
            return self.start
        if self.mode == TrsMode.SUBSET and self.end is not None and self.end != float("inf"):
            return self.end - self.start
        return None

    def __repr__(self) -> str:
        if self.is_full():
            return "TimeRangeSpec.full()"
        if self.mode == TrsMode.FROM_HEAD:
            return f"TimeRangeSpec.lookback({self.start})"
        if self.mode == TrsMode.TO_HEAD:
            return f"TimeRangeSpec.to_head({self.start})"
        return f"TimeRangeSpec.subset({self.start}, {self.end})"


class SignalSpec(BaseModel):
    """Semantic signal selection specification.

    Use factory methods for clean channel selection:

    >>> spec = SignalSpec.amp([0, 1, 2, 3])
    >>> spec = SignalSpec.ain("all")
    >>> spec = SignalSpec.multi(amp=[0, 1], din="all")
    """

    model_config = ConfigDict(frozen=True)

    sig_type: SignalType | None
    channels: tuple[int, ...] | str
    multi_select: dict[SignalType, tuple[int, ...] | str] | None = None

    # --- Single-type factory methods ---

    @classmethod
    def amp(cls, channels: list[int] | str = "all") -> SignalSpec:
        """Select amplifier channels."""
        return cls._single(SignalType.AMP, channels)

    @classmethod
    def ain(cls, channels: list[int] | str = "all") -> SignalSpec:
        """Select analog input channels."""
        return cls._single(SignalType.AIN, channels)

    @classmethod
    def din(cls, channels: list[int] | str = "all") -> SignalSpec:
        """Select digital input channels."""
        return cls._single(SignalType.DIN, channels)

    @classmethod
    def dout(cls, channels: list[int] | str = "all") -> SignalSpec:
        """Select digital output channels."""
        return cls._single(SignalType.DOUT, channels)

    @classmethod
    def all_amp(cls) -> SignalSpec:
        """Convenience: all amplifier channels."""
        return cls.amp("all")

    @classmethod
    def _single(cls, sig_type: SignalType, channels: list[int] | str) -> SignalSpec:
        ch: tuple[int, ...] | str = tuple(channels) if isinstance(channels, list) else channels
        return cls(sig_type=sig_type, channels=ch, multi_select=None)

    # --- Multi-type factory method ---

    @classmethod
    def multi(
        cls,
        amp: list[int] | str | None = None,
        ain: list[int] | str | None = None,
        din: list[int] | str | None = None,
        dout: list[int] | str | None = None,
    ) -> SignalSpec:
        """Select channels from multiple signal types."""
        multi_dict: dict[SignalType, tuple[int, ...] | str] = {}
        for sig_type, ch in [
            (SignalType.AMP, amp),
            (SignalType.AIN, ain),
            (SignalType.DIN, din),
            (SignalType.DOUT, dout),
        ]:
            if ch is not None:
                multi_dict[sig_type] = tuple(ch) if isinstance(ch, list) else ch
        if not multi_dict:
            raise ValueError("At least one signal type must be specified")
        return cls(sig_type=None, channels="multi", multi_select=multi_dict)

    # --- Resolution ---

    def is_multi(self) -> bool:
        """Check if this is a multi-type specification."""
        return self.multi_select is not None

    def to_sig_select(self, channel_metadata: ChannelMetadata) -> SigSelect:
        """Convert to SigSelect for API calls."""
        if self.is_multi():
            return self._resolve_multi(channel_metadata)
        return self._resolve_single(channel_metadata)

    def _resolve_single(self, channel_metadata: ChannelMetadata) -> SigSelect:
        arrays: dict[SignalType, npt.NDArray[np.int64]] = {
            st: np.array([], dtype=np.int64) for st in SignalType
        }
        if self.sig_type is not None:
            arrays[self.sig_type] = self._resolve_channels(
                self.sig_type, self.channels, channel_metadata
            )
        return SigSelect(
            key_idx=KeyIndex.NTV,
            amp=arrays[SignalType.AMP],
            gpio_ain=arrays[SignalType.AIN],
            gpio_din=arrays[SignalType.DIN],
            gpio_dout=arrays[SignalType.DOUT],
        )

    def _resolve_multi(self, channel_metadata: ChannelMetadata) -> SigSelect:
        arrays: dict[SignalType, npt.NDArray[np.int64]] = {
            st: np.array([], dtype=np.int64) for st in SignalType
        }
        if self.multi_select is not None:
            for sig_type, channels in self.multi_select.items():
                arrays[sig_type] = self._resolve_channels(sig_type, channels, channel_metadata)
        return SigSelect(
            key_idx=KeyIndex.NTV,
            amp=arrays[SignalType.AMP],
            gpio_ain=arrays[SignalType.AIN],
            gpio_din=arrays[SignalType.DIN],
            gpio_dout=arrays[SignalType.DOUT],
        )

    @staticmethod
    def _resolve_channels(
        sig_type: SignalType,
        channels: tuple[int, ...] | str,
        channel_metadata: ChannelMetadata,
    ) -> npt.NDArray[np.int64]:
        available = channel_metadata.index(sig_type).ntv
        if channels == "all":
            return np.array(available, dtype=np.int64)
        if channels == "selected":
            selected = channel_metadata.selected_index(sig_type).ntv
            return np.array(selected, dtype=np.int64)
        if isinstance(channels, tuple):
            invalid = set(channels) - set(available)
            if invalid:
                raise ValueError(
                    f"Invalid {sig_type.name.lower()} channel indices: {invalid}. "
                    f"Available: {available}"
                )
            return np.array(channels, dtype=np.int64)
        raise ValueError(f"Invalid channel specification: {channels}")

    def __repr__(self) -> str:
        if self.is_multi() and self.multi_select is not None:
            parts: list[str] = []
            for sig_type, channels in self.multi_select.items():
                ch_display = list(channels) if isinstance(channels, tuple) else repr(channels)
                parts.append(f"{sig_type.name.lower()}={ch_display}")
            return f"SignalSpec.multi({', '.join(parts)})"
        ch_display_single = list(self.channels) if isinstance(self.channels, tuple) else repr(self.channels)
        name = self.sig_type.name.lower() if self.sig_type is not None else "unknown"
        return f"SignalSpec.{name}({ch_display_single})"


# ===================================================================
# Flex type aliases — Pydantic BeforeValidator coercions
# ===================================================================


def _coerce_time_range(v: object) -> TimeRangeSpec:
    """Coerce raw Python types into :class:`TimeRangeSpec`.

    Accepted inputs:
    - ``TimeRangeSpec`` — pass-through
    - ``list[float]`` / ``tuple[float, ...]`` — delegates to ``TimeRangeSpec.from_list``
    """
    if isinstance(v, TimeRangeSpec):
        return v
    if isinstance(v, (list, tuple)):
        return TimeRangeSpec.from_list(list(v))
    raise ValueError(f"Expected TimeRangeSpec, list, or tuple — got {type(v).__name__}")


def _coerce_signal_spec(v: object) -> SignalSpec:
    """Coerce raw Python types into :class:`SignalSpec`.

    Accepted inputs:
    - ``SignalSpec`` — pass-through
    - ``list[int]`` — interpreted as amplifier channel indices
    - ``dict`` — passed as kwargs to ``SignalSpec.multi()``
    """
    if isinstance(v, SignalSpec):
        return v
    if isinstance(v, list):
        return SignalSpec.amp(v)
    if isinstance(v, dict):
        return SignalSpec.multi(**v)
    raise ValueError(f"Expected SignalSpec, list[int], or dict — got {type(v).__name__}")


FlexTimeRange = Annotated[
    TimeRangeSpec | Sequence[float],
    BeforeValidator(_coerce_time_range),
]
"""``TimeRangeSpec`` that also accepts ``list``/``tuple`` shorthand.

Use in Pydantic models::

    class MyRequest(BaseModel):
        time_range: FlexTimeRange
"""

FlexSignalSpec = Annotated[
    SignalSpec | Sequence[int] | Mapping[str, Sequence[int] | str],
    BeforeValidator(_coerce_signal_spec),
]
"""``SignalSpec`` that also accepts ``list[int]`` or ``dict`` shorthand.

Use in Pydantic models::

    class MyRequest(BaseModel):
        signals: FlexSignalSpec
"""
