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

from __future__ import annotations

from typing import TYPE_CHECKING

from radiens_core._constants import ALLEGO_CORE_ADDR, PCACHE_ADDR
from radiens_core._converters._record_state import record_state_request_to_proto
from radiens_core._converters._restart import restart_request_to_proto
from radiens_core._converters._status import backbone_status_from_proto

# Import here to avoid circular imports
from radiens_core._converters._stim import stim_params_from_proto, stim_params_to_proto
from radiens_core._converters._stream_state import stream_state_request_to_proto
from radiens_core._services._allego_core import (
    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_stim_params_proto,
    set_stim_step_proto,
    set_stream_state_proto,
    set_trigger_state_proto,
)
from radiens_core._transport._auth import AuthInterceptor
from radiens_core._transport._channel_pool import ChannelPool
from radiens_core.exceptions import RadiensError
from radiens_core.models.enums import RadiensService

if TYPE_CHECKING:
    from radiens_core.models.status import (
        BackboneStatus,
        ManualStimTriggerRequest,
        ManualStimTriggerToggleRequest,
        RecordStateRequest,
        RestartRequest,
        SetTriggerStateRequest,
        StimStepRequest,
        StimStepState,
        StreamStateRequest,
        TriggerState,
    )
    from radiens_core.models.stim import StimParams


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

    Usage::

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

    def __init__(
        self,
        core_address: str = ALLEGO_CORE_ADDR,
        pcache_address: str = PCACHE_ADDR,
    ) -> None:
        self._core_address = core_address
        self._pcache_address = pcache_address
        self._pool = ChannelPool(
            interceptors=[AuthInterceptor(server_address=core_address)],
        )

    # --- Lifecycle ---

    def close(self) -> None:
        """Close all gRPC channels."""
        self._pool.close()

    def __enter__(self) -> AllegoClient:
        return self

    def __exit__(self, *args: object) -> None:
        self.close()

    # --- Internal helpers ---

    def _server_address(self, service: RadiensService) -> str:
        """Resolve server address for the given service."""
        if service == RadiensService.CORE:
            return self._core_address
        return self._pcache_address

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

        # Check for errors in the response
        if response.errMsg:
            # Import here to avoid circular imports
            from radiens_core.exceptions import RadiensError
            raise RadiensError(f"Failed to set stream state: {response.errMsg}")

    def healthcheck(self) -> None:
        """Perform lightweight system health verification.

        This is a simple health check that verifies the system is responsive.
        Use `get_status()` for comprehensive system information.

        Raises:
            RadiensError: If the health check fails or system is unresponsive
        """
        response = healthcheck_proto(self._pool, self._core_address)

        # Check for errors in the response
        if response.errMsg:
            # Import here to avoid circular imports
            from radiens_core.exceptions import RadiensError
            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)

        # Check for errors in the response
        if response.errMsg:
            # Import here to avoid circular imports
            from radiens_core.exceptions import RadiensError
            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)

        # Check for errors in the response
        if response.errMsg:
            # Import here to avoid circular imports
            from radiens_core.exceptions import RadiensError
            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=0.0001,  # 100 μs
                    second_phase_duration=0.0001,
                    interphase_delay=0.00005,  # 50 μs
                    first_phase_amplitude=100.0,  # 100 μA
                    second_phase_amplitude=-100.0,
                    baseline_voltage=0.0,
                    trigger_edge_or_level=StimTriggerEdge.EDGE,
                    trigger_high_or_low=StimTriggerLevel.HIGH,
                    enabled=True,
                    post_trigger_delay=0.001,
                    pulse_or_train=StimPulseMode.SINGLE_PULSE,
                    number_of_stim_pulses=1,
                    pulse_train_period=0.02,
                    refractory_period=0.01,
                    pre_stim_amp_settle=0.001,
                    post_stim_amp_settle=0.001,
                    maintain_amp_settle=False,
                    enable_amp_settle=True,
                    post_stim_charge_recov_on=0.001,
                    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
        """
        from radiens_core._converters._trigger_control import set_trigger_state_request_to_proto

        proto_request = set_trigger_state_request_to_proto(request)
        response = set_trigger_state_proto(self._pool, self._core_address, proto_request)

        # Check for errors in the response
        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
        """
        from radiens_core._converters._trigger_control import trigger_state_from_proto

        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
        """
        from radiens_core._converters._trigger_control import manual_stim_trigger_request_to_proto

        proto_request = manual_stim_trigger_request_to_proto(request)
        response = manual_stim_trigger_proto(self._pool, self._core_address, proto_request)

        # Check for errors in the response
        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
        """
        from radiens_core._converters._trigger_control import manual_stim_trigger_toggle_request_to_proto

        proto_request = manual_stim_trigger_toggle_request_to_proto(request)
        response = manual_stim_trigger_toggle_proto(self._pool, self._core_address, proto_request)

        # Check for errors in the response
        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
        """
        from radiens_core._converters._stim_step import stim_step_request_to_proto

        proto_request = stim_step_request_to_proto(request)
        response = set_stim_step_proto(self._pool, self._core_address, proto_request)

        # Check for errors in the response
        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
        """
        from radiens_core._converters._stim_step import stim_step_state_from_proto

        response = get_stim_step_proto(self._pool, self._core_address)
        return stim_step_state_from_proto(response)
