"""DatasetMetadata — immutable metadata for a linked dataset."""

from __future__ import annotations

import uuid
from collections.abc import Sequence
from enum import IntEnum
from pathlib import Path
from typing import Annotated

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

from radiens_core.models.channels import ChannelMetadata
from radiens_core.models.time_range import TimeRange


class TrsMode(IntEnum):
    """Time range selection mode. Values match proto ``common.TimeRangeSelMode``."""

    SUBSET = 0  # TRS_SUBSET
    TO_HEAD = 1  # TRS_TO_HEAD
    FROM_HEAD = 2  # TRS_FROM_HEAD


class RadiensFileType(IntEnum):
    """Recording file type. Values match proto ``common.RadixFileTypes``."""

    RHD = 0
    XDAT = 1
    CSV = 2
    HDF5 = 3
    NEX5 = 4
    NWB = 5
    BIOINTERFACE = 7
    KILOSORT2 = 8
    ANC = 9
    SPKTRAIN_MU = 10
    BIO_RADIENS_EXPORT_HDF5 = 11
    BIO_RADIENS_EXPORT_MAT5 = 12
    NSX = 13
    PL2 = 14
    TDT = 15
    SPIKES = 16
    KPI = 17
    UNKNOWN_FILE_TYPE = 18
    MEAREC_HDF5 = 19
    OFFLINE_SORTER_SPKS_MAT5 = 20
    SPKS_NPY = 21
    OE_DAT = 22
    SPKS_PHY2 = 23


class ProtocolSpec(BaseModel):
    """Minimal summary of a saved protocol.

    Returned by ``get_protocol`` and ``get_all_protocols``.

    Attributes:
        id: Unique identifier for the protocol.
        node_count: Number of nodes defined in the protocol.
    """

    model_config = ConfigDict(frozen=True)

    id: str
    node_count: int


class FileInfo(BaseModel):
    """Filesystem entry metadata returned by ``list_directory``.

    Attributes:
        path: Full filesystem path to the file.
        base_name: Base name of the file without directory components.
        file_type: Type of the recording file. References RadiensFileType.
        dataset_uid: Unique identifier for the dataset.
        num_channels: Number of channels in the file.
        duration_sec: Total duration of the recording in seconds.
        sample_rate: Sampling rate in Hertz.
        num_bytes: Size of the file in bytes.
    """

    model_config = ConfigDict(frozen=True)

    path: str
    base_name: str
    file_type: RadiensFileType
    dataset_uid: str
    num_channels: int
    duration_sec: float
    sample_rate: int
    num_bytes: int


class DatasetMetadata(BaseModel):
    """Metadata for a linked dataset.

    Returned by ``link_data_file()`` and similar operations.
    Contains everything needed to query signals from this dataset.

    Attributes:
        id: Unique identifier for the dataset.
        base_name: Base name of the dataset.
        path: Full filesystem path to the dataset.
        file_type: Type of the recording file. References RadiensFileType.
        time_range: Valid time range for this dataset. References TimeRange.
        channel_metadata: Channel configuration and metadata. References ChannelMetadata.
        parent_dsource_id: Identifier for the parent data source. Defaults to "".
        probe_uid: Unique identifier for the probe used. Defaults to "".
        associated_spikes_ids: Registered spikes datasource IDs discovered when the
            recording file was linked (populated by ``link_data_file()``). Empty if no
            ``.spikes`` file was found alongside the recording.
    """

    model_config = ConfigDict(frozen=True)

    id: str
    base_name: str
    path: str
    file_type: RadiensFileType
    time_range: TimeRange
    channel_metadata: ChannelMetadata
    parent_dsource_id: str = ""
    probe_uid: str = ""
    associated_spikes_ids: list[str] = []

    @property
    def duration(self) -> float:
        """Dataset duration in seconds."""
        return self.time_range.dur_sec

    @property
    def sampling_rate(self) -> float:
        """Sampling rate in Hz."""
        return self.time_range.fs


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

    Attributes:
        mode: Time range selection mode. References TrsMode.
        start: Start time or duration depending on the mode.
        end: End time for subset mode. Defaults to None.
    """

    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.

        Args:
        dataset_metadata
            Dataset metadata used to obtain ``fs`` and dataset bounds.
        fs
            Sampling rate override (used when ``dataset_metadata`` is not provided).

        Returns:
        TimeRange or None
            ``None`` for ``FROM_HEAD`` mode — caller must use ``to_array()`` instead.
        """
        if self.mode == TrsMode.FROM_HEAD:
            return None

        if not isinstance(dataset_metadata, DatasetMetadata):
            raise ValueError("dataset_metadata required to resolve time range")

        dataset_tr = dataset_metadata.time_range
        effective_fs = fs if fs is not None else dataset_tr.fs

        if self.is_full():
            return dataset_tr

        if self.mode == TrsMode.SUBSET:
            end = self.end if (self.end is not None and self.end != float("inf")) else dataset_tr.sec[1]
            ts_start = int(self.start * effective_fs)
            ts_end = int(end * effective_fs)
            return TimeRange(
                sec=(self.start, end),
                timestamp=(ts_start, ts_end),
                fs=effective_fs,
                dur_sec=end - self.start,
                walltime_str=dataset_tr.walltime_str,
                n=ts_end - ts_start,
            )

        if self.mode == TrsMode.TO_HEAD:
            ts_start = int(self.start * effective_fs)
            ts_end = dataset_tr.timestamp[1]
            end = dataset_tr.sec[1]
            return TimeRange(
                sec=(self.start, end),
                timestamp=(ts_start, ts_end),
                fs=effective_fs,
                dur_sec=end - self.start,
                walltime_str=dataset_tr.walltime_str,
                n=ts_end - ts_start,
            )

        raise ValueError(f"Cannot resolve time range with mode: {self.mode}")

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


def _coerce_time_range(v: object) -> TimeRangeSpec:
    """Coerce raw Python types into `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__}")


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


def _coerce_dataset(v: object) -> DatasetMetadata | str | Path:
    """Coerce raw Python types into a dataset reference.

    Accepted inputs:

    - ``DatasetMetadata`` — pass-through
    - ``str`` — dataset ID or file path string
    - ``Path`` — file path
    """
    if isinstance(v, (DatasetMetadata, str, Path)):
        return v
    raise ValueError(f"Expected DatasetMetadata, str, or Path — got {type(v).__name__}")


FlexDataset = Annotated[
    DatasetMetadata | str | Path,
    BeforeValidator(_coerce_dataset),
]
"""``DatasetMetadata`` that also accepts an ID string or file ``Path``.

Resolution from string/path to ``DatasetMetadata`` happens inside the client
via ``_resolve_dataset()``.  The ``BeforeValidator`` validates the input type
and provides clear errors for invalid arguments.

Use with ``@validate_call``::

    @validate_call(config=ConfigDict(arbitrary_types_allowed=True))
    def get_signals(self, dataset: FlexDataset, ...) -> ...:
        meta = self._resolve_dataset(dataset)
        ...
"""


# ===================================================================
# Protocol / Transform Graph Models
# ===================================================================


class TransformNodeType(IntEnum):
    """Transform node type for curation protocol graphs.

    Values match proto ``common.TransformNodeType``.
    """

    NOTCH_FILTER = 0
    BANDPASS_FILTER = 1
    BANDSTOP_FILTER = 2
    HIGHPASS_FILTER = 3
    LOWPASS_FILTER = 4
    PAIRED_REF = 5
    VIRTUAL_REF = 6
    CAR_REF = 7
    SLICE_TIME = 9
    SLICE_CHANNELS = 10
    SOURCE = 11
    SINK = 12
    DOWNSAMPLE = 13
    REMAP_SENSOR = 14
    BULK_SOURCE = 15
    BULK_SINK = 16


class HighLowPassParams(BaseModel):
    """Parameters for high-pass or low-pass filter transform nodes."""

    model_config = ConfigDict(frozen=True)

    frequency: float
    """Cutoff frequency in Hz."""
    order: int | None = None
    """Filter order; ``None`` uses the server default."""


class BandParams(BaseModel):
    """Parameters for band-pass or band-stop filter transform nodes."""

    model_config = ConfigDict(frozen=True)

    low_frequency: float
    """Lower cutoff frequency in Hz."""
    high_frequency: float
    """Upper cutoff frequency in Hz."""
    order: int | None = None
    """Filter order; ``None`` uses the server default."""


class NotchParams(BaseModel):
    """Parameters for notch filter transform nodes."""

    model_config = ConfigDict(frozen=True)

    notch_frequency: float
    """Center frequency to reject in Hz."""
    bandwidth: float
    """Width of the rejection band in Hz."""
    order: int | None = None
    """Filter order; ``None`` uses the server default."""


class CARParams(BaseModel):
    """Parameters for Common Average Reference transform nodes (no fields)."""

    model_config = ConfigDict(frozen=True)


class VirtualRefParams(BaseModel):
    """Parameters for virtual reference transform nodes."""

    model_config = ConfigDict(frozen=True)

    ref_channel: int
    """Native channel index of the reference electrode."""


class PairedRefParams(BaseModel):
    """Parameters for paired reference transform nodes."""

    model_config = ConfigDict(frozen=True)

    ref_channel: int
    """Native channel index of the reference electrode."""
    target_channel: int
    """Native channel index of the channel to re-reference."""


class SliceTimeParams(BaseModel):
    """Parameters for time-slice transform nodes."""

    model_config = ConfigDict(frozen=True)

    start_sec: float
    """Start of the time window in seconds."""
    end_sec: float
    """End of the time window in seconds."""


class SliceChannelsParams(BaseModel):
    """Parameters for channel-slice transform nodes."""

    model_config = ConfigDict(frozen=True)

    channels: list[int]
    """System-wide channel indices to keep."""


class DownsampleParams(BaseModel):
    """Parameters for downsample transform nodes."""

    model_config = ConfigDict(frozen=True)

    factor: int
    """Downsampling factor (e.g., 2, 4, 10)."""


class SourceParams(BaseModel):
    """Parameters for SOURCE transform nodes."""

    model_config = ConfigDict(frozen=True)

    dataset_id: str = ""
    """ID of the linked dataset (empty for new sources)."""
    base_name: str
    """Base filename (without extension or directory)."""
    path: str
    """Directory path containing the dataset."""
    file_type: RadiensFileType
    """Recording file format."""


class SinkParams(BaseModel):
    """Parameters for SINK transform nodes."""

    model_config = ConfigDict(frozen=True)

    base_name: str
    """Output base filename (without extension or directory)."""
    path: str
    """Output directory path."""
    file_type: RadiensFileType
    """Output recording file format."""


class BulkSourceParams(BaseModel):
    """Parameters for BULK_SOURCE transform nodes."""

    model_config = ConfigDict(frozen=True)

    sources: list[SourceParams]
    """One entry per input dataset."""


class BulkSinkParams(BaseModel):
    """Parameters for BULK_SINK transform nodes."""

    model_config = ConfigDict(frozen=True)

    suffix: str
    """Suffix appended to each output file's base name."""
    path: str
    """Output directory path."""
    file_type: RadiensFileType
    """Output recording file format."""


class TransformEdge(BaseModel):
    """A directed edge connecting two nodes in a protocol graph.

    Use the ``source`` and ``target`` fields to specify connectivity
    by node ID.
    """

    model_config = ConfigDict(frozen=True)

    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
    """Edge ID (auto-generated UUID if not provided)."""
    source: str
    """ID of the source transform node."""
    target: str
    """ID of the target transform node."""


type TransformParams = (
    HighLowPassParams
    | BandParams
    | NotchParams
    | CARParams
    | VirtualRefParams
    | PairedRefParams
    | SliceTimeParams
    | SliceChannelsParams
    | DownsampleParams
    | SourceParams
    | SinkParams
    | BulkSourceParams
    | BulkSinkParams
    | None
)
"""Union of all transform parameter types (or ``None`` for parameter-free nodes)."""


class TransformNode(BaseModel):
    """A node in a curation protocol transform graph.

    Construct nodes using the factory class methods for convenience::

        src = TransformNode.source(dataset_meta)
        hp  = TransformNode.highpass(frequency=300.0)
        snk = TransformNode.sink("filtered", "/output", RadiensFileType.XDAT)

        edges = [
            TransformEdge(source=src.id, target=hp.id),
            TransformEdge(source=hp.id, target=snk.id),
        ]

    Then pass ``[src, hp, snk]`` and ``edges`` to
    `CurateClient.set_protocol()`.

    Notes
    -----
    Each factory method sets the correct ``type`` and ``params`` for its
    transform. When constructing manually, the ``params`` object must match
    the ``type`` (e.g., ``CAR_REF`` requires ``params=CARParams()``).
    """

    model_config = ConfigDict(frozen=True)

    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
    """Node ID (auto-generated UUID if not provided)."""
    type: TransformNodeType
    """Transform operation this node performs."""
    params: TransformParams = None
    """Operation parameters; see each factory method for the expected type."""

    # --- Filter factory methods ---

    @classmethod
    def highpass(cls, frequency: float, order: int | None = None) -> TransformNode:
        """Create a ``HIGHPASS_FILTER`` node.

        Args:
        frequency:
            Cutoff frequency in Hz.
        order:
            Filter order; ``None`` uses the server default.
        """
        return cls(
            type=TransformNodeType.HIGHPASS_FILTER,
            params=HighLowPassParams(frequency=frequency, order=order),
        )

    @classmethod
    def lowpass(cls, frequency: float, order: int | None = None) -> TransformNode:
        """Create a ``LOWPASS_FILTER`` node.

        Args:
        frequency:
            Cutoff frequency in Hz.
        order:
            Filter order; ``None`` uses the server default.
        """
        return cls(
            type=TransformNodeType.LOWPASS_FILTER,
            params=HighLowPassParams(frequency=frequency, order=order),
        )

    @classmethod
    def bandpass(
        cls,
        low_frequency: float,
        high_frequency: float,
        order: int | None = None,
    ) -> TransformNode:
        """Create a ``BANDPASS_FILTER`` node.

        Args:
        low_frequency:
            Lower cutoff frequency in Hz.
        high_frequency:
            Upper cutoff frequency in Hz.
        order:
            Filter order; ``None`` uses the server default.
        """
        return cls(
            type=TransformNodeType.BANDPASS_FILTER,
            params=BandParams(
                low_frequency=low_frequency,
                high_frequency=high_frequency,
                order=order,
            ),
        )

    @classmethod
    def bandstop(
        cls,
        low_frequency: float,
        high_frequency: float,
        order: int | None = None,
    ) -> TransformNode:
        """Create a ``BANDSTOP_FILTER`` node.

        Args:
        low_frequency:
            Lower cutoff frequency in Hz.
        high_frequency:
            Upper cutoff frequency in Hz.
        order:
            Filter order; ``None`` uses the server default.
        """
        return cls(
            type=TransformNodeType.BANDSTOP_FILTER,
            params=BandParams(
                low_frequency=low_frequency,
                high_frequency=high_frequency,
                order=order,
            ),
        )

    @classmethod
    def notch(
        cls,
        notch_frequency: float,
        bandwidth: float,
        order: int | None = None,
    ) -> TransformNode:
        """Create a ``NOTCH_FILTER`` node.

        Args:
        notch_frequency:
            Center frequency to reject in Hz.
        bandwidth:
            Width of the rejection band in Hz.
        order:
            Filter order; ``None`` uses the server default.
        """
        return cls(
            type=TransformNodeType.NOTCH_FILTER,
            params=NotchParams(
                notch_frequency=notch_frequency,
                bandwidth=bandwidth,
                order=order,
            ),
        )

    # --- Re-referencing factory methods ---

    @classmethod
    def car(cls) -> TransformNode:
        """Create a ``CAR_REF`` (Common Average Reference) node."""
        return cls(type=TransformNodeType.CAR_REF, params=CARParams())

    @classmethod
    def virtual_ref(cls, ref_channel: int) -> TransformNode:
        """Create a ``VIRTUAL_REF`` node.

        Args:
        ref_channel:
            Native channel index of the reference electrode.
        """
        return cls(
            type=TransformNodeType.VIRTUAL_REF,
            params=VirtualRefParams(ref_channel=ref_channel),
        )

    @classmethod
    def paired_ref(cls, ref_channel: int, target_channel: int) -> TransformNode:
        """Create a ``PAIRED_REF`` node.

        Args:
        ref_channel:
            Native channel index of the reference electrode.
        target_channel:
            Native channel index of the channel to re-reference.
        """
        return cls(
            type=TransformNodeType.PAIRED_REF,
            params=PairedRefParams(
                ref_channel=ref_channel,
                target_channel=target_channel,
            ),
        )

    # --- Slicing / resampling factory methods ---

    @classmethod
    def slice_time(cls, start_sec: float, end_sec: float) -> TransformNode:
        """Create a ``SLICE_TIME`` node.

        Args:
        start_sec:
            Start of the time window in seconds.
        end_sec:
            End of the time window in seconds.
        """
        return cls(
            type=TransformNodeType.SLICE_TIME,
            params=SliceTimeParams(start_sec=start_sec, end_sec=end_sec),
        )

    @classmethod
    def slice_channels(cls, channels: list[int]) -> TransformNode:
        """Create a ``SLICE_CHANNELS`` node.

        Args:
        channels:
            System-wide channel indices to keep.
        """
        return cls(
            type=TransformNodeType.SLICE_CHANNELS,
            params=SliceChannelsParams(channels=channels),
        )

    @classmethod
    def downsample(cls, factor: int) -> TransformNode:
        """Create a ``DOWNSAMPLE`` node.

        Args:
        factor:
            Downsampling factor (e.g., 2, 4, 10).
        """
        return cls(
            type=TransformNodeType.DOWNSAMPLE,
            params=DownsampleParams(factor=factor),
        )

    # --- Source / sink factory methods ---

    @classmethod
    def source(cls, dataset: DatasetMetadata) -> TransformNode:
        """Create a ``SOURCE`` node from a linked dataset.

        Args:
        dataset:
            Metadata for the input dataset (returned by ``link_data_file``).
        """
        return cls(
            type=TransformNodeType.SOURCE,
            params=SourceParams(
                dataset_id=dataset.id,
                base_name=dataset.base_name,
                path=dataset.path,
                file_type=dataset.file_type,
            ),
        )

    @classmethod
    def sink(cls, base_name: str, path: str, file_type: RadiensFileType) -> TransformNode:
        """Create a ``SINK`` node.

        Args:
        base_name:
            Output base filename (without extension or directory).
        path:
            Output directory path.
        file_type:
            Output recording file format.
        """
        return cls(
            type=TransformNodeType.SINK,
            params=SinkParams(base_name=base_name, path=path, file_type=file_type),
        )

    @classmethod
    def bulk_source(cls, datasets: list[DatasetMetadata]) -> TransformNode:
        """Create a ``BULK_SOURCE`` node from multiple linked datasets.

        Args:
        datasets:
            Metadata for all input datasets.
        """
        return cls(
            type=TransformNodeType.BULK_SOURCE,
            params=BulkSourceParams(
                sources=[
                    SourceParams(
                        dataset_id=d.id,
                        base_name=d.base_name,
                        path=d.path,
                        file_type=d.file_type,
                    )
                    for d in datasets
                ]
            ),
        )

    @classmethod
    def bulk_sink(cls, suffix: str, path: str, file_type: RadiensFileType) -> TransformNode:
        """Create a ``BULK_SINK`` node.

        Args:
        suffix:
            Suffix appended to each output file's base name.
        path:
            Output directory path.
        file_type:
            Output recording file format.
        """
        return cls(
            type=TransformNodeType.BULK_SINK,
            params=BulkSinkParams(suffix=suffix, path=path, file_type=file_type),
        )
