"""Tests for KPI converters and client methods."""

from __future__ import annotations

from typing import TYPE_CHECKING
from unittest.mock import patch

import numpy as np
import pytest

from radiens_core._constants import PRIMARY_CACHE_STREAM_GROUP_ID
from radiens_core._converters._kpi import (
    _dense_matrix_to_numpy,
    build_bundle_req,
    kpi_metric_id_from_proto,
    kpi_metric_id_to_proto,
    kpi_metrics_result_from_proto,
    kpi_mode_to_proto,
    kpi_status_from_proto,
)
from radiens_core._generated import common_pb2
from radiens_core.models.kpi import KpiMetric, KpiMetricId, KpiMetricsResult, KpiMode, KpiStatus
from radiens_core.models.signals import SignalType
from radiens_core.models.time_range import TimeRange

if TYPE_CHECKING:
    from radiens_core.models.dataset import DatasetMetadata

# =====================================================================
# Converter tests
# =====================================================================


def test_kpi_mode_to_proto_all_values() -> None:
    assert kpi_mode_to_proto(KpiMode.BASE) == common_pb2.Kpi_MetricsMode_Base
    assert kpi_mode_to_proto(KpiMode.AVG) == common_pb2.Kpi_MetricsMode_Avg
    assert kpi_mode_to_proto(KpiMode.STREAM) == common_pb2.Kpi_MetricsMode_Stream


def test_kpi_metric_id_roundtrip() -> None:
    mid = KpiMetricId(mode=KpiMode.BASE, name=KpiMetric.MEAN)
    proto = kpi_metric_id_to_proto(mid)
    assert proto.mode == common_pb2.Kpi_MetricsMode_Base
    assert proto.name == common_pb2.KPI_MEAN
    back = kpi_metric_id_from_proto(proto)
    assert back == mid


def test_kpi_metric_id_roundtrip_avg_rms() -> None:
    mid = KpiMetricId(mode=KpiMode.AVG, name=KpiMetric.RMS)
    back = kpi_metric_id_from_proto(kpi_metric_id_to_proto(mid))
    assert back == mid


def test_dense_matrix_to_numpy_normal() -> None:
    arr = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=np.float64)
    matrix = common_pb2.DenseMatrix(rows=2, cols=3, data=arr.tobytes())
    result = _dense_matrix_to_numpy(matrix)
    np.testing.assert_array_equal(result, arr)
    assert result.shape == (2, 3)


def test_dense_matrix_to_numpy_empty_rows() -> None:
    matrix = common_pb2.DenseMatrix(rows=0, cols=3, data=b"")
    result = _dense_matrix_to_numpy(matrix)
    assert result.shape == (0, 3)


def test_dense_matrix_to_numpy_empty_cols() -> None:
    matrix = common_pb2.DenseMatrix(rows=4, cols=0, data=b"")
    result = _dense_matrix_to_numpy(matrix)
    assert result.shape == (4, 0)


def test_kpi_metrics_result_from_proto() -> None:
    vals = np.array([[1.5, 2.5], [3.5, 4.5]], dtype=np.float64)
    reply = common_pb2.SignalMetricsPacketReply(
        ntvIdxs=[0, 1],
        sigType=common_pb2.PRI,
        metrics=[
            common_pb2.KpiMetricID(mode=common_pb2.Kpi_MetricsMode_Base, name=common_pb2.KPI_MEAN),
            common_pb2.KpiMetricID(mode=common_pb2.Kpi_MetricsMode_Base, name=common_pb2.KPI_RMS),
        ],
        metricVals=common_pb2.DenseMatrix(rows=2, cols=2, data=vals.tobytes()),
        tR=common_pb2.TimeRange(
            timeRangeSec=[0.0, 1.0],
            timestamp=[0, 30000],
            fs=30000.0,
            wallTimeStart="2024-01-01 00:00:00",
        ),
    )
    result = kpi_metrics_result_from_proto(reply)
    assert isinstance(result, KpiMetricsResult)
    assert result.channel_indices == (0, 1)
    assert result.signal_type == SignalType.AMP
    assert len(result.metrics) == 2
    assert result.metrics[0].name == KpiMetric.MEAN
    assert result.metrics[1].name == KpiMetric.RMS
    np.testing.assert_array_equal(result.values, vals)
    assert result.time_range.sec == (0.0, 1.0)
    assert result.time_range.fs == 30000.0


def test_kpi_metrics_result_ain_signal_type() -> None:
    vals = np.array([[7.0]], dtype=np.float64)
    reply = common_pb2.SignalMetricsPacketReply(
        ntvIdxs=[2],
        sigType=common_pb2.AUX,
        metrics=[common_pb2.KpiMetricID(mode=common_pb2.Kpi_MetricsMode_Base, name=common_pb2.KPI_SNR)],
        metricVals=common_pb2.DenseMatrix(rows=1, cols=1, data=vals.tobytes()),
        tR=common_pb2.TimeRange(timeRangeSec=[0.0, 0.5], fs=30000.0),
    )
    result = kpi_metrics_result_from_proto(reply)
    assert result.signal_type == SignalType.AIN


def test_kpi_status_from_proto_no_optional_fields() -> None:
    pb = common_pb2.KpiStatusReply(
        streamGroupId="Live Signals",
        numPacketsMemory=10,
        packetDur=0.5,
        updatePeriod=0.1,
        persistence=0.9,
        beta=0.1,
        isTrackingSignalCache=True,
        numPackets=100,
        memoryBytes=1024.0,
    )
    status = kpi_status_from_proto(pb)
    assert isinstance(status, KpiStatus)
    assert status.stream_group_id == PRIMARY_CACHE_STREAM_GROUP_ID
    assert status.num_packets_memory == 10
    assert status.packet_dur == 0.5
    assert status.update_period == 0.1
    assert status.is_tracking_signal_cache is True
    assert status.num_packets == 100
    assert status.time_range is None
    assert status.num_sigs is None


def test_kpi_status_from_proto_with_optional_fields() -> None:
    pb = common_pb2.KpiStatusReply(
        streamGroupId="ds_001",
        numPacketsMemory=5,
        packetDur=1.0,
        updatePeriod=0.5,
        persistence=0.8,
        beta=0.2,
        isTrackingSignalCache=False,
        numPackets=50,
        memoryBytes=512.0,
        tR=common_pb2.TimeRange(timeRangeSec=[0.0, 50.0], fs=30000.0),
        numSigs=64,
    )
    status = kpi_status_from_proto(pb)
    assert status.num_sigs == 64
    assert status.time_range is not None
    assert status.time_range.sec == (0.0, 50.0)


def test_build_bundle_req() -> None:
    metrics = [KpiMetricId(mode=KpiMode.BASE, name=KpiMetric.MEAN)]
    req = build_bundle_req([0, 1, 2], [0.0, 5.0], metrics)
    assert list(req.ntvIdxs) == [0, 1, 2]
    assert list(req.tR.timeRangeSec) == [0.0, 5.0]
    assert len(req.metrics) == 1
    assert req.metrics[0].name == common_pb2.KPI_MEAN


# =====================================================================
# AllegoClient KPI tests
# =====================================================================


@pytest.fixture
def mock_allego_discovery() -> object:
    with patch("radiens_core.allego_client.AllegoClient._discover") as mock:
        mock.return_value = {
            "ALLEGO_CORE": "localhost:50051",
            "ALLEGO_PCACHE": "localhost:50052",
            "ALLEGO_KPI": "localhost:50053",
        }
        yield mock


@pytest.fixture
def mock_allego_auth() -> object:
    with patch("radiens_core._base_client.AuthInterceptor") as mock:
        yield mock


def test_allego_get_kpi_metrics_builds_correct_request(
    mock_allego_discovery: object,
    mock_allego_auth: object,
) -> None:
    from radiens_core.allego_client import AllegoClient

    client = AllegoClient()
    vals = np.array([[1.0, 2.0]], dtype=np.float64)
    mock_reply = common_pb2.SignalMetricsPacketReply(
        ntvIdxs=[0],
        sigType=common_pb2.PRI,
        metrics=[common_pb2.KpiMetricID(mode=common_pb2.Kpi_MetricsMode_Base, name=common_pb2.KPI_RMS)],
        metricVals=common_pb2.DenseMatrix(rows=1, cols=2, data=vals.tobytes()),
        tR=common_pb2.TimeRange(timeRangeSec=[0.0, 1.0], fs=30000.0),
    )

    with patch("radiens_core.allego_client.kpi_get_metrics_allego_proto", return_value=mock_reply) as mock_rpc:
        result = client.get_kpi_metrics(
            metrics=[KpiMetricId(mode=KpiMode.BASE, name=KpiMetric.RMS)],
            channel_indices=[0],
            time_range=[1.0, float("nan")],
        )

    assert isinstance(result, KpiMetricsResult)
    mock_rpc.assert_called_once()
    _pool, _addr, req = mock_rpc.call_args[0]
    assert req.streamGroupId == PRIMARY_CACHE_STREAM_GROUP_ID
    assert req.stype == common_pb2.PRI
    assert list(req.arg.ntvIdxs) == [0]
    assert req.arg.metrics[0].name == common_pb2.KPI_RMS


def test_allego_get_kpi_status(
    mock_allego_discovery: object,
    mock_allego_auth: object,
) -> None:
    from radiens_core.allego_client import AllegoClient

    client = AllegoClient()
    mock_reply = common_pb2.KpiStatusReply(
        streamGroupId=PRIMARY_CACHE_STREAM_GROUP_ID,
        numPacketsMemory=20,
        packetDur=0.5,
        updatePeriod=0.1,
        persistence=0.9,
        beta=0.1,
        isTrackingSignalCache=True,
        numPackets=200,
        memoryBytes=4096.0,
    )

    with patch("radiens_core.allego_client.get_kpi_status_allego_proto", return_value=mock_reply):
        status = client.get_kpi_status()

    assert isinstance(status, KpiStatus)
    assert status.num_packets_memory == 20


def test_allego_set_kpi_packet_dur_raises_on_error(
    mock_allego_discovery: object,
    mock_allego_auth: object,
) -> None:
    from radiens_core.allego_client import AllegoClient
    from radiens_core.exceptions import RadiensError

    client = AllegoClient()
    mock_reply = common_pb2.StandardReply(errMsg="server error")

    with (
        patch("radiens_core.allego_client.set_kpi_packet_dur_allego_proto", return_value=mock_reply),
        pytest.raises(RadiensError, match="Failed to set KPI packet duration"),
    ):
        client.set_kpi_packet_dur(0.5)


def test_allego_set_kpi_update_period_raises_on_error(
    mock_allego_discovery: object,
    mock_allego_auth: object,
) -> None:
    from radiens_core.allego_client import AllegoClient
    from radiens_core.exceptions import RadiensError

    client = AllegoClient()
    mock_reply = common_pb2.StandardReply(errMsg="bad param")

    with (
        patch("radiens_core.allego_client.set_kpi_update_period_proto", return_value=mock_reply),
        pytest.raises(RadiensError, match="Failed to set KPI update period"),
    ):
        client.set_kpi_update_period(0.1)


# =====================================================================
# VidereClient KPI tests
# =====================================================================


@pytest.fixture
def mock_videre_discovery() -> object:
    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_videre_auth() -> object:
    with patch("radiens_core._base_client.AuthInterceptor") as mock:
        yield mock


def _make_videre_meta() -> DatasetMetadata:
    from radiens_core.models.channels import ChannelMetadata, KeyIdxs, SignalUnits
    from radiens_core.models.dataset import DatasetMetadata, RadiensFileType

    sig_map = {st: KeyIdxs(ntv=[0], dset=[0], sys=[0]) for st in SignalType}
    channel_metadata = ChannelMetadata(
        dataset_uid="uid_001",
        fs=30000.0,
        source_label="test",
        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,
    )
    return DatasetMetadata(
        id="ds_001",
        base_name="recording",
        path="/path/to/file.xdat",
        file_type=RadiensFileType.XDAT,
        time_range=time_range,
        channel_metadata=channel_metadata,
    )


def test_videre_get_kpi_metrics_uses_dataset_id(
    mock_videre_discovery: object,
    mock_videre_auth: object,
) -> None:
    from radiens_core.videre_client import VidereClient

    client = VidereClient()
    meta = _make_videre_meta()
    client._datasets[meta.id] = meta

    vals = np.array([[5.0]], dtype=np.float64)
    mock_reply = common_pb2.SignalMetricsPacketReply(
        ntvIdxs=[0],
        sigType=common_pb2.PRI,
        metrics=[common_pb2.KpiMetricID(mode=common_pb2.Kpi_MetricsMode_Base, name=common_pb2.KPI_MEAN)],
        metricVals=common_pb2.DenseMatrix(rows=1, cols=1, data=vals.tobytes()),
        tR=common_pb2.TimeRange(timeRangeSec=[0.0, 1.0], fs=30000.0),
    )

    with patch("radiens_core.videre_client.kpi_get_metrics_radiens_proto", return_value=mock_reply) as mock_rpc:
        result = client.get_kpi_metrics(
            dataset=meta,
            metrics=[KpiMetricId(mode=KpiMode.BASE, name=KpiMetric.MEAN)],
            channel_indices=[0],
            time_range=[0.0, 1.0],
        )

    assert isinstance(result, KpiMetricsResult)
    mock_rpc.assert_called_once()
    _pool, _addr, req = mock_rpc.call_args[0]
    assert req.streamGroupId == "ds_001"
    assert list(req.arg.ntvIdxs) == [0]


def test_videre_kpi_calculate_uses_dataset_id(
    mock_videre_discovery: object,
    mock_videre_auth: object,
) -> None:
    from radiens_core.videre_client import VidereClient

    client = VidereClient()
    meta = _make_videre_meta()
    client._datasets[meta.id] = meta

    with patch("radiens_core.videre_client.kpi_calculate_proto", return_value=common_pb2.StandardReply()) as mock_rpc:
        client.kpi_calculate(meta)

    mock_rpc.assert_called_once()
    _pool, _addr, req = mock_rpc.call_args[0]
    assert list(req.dsourceID) == ["ds_001"]


def test_videre_kpi_clear_uses_dataset_id(
    mock_videre_discovery: object,
    mock_videre_auth: object,
) -> None:
    from radiens_core.videre_client import VidereClient

    client = VidereClient()
    meta = _make_videre_meta()
    client._datasets[meta.id] = meta

    with patch("radiens_core.videre_client.kpi_clear_proto", return_value=common_pb2.StandardReply()) as mock_rpc:
        client.kpi_clear(meta)

    mock_rpc.assert_called_once()
    _pool, _addr, req = mock_rpc.call_args[0]
    assert list(req.dsourceID) == ["ds_001"]


def test_videre_set_kpi_packet_dur(
    mock_videre_discovery: object,
    mock_videre_auth: object,
) -> None:
    from radiens_core.videre_client import VidereClient

    client = VidereClient()
    meta = _make_videre_meta()
    client._datasets[meta.id] = meta

    with patch(
        "radiens_core.videre_client.set_kpi_packet_dur_radiens_proto",
        return_value=common_pb2.StandardReply(),
    ) as mock_rpc:
        client.set_kpi_packet_dur(meta, 2.0)

    mock_rpc.assert_called_once()
    _pool, _addr, req = mock_rpc.call_args[0]
    assert req.streamGroupId == "ds_001"
    assert req.param == 2.0


def test_videre_kpi_calculate_raises_on_error(
    mock_videre_discovery: object,
    mock_videre_auth: object,
) -> None:
    from radiens_core.exceptions import RadiensError
    from radiens_core.videre_client import VidereClient

    client = VidereClient()
    meta = _make_videre_meta()
    client._datasets[meta.id] = meta

    with (
        patch(
            "radiens_core.videre_client.kpi_calculate_proto",
            return_value=common_pb2.StandardReply(errMsg="compute failed"),
        ),
        pytest.raises(RadiensError, match="KPI calculate failed"),
    ):
        client.kpi_calculate(meta)
