"""Tests for CurateClient."""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

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

import pytest

from radiens_core._generated import common_pb2
from radiens_core.curate_client import CurateClient
from radiens_core.models.dataset import RadiensFileType


@pytest.fixture
def mock_discovery() -> Generator[Any, None, None]:  # type: ignore[explicit-any]
    """Mock server discovery."""
    with patch("radiens_core.curate_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_execute_protocol_progress(mock_discovery: Any, mock_auth: Any) -> None:  # type: ignore[explicit-any]
    """Test _execute_protocol with a mocked progress stream."""
    client = CurateClient()

    # Mock set_protocol_proto
    mock_set = MagicMock()

    # Mock apply_protocol_proto to return a stream of progress messages
    mock_apply = MagicMock()

    def progress_stream(*args: Any, **kwargs: Any) -> Generator[common_pb2.ApplyProtocolProgress, None, None]:  # type: ignore[explicit-any]
        yield common_pb2.ApplyProtocolProgress(fracComplete=0.2)
        yield common_pb2.ApplyProtocolProgress(fracComplete=0.5)
        yield common_pb2.ApplyProtocolProgress(fracComplete=1.0)

    mock_apply.return_value = progress_stream()

    with (
        patch("radiens_core.curate_client.set_protocol_proto", mock_set),
        patch("radiens_core.curate_client.apply_protocol_proto", mock_apply),
        patch("radiens_core.curate_client.tqdm") as mock_tqdm,
    ):
        # Mock tqdm instance
        pbar = mock_tqdm.return_value.__enter__.return_value

        nodes = [common_pb2.TransformNode(id="node1")]
        edges = [common_pb2.TransformEdge(id="edge1")]

        client._execute_protocol(nodes, edges, desc="Test Task")

        # Verify set_protocol was called
        mock_set.assert_called_once()
        protocol = mock_set.call_args[0][2]
        assert protocol.transformNodes[0].id == "node1"
        assert protocol.transformEdges[0].id == "edge1"

        # Verify apply_protocol was called
        mock_apply.assert_called_once()

        # Verify tqdm updates
        # Expected updates: 0.2*100=20, (0.5-0.2)*100=30, (1.0-0.5)*100=50
        assert pbar.update.call_count == 3
        calls = pbar.update.call_args_list
        assert calls[0][0][0] == pytest.approx(20.0)
        assert calls[1][0][0] == pytest.approx(30.0)
        assert calls[2][0][0] == pytest.approx(50.0)


def test_link_data_file(mock_discovery: Any, mock_auth: Any) -> None:  # type: ignore[explicit-any]
    """Test link_data_file method calls the correct service."""
    client = CurateClient()
    mock_meta = MagicMock()
    mock_meta.id = "test_ds"

    with patch("radiens_core.curate_client.set_datasource", return_value=mock_meta) as mock_set_ds:
        meta = client.link_data_file("/path/to/test.xdat")

        assert meta == mock_meta
        assert client._datasets["test_ds"] == mock_meta
        mock_set_ds.assert_called_once()
        assert mock_set_ds.call_args[0][2] == "/path/to/test.xdat"


def test_slice_time(mock_discovery: Any, mock_auth: Any) -> None:  # type: ignore[explicit-any]
    """Test slice_time method builds correct graph and executes."""
    client = CurateClient()

    mock_input_meta = MagicMock()
    mock_input_meta.id = "input_ds"
    mock_input_meta.base_name = "test_recording"
    mock_input_meta.path = "/path/to"
    mock_input_meta.file_type = RadiensFileType.XDAT
    client._datasets["input_ds"] = mock_input_meta

    mock_output_meta = MagicMock()

    with (
        patch.object(client, "_execute_protocol") as mock_exec,
        patch.object(client, "link_data_file", return_value=mock_output_meta),
    ):
        res = client.slice_time("input_ds", 10.5, 20.5, "output.xdat")

        assert res == mock_output_meta
        mock_exec.assert_called_once()
        args, kwargs = mock_exec.call_args
        nodes, edges = args
        assert "Slicing output.xdat" in kwargs["desc"]

        # Verify nodes: Source, Transform, Sink
        assert len(nodes) == 3
        assert nodes[0].type == common_pb2.SOURCE
        assert nodes[0].datasourceSinkParams.id == "input_ds"

        assert nodes[1].type == common_pb2.SLICE_TIME
        assert nodes[1].sliceTimeParams.timeStart == 10.5
        assert nodes[1].sliceTimeParams.timeEnd == 20.5

        assert nodes[2].type == common_pb2.SINK
        assert nodes[2].datasourceSinkParams.dsrcName == "output"

        # Verify edges
        assert len(edges) == 2
        assert edges[0].source == nodes[0].id
        assert edges[0].target == nodes[1].id
        assert edges[1].source == nodes[1].id
        assert edges[1].target == nodes[2].id


def test_downsample(mock_discovery: Any, mock_auth: Any) -> None:  # type: ignore[explicit-any]
    """Test downsample method builds correct graph and executes."""
    client = CurateClient()

    mock_input_meta = MagicMock()
    mock_input_meta.id = "input_ds"
    mock_input_meta.base_name = "test_recording"
    mock_input_meta.path = "/path/to"
    mock_input_meta.file_type = RadiensFileType.XDAT
    client._datasets["input_ds"] = mock_input_meta

    mock_output_meta = MagicMock()

    with (
        patch.object(client, "_execute_protocol") as mock_exec,
        patch.object(client, "link_data_file", return_value=mock_output_meta),
    ):
        res = client.downsample("input_ds", 4, "downsampled.rhd")

        assert res == mock_output_meta
        mock_exec.assert_called_once()
        args, kwargs = mock_exec.call_args
        nodes, _edges = args
        assert "Downsampling downsampled.rhd" in kwargs["desc"]

        assert len(nodes) == 3
        assert nodes[1].type == common_pb2.DOWNSAMPLE
        assert nodes[1].downsampleParams.sampleFactor == 4

        assert nodes[2].datasourceSinkParams.fileType == common_pb2.RHD
