"""VidereClient — offline file-based analysis client."""

from __future__ import annotations

import contextlib
from typing import TYPE_CHECKING, ClassVar, cast

from pydantic import ConfigDict, validate_call

from radiens_core._base_file_client import BaseFileClient
from radiens_core._converters._dsp import (
    domain_to_proto_dsp_params,
    proto_to_domain_dsp_group,
)
from radiens_core._converters._kpi import (
    build_bundle_req,
    kpi_metrics_result_from_proto,
    kpi_status_from_proto,
)
from radiens_core._converters._signals import sig_select_to_proto, signal_arrays_from_signals2
from radiens_core._converters._spikes import (
    channel_spike_data_from_raster,
    merge_waveforms_channel,
    merge_waveforms_neuron,
    neuron_spike_data_from_raster,
    neurons_result_from_proto,
    spikes_spec_from_proto,
)
from radiens_core._converters._time_range import time_range_to_proto, trs_mode_to_proto
from radiens_core._converters._validation import validate_proto
from radiens_core._generated import biointerface_pb2, common_pb2, spikesorter_pb2
from radiens_core._services._radiens_core import (
    biointerface_get_neurons_proto,
    get_dsp_group_proto,
    get_hd_snapshot_py_proto,
    get_kpi_status_radiens_proto,
    kpi_calculate_proto,
    kpi_clear_proto,
    kpi_get_metrics_radiens_proto,
    set_dsp_group_proto,
    set_kpi_packet_dur_radiens_proto,
    spikes_get_dense_proto,
    spikes_get_raster_data_proto,
    spikes_get_spec_proto,
)
from radiens_core._typing_helpers import validated
from radiens_core.exceptions import RadiensError
from radiens_core.models.common import ClientType, ServiceEndpoint
from radiens_core.models.dataset import DatasetMetadata, FlexDataset, FlexTimeRange, TimeRangeSpec
from radiens_core.models.kpi import KpiMetricId, KpiMetricsResult, KpiStatus
from radiens_core.models.signals import FlexSignalSpec, SignalArrays, SignalSpec, SignalType
from radiens_core.models.spikes import ChannelSpikeData, NeuronSpikeData, NeuronsResult, SpikesSpec

if TYPE_CHECKING:
    from radiens_core.models.dsp import DSPGroup, DSPParams, FilterStage


class VidereClient(BaseFileClient):
    """Client for offline file-based analysis.

    Auto-discovers a running ``radiensserver`` on the local machine.

    Usage::

        with VidereClient() as client:
            meta = client.link_data_file("/path/to/recording.xdat")
            sigs = client.get_signals(meta, time_range=[0, 1], signals=[0, 1, 2])
    """

    _client_type: ClassVar[ClientType] = ClientType.VIDERE

    # --- Public API ---

    @validate_call(config=ConfigDict(arbitrary_types_allowed=True))
    def get_signals(
        self,
        dataset: FlexDataset,
        time_range: FlexTimeRange | None = None,
        signals: FlexSignalSpec | None = None,
    ) -> SignalArrays:
        """Retrieve signal data for a specific time range and channel selection.

        Args:
            dataset: Dataset to query (metadata object, ID string, or file path).
            time_range: Time range specification (e.g. ``[start, end]`` in seconds).
                ``None`` returns the full recording. Lookback mode is not supported for offline files.
            signals: Signal selection (e.g. list of amplifier channel indices).
                ``None`` returns all amplifier channels.

        Returns:
            The requested signal data as numpy arrays.

        Raises:
            ValueError: If lookback time range mode is used (not supported offline).
            RadiensError: On communication failure or server error.
        """

        meta = self._resolve_dataset(dataset)

        if time_range is not None:
            tr_spec = validated(time_range, TimeRangeSpec)
            concrete_tr = tr_spec.resolve(dataset_metadata=meta)
            if concrete_tr is None:
                # FROM_HEAD (lookback) mode is not meaningful for offline files
                raise ValueError(
                    f"lookback() time range is not supported for offline files; "
                    f"use TimeRangeSpec.subset(start, end) or TimeRangeSpec.full() instead. Got: {time_range}"
                )
        else:
            tr_spec = TimeRangeSpec.full()
            concrete_tr = tr_spec.resolve(dataset_metadata=meta)
            assert concrete_tr is not None  # full() always resolves

        sig_spec = validated(signals, SignalSpec) if signals is not None else SignalSpec.all_amp()

        sig_select = sig_spec.to_sig_select(meta.channel_metadata)

        # Build proto request
        req = common_pb2.HDSnapshotRequest2(
            dsourceID=meta.id,
            tR=time_range_to_proto(concrete_tr),
            selMode=trs_mode_to_proto(tr_spec.mode),
            sigSel=sig_select_to_proto(sig_select),
        )
        validate_proto(req)

        # Execute RPC
        addr = self._server_address(ServiceEndpoint.RADIENS_CORE)
        response = get_hd_snapshot_py_proto(self._pool, addr, req)

        return signal_arrays_from_signals2(response.sigs)

    def get_dsp_group(self, dataset: FlexDataset) -> DSPGroup:
        """Get DSP filter configuration for a dataset.

        Args:
            dataset: Dataset to query (metadata object, ID string, or file path).

        Returns:
            The current hardware and software filter parameters.
        """
        meta = self._resolve_dataset(dataset)
        addr = self._server_address(ServiceEndpoint.RADIENS_CORE)
        response = get_dsp_group_proto(self._pool, addr, meta.id)
        return proto_to_domain_dsp_group(response)

    def set_dsp_group(
        self,
        dataset: FlexDataset,
        *,
        stage: FilterStage,
        params: list[DSPParams],
    ) -> None:
        """Update DSP filter configuration for a specific dataset and stage.

        Args:
            dataset: Dataset to update (metadata object, ID string, or file path).
            stage: The filter stage to update (e.g. ``STAGE1``, ``STAGE2``).
            params: List of filter parameters for this stage.
        """
        meta = self._resolve_dataset(dataset)
        req = common_pb2.SetDSPGroupRequest(
            streamGroupId=meta.id,
            stage=cast("common_pb2.FilterStage.ValueType", stage.value),
            dspParams=[domain_to_proto_dsp_params(p) for p in params],
        )
        validate_proto(req)

        addr = self._server_address(ServiceEndpoint.RADIENS_CORE)
        set_dsp_group_proto(self._pool, addr, req)

    # ===================================================================
    # KPI Metrics
    # ===================================================================

    @validate_call(config=ConfigDict(arbitrary_types_allowed=True))
    def get_kpi_metrics(
        self,
        dataset: FlexDataset,
        metrics: list[KpiMetricId],
        channel_indices: list[int] | None = None,
        time_range: FlexTimeRange | None = None,
        signal_type: SignalType = SignalType.AMP,
    ) -> KpiMetricsResult:
        """Retrieve computed KPI signal metrics for a linked dataset.

        Args:
            dataset: Dataset to query (metadata object, ID string, or file path).
            metrics: List of KPI metric identifiers to retrieve.
            channel_indices: Native channel indices (ntvIdxs) to query.
                ``None`` returns metrics for all channels.
            time_range: Time range within the dataset (e.g. ``[start, end]`` in seconds).
                ``None`` returns metrics over the full recording. Lookback mode is not supported for offline files.
            signal_type: Signal type to query (default ``AMP``).

        Returns:
            Metric values as a ``(n_channels, n_metrics)`` numpy array with row/column labels.

        Raises:
            ValueError: If lookback time range mode is used (not supported offline).
            RadiensError: On communication failure or server error.
        """
        meta = self._resolve_dataset(dataset)

        if time_range is not None:
            tr_spec = validated(time_range, TimeRangeSpec)
            concrete_tr = tr_spec.resolve(dataset_metadata=meta)
            if concrete_tr is None:
                raise ValueError(
                    "lookback() time range is not supported for offline files; "
                    f"use TimeRangeSpec.subset(start, end) or TimeRangeSpec.full() instead. Got: {time_range}"
                )
        else:
            tr_spec = TimeRangeSpec.full()
            concrete_tr = tr_spec.resolve(dataset_metadata=meta)
            assert concrete_tr is not None  # full() always resolves

        chan_list = channel_indices if channel_indices is not None else []
        req = common_pb2.KpiMetricsReq(
            streamGroupId=meta.id,
            stype=cast("common_pb2.SignalType.ValueType", signal_type.value),
            arg=build_bundle_req(chan_list, list(concrete_tr.sec), metrics),
        )
        addr = self._server_address(ServiceEndpoint.RADIENS_CORE)
        response = kpi_get_metrics_radiens_proto(self._pool, addr, req)
        return kpi_metrics_result_from_proto(response)

    def get_kpi_status(self, dataset: FlexDataset) -> KpiStatus:
        """Get the current KPI cache status for a dataset.

        Args:
            dataset: Dataset to query (metadata object, ID string, or file path).

        Returns:
            Number of packets in cache, packet duration, memory usage, and availability state.
        """
        meta = self._resolve_dataset(dataset)
        req = common_pb2.GetKpiStatusRequest(streamGroupId=meta.id)
        addr = self._server_address(ServiceEndpoint.RADIENS_CORE)
        response = get_kpi_status_radiens_proto(self._pool, addr, req)
        return kpi_status_from_proto(response)

    def kpi_calculate(self, dataset: FlexDataset) -> None:
        """Trigger KPI computation for a linked dataset.

        Must be called before ``get_kpi_metrics`` returns useful data
        on a freshly linked file.

        Args:
            dataset: Dataset to compute KPIs for (metadata object, ID string, or file path).

        Raises:
            RadiensError: On communication failure or server error.
        """
        meta = self._resolve_dataset(dataset)
        req = common_pb2.KpiStandardRequest(dsourceID=[meta.id])
        addr = self._server_address(ServiceEndpoint.RADIENS_CORE)
        response = kpi_calculate_proto(self._pool, addr, req)
        if response.errMsg:
            raise RadiensError(f"KPI calculate failed: {response.errMsg}")

    def kpi_clear(self, dataset: FlexDataset) -> None:
        """Clear the KPI cache for a dataset.

        Resets accumulated KPI state. Useful before rerunning with
        changed parameters.

        Args:
            dataset: Dataset to clear KPI state for (metadata object, ID string, or file path).

        Raises:
            RadiensError: On communication failure or server error.
        """
        meta = self._resolve_dataset(dataset)
        req = common_pb2.KpiStandardRequest(dsourceID=[meta.id])
        addr = self._server_address(ServiceEndpoint.RADIENS_CORE)
        response = kpi_clear_proto(self._pool, addr, req)
        if response.errMsg:
            raise RadiensError(f"KPI clear failed: {response.errMsg}")

    def set_kpi_packet_dur(self, dataset: FlexDataset, packet_dur: float) -> None:
        """Set the KPI packet duration (integration window) for a dataset.

        Args:
            dataset: Dataset to configure (metadata object, ID string, or file path).
            packet_dur: Duration of each KPI packet in seconds.

        Raises:
            RadiensError: On communication failure or server error.
        """
        meta = self._resolve_dataset(dataset)
        req = common_pb2.SetKpiParamRequest(streamGroupId=meta.id, param=packet_dur)
        addr = self._server_address(ServiceEndpoint.RADIENS_CORE)
        response = set_kpi_packet_dur_radiens_proto(self._pool, addr, req)
        if response.errMsg:
            raise RadiensError(f"Failed to set KPI packet duration: {response.errMsg}")

    # ===================================================================
    # Spike Data
    # ===================================================================

    def _resolve_spikes_id(self, meta: DatasetMetadata) -> str:
        """Return the registered spikes datasource ID for a dataset.

        Raises RadiensError if no .spikes file was found when the dataset was linked.
        """
        if meta.associated_spikes_ids:
            return meta.associated_spikes_ids[0]
        raise RadiensError(
            f"No spike data registered for dataset '{meta.id}'. "
            "Ensure a .spikes file is present alongside the recording file."
        )

    def _resolve_amp_channels(
        self,
        channels: FlexSignalSpec | None,
        meta: DatasetMetadata,
    ) -> list[int]:
        """Resolve a channel specification to a list of AMP native indices.

        Args:
            channels: Channel specification (list[int], SignalSpec, "all", or None).
                ``None`` returns all amplifier channels.
            meta: Dataset metadata for channel resolution.

        Returns:
            Sorted list of native AMP channel indices.
        """
        sig_spec = validated(channels, SignalSpec) if channels is not None else SignalSpec.all_amp()
        sig_select = sig_spec.to_sig_select(meta.channel_metadata)
        return sorted(int(i) for i in sig_select.amp)

    def get_spikes_spec(self, dataset: FlexDataset) -> SpikesSpec:
        """Retrieve spike interface metadata for a linked dataset.

        Args:
            dataset: Dataset to query (metadata object, ID string, or file path).

        Returns:
            Spike metadata including the biointerface ID, enabled channel indices, and spike time range.

        Raises:
            RadiensError: On communication failure or server error.
        """
        meta = self._resolve_dataset(dataset)
        spikes_id = self._resolve_spikes_id(meta)
        addr = self._server_address(ServiceEndpoint.RADIENS_CORE)
        response = spikes_get_spec_proto(self._pool, addr, spikes_id)
        return spikes_spec_from_proto(response)

    @validate_call(config=ConfigDict(arbitrary_types_allowed=True))
    def get_spikes_by_channel(
        self,
        dataset: FlexDataset,
        channels: FlexSignalSpec | None = None,
        time_range: FlexTimeRange | None = None,
        include_waveforms: bool = False,
    ) -> dict[int, ChannelSpikeData]:
        """Retrieve spike timestamps and neuron labels per recording channel.

        Returns one ``ChannelSpikeData`` entry per channel that has spikes in the requested time range.

        Args:
            dataset: Dataset to query (metadata object, ID string, or file path).
            channels: Channels to include. Accepts a ``list[int]`` of amplifier
                channel indices, a ``SignalSpec``, ``"all"``, or ``None``.
                ``None`` returns all amplifier channels.
            time_range: Time range within the dataset (e.g. ``[start, end]`` in seconds).
                ``None`` returns spikes over the full recording. Lookback mode is not supported for offline files.
            include_waveforms: If ``True``, also fetch waveform snippets for each spike.
                Adds a second RPC call. Default ``False``.

        Returns:
            Map from native channel index to spike data.

        Raises:
            ValueError: If lookback time range mode is used (not supported offline).
            RadiensError: On communication failure or server error.
        """
        meta = self._resolve_dataset(dataset)
        amp_indices = self._resolve_amp_channels(channels, meta)

        if time_range is not None:
            tr_spec = validated(time_range, TimeRangeSpec)
            concrete_tr = tr_spec.resolve(dataset_metadata=meta)
            if concrete_tr is None:
                raise ValueError(
                    "lookback() time range is not supported for offline files; "
                    f"use TimeRangeSpec.subset(start, end) or TimeRangeSpec.full() instead. Got: {time_range}"
                )
            start, end = concrete_tr.sec
        else:
            start, end = 0.0, meta.duration

        spikes_id = self._resolve_spikes_id(meta)
        addr = self._server_address(ServiceEndpoint.RADIENS_CORE)
        spec_reply = spikes_get_spec_proto(self._pool, addr, spikes_id)

        duration = end - start
        raster_req = spikesorter_pb2.SpikeSorterGetRasterDataRequest(
            spikeSorterID=spec_reply.biointerfaceID,
            timeRange=[start, end],
            mode=spikesorter_pb2.RasterMode.CHANNELS,
            timeWindow=duration,
            plotWidthPoints=max(meta.sampling_rate * duration, 1.0),
            componentID="radiens_core_python",
            ntvChanIdx=amp_indices,
        )
        raster_reply = spikes_get_raster_data_proto(self._pool, addr, raster_req)
        result = channel_spike_data_from_raster(raster_reply)

        if include_waveforms:
            dense_req = biointerface_pb2.SpikesGetSpikesRequest(
                dsourceID=spikes_id,
                ntvChanIdx=amp_indices,
                timeStart=start,
                timeStop=end,
                mode=biointerface_pb2.NeuronsSignalMode.SpikeWaveforms,
                maxSpikesPerChannel=-1,
            )
            dense_reply = spikes_get_dense_proto(self._pool, addr, dense_req)
            result = merge_waveforms_channel(result, raster_reply, dense_reply)

        return result

    @validate_call(config=ConfigDict(arbitrary_types_allowed=True))
    def get_spikes_by_neuron(
        self,
        dataset: FlexDataset,
        channels: FlexSignalSpec | None = None,
        time_range: FlexTimeRange | None = None,
        include_waveforms: bool = False,
    ) -> dict[int, NeuronSpikeData]:
        """Retrieve spike timestamps per sorted neuron.

        Same data as ``get_spikes_by_channel()`` but regrouped by neuron identity rather than by channel.

        Args:
            dataset: Dataset to query (metadata object, ID string, or file path).
            channels: Channels to include. Accepts a ``list[int]`` of amplifier
                channel indices, a ``SignalSpec``, ``"all"``, or ``None``.
                ``None`` returns all amplifier channels.
            time_range: Time range within the dataset (e.g. ``[start, end]`` in seconds).
                ``None`` returns spikes over the full recording. Lookback mode is not supported for offline files.
            include_waveforms: If ``True``, also fetch waveform snippets for each spike.
                Adds a second RPC call. Default ``False``.

        Returns:
            Map from integer neuron label to spike data.

        Raises:
            ValueError: If lookback time range mode is used (not supported offline).
            RadiensError: On communication failure or server error.
        """
        meta = self._resolve_dataset(dataset)
        amp_indices = self._resolve_amp_channels(channels, meta)

        if time_range is not None:
            tr_spec = validated(time_range, TimeRangeSpec)
            concrete_tr = tr_spec.resolve(dataset_metadata=meta)
            if concrete_tr is None:
                raise ValueError(
                    "lookback() time range is not supported for offline files; "
                    f"use TimeRangeSpec.subset(start, end) or TimeRangeSpec.full() instead. Got: {time_range}"
                )
            start, end = concrete_tr.sec
        else:
            start, end = 0.0, meta.duration

        spikes_id = self._resolve_spikes_id(meta)
        addr = self._server_address(ServiceEndpoint.RADIENS_CORE)
        spec_reply = spikes_get_spec_proto(self._pool, addr, spikes_id)

        # Build neuron-label → channel mapping from the spec
        label_to_chan: dict[int, int] = {}
        for site in spec_reply.site:
            for nd in site.neuron:
                with contextlib.suppress(ValueError):
                    label_to_chan[int(nd.spikeLabel)] = nd.siteNtvChanIdx

        duration = end - start
        raster_req = spikesorter_pb2.SpikeSorterGetRasterDataRequest(
            spikeSorterID=spec_reply.biointerfaceID,
            timeRange=[start, end],
            mode=spikesorter_pb2.RasterMode.NEURONS,
            timeWindow=duration,
            plotWidthPoints=max(meta.sampling_rate * duration, 1.0),
            componentID="radiens_core_python",
            ntvChanIdx=amp_indices,
        )
        raster_reply = spikes_get_raster_data_proto(self._pool, addr, raster_req)
        result = neuron_spike_data_from_raster(raster_reply, label_to_chan)

        if include_waveforms:
            dense_req = biointerface_pb2.SpikesGetSpikesRequest(
                dsourceID=spikes_id,
                ntvChanIdx=amp_indices,
                timeStart=start,
                timeStop=end,
                mode=biointerface_pb2.NeuronsSignalMode.SpikeWaveforms,
                maxSpikesPerChannel=-1,
            )
            dense_reply = spikes_get_dense_proto(self._pool, addr, dense_req)
            result = merge_waveforms_neuron(result, raster_reply, dense_reply)

        return result

    @validate_call(config=ConfigDict(arbitrary_types_allowed=True))
    def get_neurons(
        self,
        dataset: FlexDataset,
        channels: FlexSignalSpec | None = None,
        time_range: FlexTimeRange | None = None,
    ) -> NeuronsResult:
        """Retrieve neuron descriptors for a linked dataset.

        Returns spike count, rate, SNR, waveform peak amplitude, and probe
        position for each sorted neuron, organized by recording channel.

        Args:
            dataset: Dataset to query (metadata object, ID string, or file path).
            channels: Channels to include. Accepts a ``list[int]`` of amplifier
                channel indices, a ``SignalSpec``, ``"all"``, or ``None``.
                ``None`` returns all amplifier channels.
            time_range: Time range within the dataset (e.g. ``[start, end]`` in seconds).
                ``None`` returns neurons over the full recording. Lookback mode is not
                supported for offline files.

        Returns:
            Neuron descriptors grouped by native channel index.

        Raises:
            ValueError: If lookback time range mode is used (not supported offline).
            RadiensError: On communication failure or server error.
        """
        meta = self._resolve_dataset(dataset)
        amp_indices = self._resolve_amp_channels(channels, meta)
        spikes_id = self._resolve_spikes_id(meta)

        if time_range is not None:
            tr_spec = validated(time_range, TimeRangeSpec)
            concrete_tr = tr_spec.resolve(dataset_metadata=meta)
            if concrete_tr is None:
                raise ValueError(
                    "lookback() time range is not supported for offline files; "
                    "use TimeRangeSpec.subset(start, end) or TimeRangeSpec.full() instead."
                )
            tr_list: list[float] = list(concrete_tr.sec)
        else:
            tr_list = []

        req = biointerface_pb2.BiointerfaceGetNeuronsRequest(
            dsourceID=spikes_id,
            ntvChanIdx=amp_indices,
            timeRange=tr_list,
        )
        addr = self._server_address(ServiceEndpoint.RADIENS_CORE)
        response = biointerface_get_neurons_proto(self._pool, addr, req)
        return neurons_result_from_proto(response)
