"""AllegoClient — real-time data acquisition client."""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, cast

from pydantic import ConfigDict, validate_call

from radiens_core._base_client import BaseClient
from radiens_core._constants import (
    ALLEGO_SERVICE_DEFAULTS,
    PRIMARY_CACHE_STREAM_GROUP_ID,
)
from radiens_core._converters._channels import channel_metadata_from_proto
from radiens_core._converters._record_state import record_state_request_to_proto
from radiens_core._converters._recording_config import recording_config_from_proto, recording_config_to_proto
from radiens_core._converters._restart import restart_request_to_proto
from radiens_core._converters._signals import signal_arrays_from_hd_snapshot
from radiens_core._converters._status import backbone_status_from_proto
from radiens_core._converters._stim import stim_params_from_proto, stim_params_to_proto
from radiens_core._converters._stim_step import stim_step_request_to_proto, stim_step_state_from_proto
from radiens_core._converters._stream_state import stream_state_request_to_proto
from radiens_core._converters._trigger_control import (
    manual_stim_trigger_request_to_proto,
    manual_stim_trigger_toggle_request_to_proto,
    set_trigger_state_request_to_proto,
    trigger_state_from_proto,
)
from radiens_core._generated import common_pb2
from radiens_core._services._allego_core import (
    get_recording_config_proto,
    get_signal_group_proto,
    get_status_proto,
    get_stim_params_proto,
    get_stim_step_proto,
    get_trigger_state_proto,
    healthcheck_proto,
    manual_stim_trigger_proto,
    manual_stim_trigger_toggle_proto,
    restart_proto,
    set_record_state_proto,
    set_recording_config_proto,
    set_stim_params_proto,
    set_stim_step_proto,
    set_stream_state_proto,
    set_trigger_state_proto,
)
from radiens_core._services._allego_pcache import get_hd_snapshot_proto
from radiens_core.exceptions import RadiensError
from radiens_core.models.common import ServiceEndpoint
from radiens_core.models.dataset import FlexTimeRange, TrsMode
from radiens_core.models.signals import FlexSignalSpec, SignalArrays

if TYPE_CHECKING:
    from radiens_core.models.channels import ChannelMetadata
    from radiens_core.models.dataset import TimeRangeSpec
    from radiens_core.models.signals import SignalSpec
    from radiens_core.models.status import (
        BackboneStatus,
        RecordingConfig,
        RecordStateRequest,
        RestartRequest,
        StreamStateRequest,
    )
    from radiens_core.models.stim import StimParams, StimStepRequest, StimStepState
    from radiens_core.models.triggers import (
        ManualStimTriggerRequest,
        ManualStimTriggerToggleRequest,
        SetTriggerStateRequest,
        TriggerState,
    )

_log = logging.getLogger(__name__)


def _time_range_floats(tr_spec: TimeRangeSpec) -> list[float]:
    """Convert a TimeRangeSpec to [start, end_or_nan] for HDSnapshotRequest.timeRange.

    FROM_HEAD: [duration_sec, nan] — server reads back ``duration`` seconds from cache head.
    SUBSET:    [start_sec, end_sec].
    TO_HEAD:   [start_sec, nan].
    """
    if tr_spec.mode == TrsMode.SUBSET:
        return [tr_spec.start, tr_spec.end if tr_spec.end is not None else 0.0]
    return [tr_spec.start, float("nan")]


class AllegoClient(BaseClient):
    """Client for real-time data acquisition and streaming.

    Auto-discovers running servers.  Falls back to well-known allego
    ports if server discovery is unavailable (no radiensserver running).

    Usage::

        with AllegoClient() as client:
            params = client.get_stim_params()
    """

    def __init__(self) -> None:
        addresses = self._discover()
        core_addr = addresses[ServiceEndpoint.ALLEGO_CORE]
        super().__init__(
            addresses=addresses,
            auth_server_address=core_addr,
            fallback_service=ServiceEndpoint.ALLEGO_CORE,
        )

    @staticmethod
    def _discover() -> dict[str, str]:
        return dict(ALLEGO_SERVICE_DEFAULTS)

    # --- Internal helpers ---

    @property
    def _core_address(self) -> str:
        return self._server_address(ServiceEndpoint.ALLEGO_CORE)

    @property
    def _pcache_address(self) -> str:
        return self._server_address(ServiceEndpoint.ALLEGO_PCACHE)

    # --- Public API ---

    def get_stim_params(self) -> dict[int, StimParams]:
        """Get stimulation parameters for all configured channels.

        Returns a dict keyed by system channel index (ntvAbsKeyIdx).
        """
        response = get_stim_params_proto(self._pool, self._core_address)
        return stim_params_from_proto(response)

    def get_status(self) -> BackboneStatus:
        """Get comprehensive system status.

        Returns current streaming status, recording status, port information,
        connection state, and GPIO channel counts.
        """
        response = get_status_proto(self._pool, self._core_address)
        return backbone_status_from_proto(response)

    def get_channel_metadata(self) -> ChannelMetadata:
        """Get current channel metadata for all connected channels.

        Returns:
            Complete channel metadata including probe geometry, signal types,
            channel indices, and site positions for all channels.

        Raises:
            RadiensError: On communication failure or server error
        """
        response = get_signal_group_proto(self._pool, self._core_address)
        return channel_metadata_from_proto(response)

    @validate_call(config=ConfigDict(arbitrary_types_allowed=True))
    def get_signals(
        self,
        time_range: FlexTimeRange,
        signals: FlexSignalSpec,
    ) -> SignalArrays:
        """Retrieve signal data from the primary live cache.

        Parameters
        ----------
        time_range
            Time range spec. Most common: ``TimeRangeSpec.lookback(5.0)`` or
            shorthand ``[duration, float('nan')]``. Also accepts
            ``[start, end]`` for absolute SUBSET requests.
        signals
            Channel selection — list of amp indices, ``SignalSpec.amp([...])``,
            or dict for multi-type.

        Returns
        -------
        SignalArrays
            Signal data as numpy arrays (channels x samples).
        """
        tr_spec = cast("TimeRangeSpec", time_range)
        sig_spec = cast("SignalSpec", signals)

        channel_meta = self.get_channel_metadata()
        sig_select = sig_spec.to_sig_select(channel_meta)

        def _sel(idxs: list[int]) -> common_pb2.HDSnapshotRequest.SelectedSignals | None:
            return common_pb2.HDSnapshotRequest.SelectedSignals(ntvChanIdxs=idxs) if idxs else None

        req = common_pb2.HDSnapshotRequest(
            streamGroupId=PRIMARY_CACHE_STREAM_GROUP_ID,
            timeRange=_time_range_floats(tr_spec),
            priSignals=_sel(sig_select.amp.tolist()),
            auxSignals=_sel(sig_select.gpio_ain.tolist()),
            dinSignals=_sel(sig_select.gpio_din.tolist()),
            doutSignals=_sel(sig_select.gpio_dout.tolist()),
        )

        response = get_hd_snapshot_proto(self._pool, self._pcache_address, req)
        return signal_arrays_from_hd_snapshot(response)

    def set_stream_state(self, request: StreamStateRequest) -> None:
        """Set data streaming state.

        Args:
            request: Stream state configuration specifying the desired mode

        Raises:
            RadiensError: If the stream state change fails
        """
        proto_request = stream_state_request_to_proto(request)
        response = set_stream_state_proto(self._pool, self._core_address, proto_request)

        if response.errMsg:
            raise RadiensError(f"Failed to set stream state: {response.errMsg}")

    def healthcheck(self, service: ServiceEndpoint = ServiceEndpoint.ALLEGO_LIFECYCLE) -> None:
        """Perform lightweight system health verification.

        Parameters
        ----------
        service
            Which gRPC service to probe (default ``ALLEGO_LIFECYCLE``).
            ``ALLEGO_CORE`` is also supported.

        Raises:
            RadiensError: If the health check fails or system is unresponsive
        """
        response = healthcheck_proto(self._pool, self._server_address(service), service=service)
        if response.errMsg:
            raise RadiensError(f"Health check failed: {response.errMsg}")

    def set_record_state(self, request: RecordStateRequest) -> None:
        """Set data recording state.

        Args:
            request: Record state configuration specifying the desired mode and optional port

        Raises:
            RadiensError: If the record state change fails
        """
        proto_request = record_state_request_to_proto(request)
        response = set_record_state_proto(self._pool, self._core_address, proto_request)
        if response.errMsg:
            raise RadiensError(f"Failed to set record state: {response.errMsg}")

    def restart(self, request: RestartRequest) -> None:
        """Restart the system with specified backbone mode.

        Args:
            request: Restart configuration specifying the desired backbone mode

        Raises:
            RadiensError: If the restart fails
        """
        proto_request = restart_request_to_proto(request)
        response = restart_proto(self._pool, self._core_address, proto_request)
        if response.errMsg:
            raise RadiensError(f"System restart failed: {response.errMsg}")

    def set_stim_params(self, params: dict[int, StimParams]) -> None:
        """Update stimulation parameters for specified channels.

        Args:
            params: Map from system channel indices to their new stimulation parameters.
                   Each StimParams defines waveform shape, timing, amplitudes, trigger
                   settings, pulse trains, and recovery parameters for that channel.

        Raises:
            RadiensError: On communication failure or server error

        Note:
            The underlying gRPC API sets parameters for one channel at a time.
            This method calls SetStimParams once per channel in the provided dictionary.

        Usage:
            Update stimulation parameters for multiple channels::

                # Define stimulation parameters
                stim_params = StimParams(
                    stim_shape=StimShape.BIPHASIC,
                    stim_polarity=StimPolarity.CATHODIC_FIRST,
                    first_phase_duration=100.0,  # 100 μs
                    second_phase_duration=100.0,  # 100 μs
                    interphase_delay=50.0,  # 50 μs
                    first_phase_amplitude=100.0,  # 100 μA
                    second_phase_amplitude=-100.0,  # 100 μA (opposite polarity)
                    baseline_voltage=0.0,
                    trigger_edge_or_level=StimTriggerEdge.EDGE,
                    trigger_high_or_low=StimTriggerLevel.HIGH,
                    enabled=True,
                    post_trigger_delay=1000.0,  # 1 ms = 1000 μs
                    pulse_or_train=StimPulseMode.SINGLE_PULSE,
                    number_of_stim_pulses=1,
                    pulse_train_period=20000.0,  # 20 ms = 50 Hz
                    refractory_period=10000.0,  # 10 ms
                    pre_stim_amp_settle=1000.0,  # 1 ms
                    post_stim_amp_settle=1000.0,  # 1 ms
                    maintain_amp_settle=False,
                    enable_amp_settle=True,
                    post_stim_charge_recov_on=1000.0,  # 1 ms
                    post_stim_charge_recov_off=0.0,
                    enable_charge_recovery=True,
                    trigger_source_is_keypress=False,
                    trigger_source_idx=0,
                    stim_sys_chan_idx=42,
                )

                # Update channels 42 and 43
                client.set_stim_params({
                    42: stim_params,
                    43: stim_params.model_copy(update={'first_phase_amplitude': 150.0}),
                })
        """

        if not params:
            msg = "params cannot be empty - must specify at least one channel"
            raise RadiensError(msg)

        # SetStimParams works on one channel at a time
        # Validate and process each channel
        for sys_chan_idx, stim_params in params.items():
            # Ensure the sys_chan_idx matches the one in the StimParams
            if stim_params.stim_sys_chan_idx != sys_chan_idx:
                # Create updated stim_params with correct sys_chan_idx
                updated_params = stim_params.model_copy(update={"stim_sys_chan_idx": sys_chan_idx})
            else:
                updated_params = stim_params

            proto_request = stim_params_to_proto(updated_params)
            response = set_stim_params_proto(self._pool, self._core_address, proto_request)

            # Check for errors in the response
            if response.errMsg:
                raise RadiensError(f"Set stimulation parameters failed for channel {sys_chan_idx}: {response.errMsg}")

    def set_trigger_state(self, request: SetTriggerStateRequest) -> None:
        """Set trigger channel state for stimulation control.

        Args:
            request: Configuration for trigger channel state

        Raises:
            RadiensError: On communication failure or server error
        """
        proto_request = set_trigger_state_request_to_proto(request)
        response = set_trigger_state_proto(self._pool, self._core_address, proto_request)
        if response.errMsg:
            raise RadiensError(f"Set trigger state failed: {response.errMsg}")

    def get_trigger_state(self) -> TriggerState:
        """Get current trigger channel configuration.

        Returns:
            Current trigger state showing enabled channels and ports

        Raises:
            RadiensError: On communication failure or server error
        """
        response = get_trigger_state_proto(self._pool, self._core_address)
        return trigger_state_from_proto(response)

    def manual_stim_trigger(self, request: ManualStimTriggerRequest) -> None:
        """Manually trigger stimulation.

        Args:
            request: Manual trigger configuration

        Raises:
            RadiensError: On communication failure or server error
        """
        proto_request = manual_stim_trigger_request_to_proto(request)
        response = manual_stim_trigger_proto(self._pool, self._core_address, proto_request)
        if response.errMsg:
            raise RadiensError(f"Manual stim trigger failed: {response.errMsg}")

    def manual_stim_trigger_toggle(self, request: ManualStimTriggerToggleRequest) -> None:
        """Toggle manual stimulation trigger.

        Args:
            request: Manual trigger toggle configuration

        Raises:
            RadiensError: On communication failure or server error
        """
        proto_request = manual_stim_trigger_toggle_request_to_proto(request)
        response = manual_stim_trigger_toggle_proto(self._pool, self._core_address, proto_request)
        if response.errMsg:
            raise RadiensError(f"Manual stim trigger toggle failed: {response.errMsg}")

    def set_stim_step(self, request: StimStepRequest) -> None:
        """Set stimulation step size.

        Args:
            request: Stimulation step configuration

        Raises:
            RadiensError: On communication failure or server error
        """
        proto_request = stim_step_request_to_proto(request)
        response = set_stim_step_proto(self._pool, self._core_address, proto_request)
        if response.errMsg:
            raise RadiensError(f"Set stim step failed: {response.errMsg}")

    def get_stim_step(self) -> StimStepState:
        """Get current stimulation step size setting.

        Returns:
            Current stimulation step configuration

        Raises:
            RadiensError: On communication failure or server error
        """
        response = get_stim_step_proto(self._pool, self._core_address)
        return stim_step_state_from_proto(response)

    def set_recording_config(self, config: RecordingConfig) -> None:
        """Set recording configuration settings.

        Configures output file paths, naming, and recording behavior.
        This affects where recorded data will be saved when recording is started.

        Args:
            config: Recording configuration specifying file paths, naming,
                   and recording type (continuous/spikes/both)

        Raises:
            RadiensError: If the recording config change fails

        Example:
            Configure recording to save continuous data with timestamps::

                config = RecordingConfig(
                    base_file_name="my_recording",
                    base_file_path="/path/to/data",
                    recording_type=DataRecordingType.CONTINUOUS,
                    time_stamp=True,
                )
                client.set_recording_config(config)
        """
        proto_request = recording_config_to_proto(config)
        response = set_recording_config_proto(self._pool, self._core_address, proto_request)
        if response.errMsg:
            raise RadiensError(f"Failed to set recording config: {response.errMsg}")

    def get_recording_config(self) -> RecordingConfig:
        """Get current recording configuration settings.

        Returns:
            Current recording configuration including file paths, naming,
            and recording type settings

        Raises:
            RadiensError: On communication failure or server error
        """
        response = get_recording_config_proto(self._pool, self._core_address)
        return recording_config_from_proto(response)
