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

from __future__ import annotations

from collections.abc import Sequence
from enum import IntEnum
from typing import Annotated

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

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
    KILOSORT2 = 8
    NSX = 13
    TDT = 15
    SPIKES = 16
    OE_DAT = 22


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

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

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

    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 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 :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__}")


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