"""Signal selection, signal data array models, and signal enums."""

from __future__ import annotations

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

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

if TYPE_CHECKING:
    from radiens_core.models.channels import ChannelMetadata


class KeyIndex(IntEnum):
    """Channel key index type.

    NTV: native position in DAQ stream per signal type (starts at 0).
    DATA: position in backing dataset per signal type.
    SYS: global position across all signal types.
    """

    NTV = 0
    DATA = 1
    SYS = 2


class SignalType(IntEnum):
    """Signal type. Values match proto ``common.SignalType``."""

    AMP = 0  # PRI in proto
    AIN = 1  # AUX in proto
    DIN = 2
    DOUT = 3


class SigSelect(BaseModel):
    """Signal selection specifying channels per signal type.

    Attributes:
        amp: Array of native channel indices for amplifier signals.
        gpio_ain: Array of native channel indices for analog input signals.
        gpio_din: Array of native channel indices for digital input signals.
        gpio_dout: Array of native channel indices for digital output signals.
        key_idx: Channel key index type. References KeyIndex.
    """

    model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True)

    amp: npt.NDArray[np.int64]
    gpio_ain: npt.NDArray[np.int64]
    gpio_din: npt.NDArray[np.int64]
    gpio_dout: npt.NDArray[np.int64]
    key_idx: KeyIndex


class SignalArrays(BaseModel):
    """Container for signal data arrays, one per signal type.

    Attributes:
        amp: 2D array (channels, samples) for amplifier data. Defaults to None.
        gpio_ain: 2D array (channels, samples) for analog input data. Defaults to None.
        gpio_din: 2D array (channels, samples) for digital input data. Defaults to None.
        gpio_dout: 2D array (channels, samples) for digital output data. Defaults to None.
    """

    model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True)

    amp: npt.NDArray[np.float64] | None = None
    gpio_ain: npt.NDArray[np.float64] | None = None
    gpio_din: npt.NDArray[np.float64] | None = None
    gpio_dout: npt.NDArray[np.float64] | None = None


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")

    Attributes:
        sig_type: Primary signal type. References SignalType.
        channels: Tuple of channel indices or selection string ("all", "selected").
        multi_select: Map of signal types to channel selections. Defaults to None.
    """

    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}. 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})"


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

    Accepted inputs:
    - ``SignalSpec`` — pass-through
    - ``"all"`` — all amplifier channels (shorthand for ``SignalSpec.all_amp()``)
    - ``list[int]`` — interpreted as amplifier channel indices
    - ``dict`` — passed as kwargs to ``SignalSpec.multi()``
    """
    if isinstance(v, SignalSpec):
        return v
    if isinstance(v, str):
        if v == "all":
            return SignalSpec.all_amp()
        raise ValueError(f"Unknown signal spec string {v!r} — only 'all' is supported")
    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__}")


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

Use in Pydantic models::

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