"""Tests for AllegoClient.get_signals."""

from __future__ import annotations

from unittest.mock import MagicMock, patch

import numpy as np
import pytest

from radiens_core._constants import PRIMARY_CACHE_STREAM_GROUP_ID
from radiens_core._generated import common_pb2, datasource_pb2
from radiens_core.allego_client import AllegoClient
from radiens_core.exceptions import RadiensError
from radiens_core.models.channels import ChannelMetadata, KeyIdxs, SignalUnits
from radiens_core.models.dataset import TimeRangeSpec
from radiens_core.models.ports import Port
from radiens_core.models.signals import SignalArrays, SignalType
from radiens_core.models.status import (
    BackboneMode,
    RecordMode,
    RecordStateRequest,
    StreamMode,
    StreamStateRequest,
)


def _make_channel_metadata(fs: float = 30000.0) -> ChannelMetadata:
    sig_map = {st: KeyIdxs(ntv=[0], dset=[0], sys=[0]) for st in SignalType}
    return ChannelMetadata(
        dataset_uid="live",
        fs=fs,
        source_label="test_source",
        neighbor_max_radius_um=100.0,
        signal_units={SignalType.AMP: SignalUnits.MICROVOLTS},
        has_discrete=False,
        has_continuous_amp=True,
        sig_map=sig_map,
        selected_sig_map=sig_map,
        num_channels={st: 1 for st in SignalType},
        selected_num_channels={st: 1 for st in SignalType},
        site_positions=[],
    )


def _make_hd_snapshot(n_channels: int = 1, n_samples: int = 30000) -> common_pb2.HDSnapshot:
    data = np.random.rand(n_channels, n_samples).astype(np.float32)
    snapshot = common_pb2.HDSnapshot()
    snapshot.priSignals.CopyFrom(
        common_pb2.RawSignals(
            data=data.tobytes(),
            shape=[n_channels, n_samples],
        )
    )
    return snapshot


def test_get_signals_lookback() -> None:
    """get_signals with FROM_HEAD mode sends correct request and returns SignalArrays."""
    with (
        patch("radiens_core._base_client.AuthInterceptor"),
        patch("radiens_core.allego_client.get_signal_group_proto"),
        patch(
            "radiens_core.allego_client.channel_metadata_from_proto",
            return_value=_make_channel_metadata(),
        ),
        patch(
            "radiens_core.allego_client.get_hd_snapshot_proto",
            return_value=_make_hd_snapshot(1, 30000),
        ) as mock_rpc,
    ):
        client = AllegoClient()
        sigs = client.get_signals(
            time_range=TimeRangeSpec.lookback(1.0),
            signals=[0],
        )

    assert isinstance(sigs, SignalArrays)
    assert sigs.amp is not None
    assert sigs.amp.shape == (1, 30000)

    mock_rpc.assert_called_once()
    args, _ = mock_rpc.call_args
    req: common_pb2.HDSnapshotRequest = args[2]
    assert req.streamGroupId == PRIMARY_CACHE_STREAM_GROUP_ID
    assert list(req.timeRange) == [1.0, float("nan")] or (req.timeRange[0] == 1.0 and np.isnan(req.timeRange[1]))
    assert list(req.priSignals.ntvChanIdxs) == [0]


def test_get_signals_subset() -> None:
    """get_signals with SUBSET mode sends [start, end] timeRange."""
    with (
        patch("radiens_core._base_client.AuthInterceptor"),
        patch("radiens_core.allego_client.get_signal_group_proto"),
        patch(
            "radiens_core.allego_client.channel_metadata_from_proto",
            return_value=_make_channel_metadata(),
        ),
        patch(
            "radiens_core.allego_client.get_hd_snapshot_proto",
            return_value=_make_hd_snapshot(1, 30000),
        ) as mock_rpc,
    ):
        client = AllegoClient()
        sigs = client.get_signals(time_range=[0.0, 1.0], signals=[0])

    assert isinstance(sigs, SignalArrays)
    args, _ = mock_rpc.call_args
    req: common_pb2.HDSnapshotRequest = args[2]
    assert list(req.timeRange) == [0.0, 1.0]
    assert req.streamGroupId == PRIMARY_CACHE_STREAM_GROUP_ID


def test_get_signals_lookback_shorthand() -> None:
    """get_signals accepts [duration, nan] shorthand and sends FROM_HEAD-style request."""
    with (
        patch("radiens_core._base_client.AuthInterceptor"),
        patch("radiens_core.allego_client.get_signal_group_proto"),
        patch(
            "radiens_core.allego_client.channel_metadata_from_proto",
            return_value=_make_channel_metadata(),
        ),
        patch(
            "radiens_core.allego_client.get_hd_snapshot_proto",
            return_value=_make_hd_snapshot(1, 30000),
        ) as mock_rpc,
    ):
        client = AllegoClient()
        sigs = client.get_signals(time_range=[1.0, float("nan")], signals=[0])

    assert isinstance(sigs, SignalArrays)
    args, _ = mock_rpc.call_args
    req: common_pb2.HDSnapshotRequest = args[2]
    assert req.timeRange[0] == 1.0
    assert np.isnan(req.timeRange[1])


# ---------------------------------------------------------------------------
# set_stream_state
# ---------------------------------------------------------------------------


def test_set_stream_state_sends_on_request() -> None:
    """set_stream_state sends correct proto mode for StreamMode.ON."""
    mock_reply = MagicMock()
    mock_reply.errMsg = ""
    with (
        patch("radiens_core._base_client.AuthInterceptor"),
        patch("radiens_core.allego_client.set_stream_state_proto", return_value=mock_reply) as mock_rpc,
    ):
        client = AllegoClient()
        client.set_stream_state(StreamStateRequest(mode=StreamMode.ON))

    mock_rpc.assert_called_once()
    args, _ = mock_rpc.call_args
    assert args[2].mode == 1  # S_ON


def test_set_stream_state_raises_on_error() -> None:
    """set_stream_state raises RadiensError when response.errMsg is non-empty."""
    mock_reply = MagicMock()
    mock_reply.errMsg = "stream error"
    with (
        patch("radiens_core._base_client.AuthInterceptor"),
        patch("radiens_core.allego_client.set_stream_state_proto", return_value=mock_reply),
    ):
        client = AllegoClient()
        with pytest.raises(RadiensError):
            client.set_stream_state(StreamStateRequest(mode=StreamMode.OFF))


# ---------------------------------------------------------------------------
# set_record_state
# ---------------------------------------------------------------------------


def test_set_record_state_on_no_port() -> None:
    """set_record_state with no port sends mode ON and leaves port field unset."""
    mock_reply = MagicMock()
    mock_reply.errMsg = ""
    with (
        patch("radiens_core._base_client.AuthInterceptor"),
        patch("radiens_core.allego_client.set_record_state_proto", return_value=mock_reply) as mock_rpc,
    ):
        client = AllegoClient()
        client.set_record_state(RecordStateRequest(mode=RecordMode.ON))

    mock_rpc.assert_called_once()
    args, _ = mock_rpc.call_args
    proto_req = args[2]
    assert proto_req.mode == 1  # R_ON
    assert not proto_req.HasField("port")


def test_set_record_state_with_port() -> None:
    """set_record_state with Port.A sets proto port field to 0."""
    mock_reply = MagicMock()
    mock_reply.errMsg = ""
    with (
        patch("radiens_core._base_client.AuthInterceptor"),
        patch("radiens_core.allego_client.set_record_state_proto", return_value=mock_reply) as mock_rpc,
    ):
        client = AllegoClient()
        client.set_record_state(RecordStateRequest(mode=RecordMode.ON, port=Port.A))

    mock_rpc.assert_called_once()
    args, _ = mock_rpc.call_args
    proto_req = args[2]
    assert proto_req.HasField("port")
    assert proto_req.port == 0  # common_pb2.A


def test_set_record_state_raises_on_error() -> None:
    """set_record_state raises RadiensError when response.errMsg is non-empty."""
    mock_reply = MagicMock()
    mock_reply.errMsg = "record error"
    with (
        patch("radiens_core._base_client.AuthInterceptor"),
        patch("radiens_core.allego_client.set_record_state_proto", return_value=mock_reply),
    ):
        client = AllegoClient()
        with pytest.raises(RadiensError):
            client.set_record_state(RecordStateRequest(mode=RecordMode.OFF))


# ---------------------------------------------------------------------------
# restart
# ---------------------------------------------------------------------------


def test_restart_sends_correct_mode() -> None:
    """restart sends proto mode 0 for BackboneMode.SMARTBOX_PRO."""
    mock_reply = MagicMock()
    mock_reply.errMsg = ""
    with (
        patch("radiens_core._base_client.AuthInterceptor"),
        patch("radiens_core.allego_client.restart_proto", return_value=mock_reply) as mock_rpc,
    ):
        client = AllegoClient()
        client.restart(mode=BackboneMode.SMARTBOX_PRO)

    mock_rpc.assert_called_once()
    args, _ = mock_rpc.call_args
    assert args[2].mode == 0  # SMARTBOX_PRO


def test_restart_raises_on_error() -> None:
    """restart raises RadiensError when response.errMsg is non-empty."""
    mock_reply = MagicMock()
    mock_reply.errMsg = "restart error"
    with (
        patch("radiens_core._base_client.AuthInterceptor"),
        patch("radiens_core.allego_client.restart_proto", return_value=mock_reply),
    ):
        client = AllegoClient()
        with pytest.raises(RadiensError):
            client.restart(mode=BackboneMode.SMARTBOX_PRO)


# ---------------------------------------------------------------------------
# Hardware Diagnostics
# ---------------------------------------------------------------------------


def test_get_intan_impedance() -> None:
    """get_intan_impedance sends correct request and returns Impedance."""
    mock_reply = datasource_pb2.Impedance()
    mock_reply.impedances.add(ntvChanIdx=0, zMag=1000.0, zPhase=0.0)

    with (
        patch("radiens_core._base_client.AuthInterceptor"),
        patch("radiens_core.allego_client.get_intan_impedance_proto", return_value=mock_reply) as mock_rpc,
    ):
        client = AllegoClient()
        result = client.get_intan_impedance(Port.A)

    mock_rpc.assert_called_once()
    assert len(result.measurements) == 1
    assert result.measurements[0].ntv_chan_idx == 0
    assert result.measurements[0].magnitude_ohm == 1000.0


def test_scan_ports() -> None:
    """scan_ports sends request and returns ChannelMetadata."""
    mock_reply = common_pb2.SignalGroup()
    with (
        patch("radiens_core._base_client.AuthInterceptor"),
        patch("radiens_core.allego_client.scan_ports_proto", return_value=mock_reply) as mock_rpc,
        patch(
            "radiens_core.allego_client.channel_metadata_from_proto",
            return_value=_make_channel_metadata(),
        ),
    ):
        client = AllegoClient()
        result = client.scan_ports()

    mock_rpc.assert_called_once()
    assert isinstance(result, ChannelMetadata)
