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

from __future__ import annotations

from unittest.mock import patch

import numpy as np

from radiens_core._constants import PRIMARY_CACHE_STREAM_GROUP_ID
from radiens_core._generated import common_pb2
from radiens_core.allego_client import AllegoClient
from radiens_core.models.channels import ChannelMetadata, KeyIdxs, SignalUnits
from radiens_core.models.dataset import TimeRangeSpec
from radiens_core.models.signals import SignalArrays, SignalType


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