"""Tests for VidereClient."""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from collections.abc import Generator
from unittest.mock import patch

import numpy as np
import pytest

from radiens_core._generated import common_pb2
from radiens_core.models.channels import ChannelMetadata, KeyIdxs, SignalUnits
from radiens_core.models.dataset import DatasetMetadata, RadiensFileType
from radiens_core.models.signals import SignalArrays, SignalType
from radiens_core.models.time_range import TimeRange
from radiens_core.videre_client import VidereClient


@pytest.fixture
def mock_discovery() -> Generator[Any, None, None]:  # type: ignore[explicit-any]
    """Mock server discovery."""
    with patch("radiens_core._base_file_client.discover_radiens_addresses") as mock:
        mock.return_value = {
            "RADIENS_CORE": "localhost:50055",
            "RADIENS_LIFECYCLE": "localhost:50055",
        }
        yield mock


@pytest.fixture
def mock_auth() -> Generator[Any, None, None]:  # type: ignore[explicit-any]
    """Mock authentication interceptor."""
    with patch("radiens_core._base_client.AuthInterceptor") as mock:
        yield mock


def test_get_signals(mock_discovery: Any, mock_auth: Any) -> None:  # type: ignore[explicit-any]
    """Test get_signals method calls the correct RPC."""
    client = VidereClient()

    # Create dummy metadata
    sig_map = {st: KeyIdxs(ntv=[0], dset=[0], sys=[0]) for st in SignalType}
    channel_metadata = ChannelMetadata(
        dataset_uid="test_uid",
        fs=30000.0,
        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=[],
    )
    time_range = TimeRange(
        sec=(0.0, 1.0),
        timestamp=(0, 30000),
        fs=30000.0,
        dur_sec=1.0,
        walltime_str="2024-01-01 00:00:00",
        n=30000,
    )
    meta = DatasetMetadata(
        id="test_ds",
        base_name="test",
        path="/path/to/test",
        file_type=RadiensFileType.XDAT,
        time_range=time_range,
        channel_metadata=channel_metadata,
    )

    # Register the metadata manually
    client._datasets[meta.id] = meta

    # Mock RPC response
    mock_res = common_pb2.HDSnapshot2()

    # Create fake signal data: 1 channel, 100 samples
    data = np.random.rand(1, 100).astype(np.float32)
    mock_res.sigs.amp.data = data.tobytes()
    mock_res.sigs.amp.shape.extend([1, 100])
    mock_res.sigs.sG.CopyFrom(common_pb2.SignalGroup())  # Minimal valid proto
    mock_res.sigs.tR.CopyFrom(common_pb2.TimeRange(fs=30000.0))

    with patch("radiens_core.videre_client.get_hd_snapshot_py_proto", return_value=mock_res) as mock_rpc:
        sigs = client.get_signals(meta, time_range=[0, 0.1], signals=[0])

        assert isinstance(sigs, SignalArrays)
        assert sigs.amp is not None
        assert sigs.amp.shape == (1, 100)
        np.testing.assert_allclose(sigs.amp, data, atol=1e-6)

        # Verify RPC call
        mock_rpc.assert_called_once()
        args, _ = mock_rpc.call_args
        req = args[2]
        assert req.dsourceID == meta.id
        assert list(req.sigSel.ampNtvIdxs) == [0]


def test_link_data_file(mock_discovery: Any, mock_auth: Any) -> None:  # type: ignore[explicit-any]
    """Test link_data_file calls set_datasource with the correct path."""
    client = VidereClient()

    sig_map = {st: KeyIdxs(ntv=[0], dset=[0], sys=[0]) for st in SignalType}
    channel_metadata = ChannelMetadata(
        dataset_uid="test_uid",
        fs=30000.0,
        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=[],
    )
    time_range = TimeRange(
        sec=(0.0, 10.0),
        timestamp=(0, 300000),
        fs=30000.0,
        dur_sec=10.0,
        walltime_str="2024-01-01 00:00:00",
        n=300000,
    )
    expected_meta = DatasetMetadata(
        id="ds_001",
        base_name="recording",
        path="/path/to/file.xdat",
        file_type=RadiensFileType.XDAT,
        time_range=time_range,
        channel_metadata=channel_metadata,
    )

    with patch("radiens_core._base_file_client.set_datasource", return_value=expected_meta) as mock_set:
        result = client.link_data_file("/path/to/file.xdat")

        assert result == expected_meta
        mock_set.assert_called_once()
        _, _, path_arg = mock_set.call_args[0]
        assert path_arg == "/path/to/file.xdat"
