"""
Tests for InferenceClient.
"""

from unittest.mock import MagicMock, patch

import grpc
import pytest

from daisi.clients.inference import InferenceClient
from daisi.config import DaisiConfig
from daisi.exceptions import DaisiError
from protos.v1.models import InferenceModels_pb2


def create_grpc_error(details: str = "Test error") -> grpc.RpcError:
    """Create a mock gRPC error with details method."""
    error = grpc.RpcError()
    error.details = MagicMock(return_value=details)
    return error


class TestInferenceClientInitialization:
    """Test InferenceClient initialization."""

    def test_init_with_config(self) -> None:
        """Test initialization with config."""
        config = DaisiConfig(client_key="test-key")
        client = InferenceClient(config)
        assert client.config == config
        assert client._stub is None
        assert client.session_id is None
        assert client.inference_id is None

    def test_init_with_session_id(self) -> None:
        """Test initialization with session ID."""
        config = DaisiConfig(client_key="test-key")
        client = InferenceClient(config, session_id="session-123")
        assert client.session_id == "session-123"


class TestInferenceClientCreate:
    """Test InferenceClient create method."""

    def test_create_inference_minimal(self) -> None:
        """Test creating inference with minimal parameters."""
        config = DaisiConfig(client_key="test-key")
        client = InferenceClient(config, session_id="session-123")

        mock_stub = MagicMock()
        mock_response = MagicMock()
        mock_response.SessionId = "session-123"
        mock_response.InferenceId = "inference-456"
        mock_stub.Create.return_value = mock_response

        client._stub = mock_stub
        session_id, inference_id = client.create()

        assert session_id == "session-123"
        assert inference_id == "inference-456"
        assert client.inference_id == "inference-456"

    def test_create_inference_with_params(self) -> None:
        """Test creating inference with full parameters."""
        config = DaisiConfig(client_key="test-key")
        client = InferenceClient(config)

        mock_stub = MagicMock()
        mock_response = MagicMock()
        mock_response.SessionId = "session-123"
        mock_response.InferenceId = "inference-456"
        mock_stub.Create.return_value = mock_response

        client._stub = mock_stub
        session_id, inference_id = client.create(
            session_id="session-123",
            model_name="llama-3",
            initialization_prompt="You are a helpful assistant",
            think_level=InferenceModels_pb2.ThinkChainOfThought,
            tool_groups=[InferenceModels_pb2.InferenceCodingTools],
        )

        call_args = mock_stub.Create.call_args
        request = call_args[0][0]
        assert request.ModelName == "llama-3"
        assert request.InitializationPrompt == "You are a helpful assistant"
        assert request.ThinkLevel == InferenceModels_pb2.ThinkChainOfThought

    def test_create_inference_no_session_id(self) -> None:
        """Test creating inference without session ID raises error."""
        config = DaisiConfig(client_key="test-key")
        client = InferenceClient(config)

        with pytest.raises(DaisiError, match="No session ID provided or available"):
            client.create()

    def test_create_inference_grpc_error(self) -> None:
        """Test create inference with gRPC error."""
        config = DaisiConfig(client_key="test-key")
        client = InferenceClient(config, session_id="session-123")

        mock_stub = MagicMock()
        mock_stub.Create.side_effect = create_grpc_error()

        client._stub = mock_stub
        with pytest.raises(DaisiError, match="Failed to create inference"):
            client.create()


class TestInferenceClientSend:
    """Test InferenceClient send method."""

    def test_send_inference_basic(self) -> None:
        """Test sending basic inference request."""
        config = DaisiConfig(client_key="test-key")
        client = InferenceClient(config, session_id="session-123")
        client.inference_id = "inference-456"

        mock_stub = MagicMock()
        mock_response1 = MagicMock()
        mock_response2 = MagicMock()
        mock_stub.Send.return_value = iter([mock_response1, mock_response2])

        client._stub = mock_stub
        responses = list(client.send("Hello"))

        assert len(responses) == 2
        mock_stub.Send.assert_called_once()

    def test_send_inference_with_params(self) -> None:
        """Test sending inference with sampling parameters."""
        config = DaisiConfig(client_key="test-key")
        client = InferenceClient(config, session_id="session-123")
        client.inference_id = "inference-456"

        mock_stub = MagicMock()
        mock_stub.Send.return_value = iter([])

        client._stub = mock_stub
        list(
            client.send(
                "Hello",
                temperature=0.7,
                top_p=0.9,
                top_k=50,
                max_tokens=2000,
            )
        )

        call_args = mock_stub.Send.call_args
        request = call_args[0][0]
        assert request.Text == "Hello"
        assert abs(request.Temperature - 0.7) < 0.01
        assert abs(request.TopP - 0.9) < 0.01
        assert request.TopK == 50
        assert request.MaxTokens == 2000

    def test_send_inference_no_session_id(self) -> None:
        """Test sending inference without session ID raises error."""
        config = DaisiConfig(client_key="test-key")
        client = InferenceClient(config)
        client.inference_id = "inference-456"

        with pytest.raises(DaisiError, match="No session ID provided or available"):
            list(client.send("Hello"))

    def test_send_inference_no_inference_id(self) -> None:
        """Test sending inference without inference ID raises error."""
        config = DaisiConfig(client_key="test-key")
        client = InferenceClient(config, session_id="session-123")

        with pytest.raises(DaisiError, match="No inference ID provided or available"):
            list(client.send("Hello"))

    def test_send_inference_grpc_error(self) -> None:
        """Test send inference with gRPC error."""
        config = DaisiConfig(client_key="test-key")
        client = InferenceClient(config, session_id="session-123")
        client.inference_id = "inference-456"

        mock_stub = MagicMock()
        mock_stub.Send.side_effect = create_grpc_error()

        client._stub = mock_stub
        with pytest.raises(DaisiError, match="Failed to send inference"):
            list(client.send("Hello"))


class TestInferenceClientStats:
    """Test InferenceClient stats method."""

    def test_stats_success(self) -> None:
        """Test getting inference stats."""
        config = DaisiConfig(client_key="test-key")
        client = InferenceClient(config, session_id="session-123")
        client.inference_id = "inference-456"

        mock_stub = MagicMock()
        mock_response = MagicMock()
        mock_stub.Stats.return_value = mock_response

        client._stub = mock_stub
        result = client.stats()

        assert result == mock_response
        mock_stub.Stats.assert_called_once()

    def test_stats_no_session_id(self) -> None:
        """Test stats without session ID raises error."""
        config = DaisiConfig(client_key="test-key")
        client = InferenceClient(config)
        client.inference_id = "inference-456"

        with pytest.raises(DaisiError, match="No session ID provided or available"):
            client.stats()

    def test_stats_no_inference_id(self) -> None:
        """Test stats without inference ID raises error."""
        config = DaisiConfig(client_key="test-key")
        client = InferenceClient(config, session_id="session-123")

        with pytest.raises(DaisiError, match="No inference ID provided or available"):
            client.stats()

    def test_stats_grpc_error(self) -> None:
        """Test stats with gRPC error."""
        config = DaisiConfig(client_key="test-key")
        client = InferenceClient(config, session_id="session-123")
        client.inference_id = "inference-456"

        mock_stub = MagicMock()
        mock_stub.Stats.side_effect = create_grpc_error()

        client._stub = mock_stub
        with pytest.raises(DaisiError, match="Failed to get inference stats"):
            client.stats()


class TestInferenceClientCloseInference:
    """Test InferenceClient close_inference method."""

    def test_close_inference_success(self) -> None:
        """Test closing inference."""
        config = DaisiConfig(client_key="test-key")
        client = InferenceClient(config, session_id="session-123")
        client.inference_id = "inference-456"

        mock_stub = MagicMock()
        mock_response = MagicMock()
        mock_response.Success = True
        mock_stub.Close.return_value = mock_response

        client._stub = mock_stub
        result = client.close_inference()

        assert result is True
        assert client.inference_id is None

    def test_close_inference_with_reason(self) -> None:
        """Test closing inference with reason."""
        config = DaisiConfig(client_key="test-key")
        client = InferenceClient(config, session_id="session-123")
        client.inference_id = "inference-456"

        mock_stub = MagicMock()
        mock_response = MagicMock()
        mock_response.Success = True
        mock_stub.Close.return_value = mock_response

        client._stub = mock_stub
        client.close_inference(reason=InferenceModels_pb2.CloseRequestedByClient)

        call_args = mock_stub.Close.call_args
        request = call_args[0][0]
        assert request.Reason == InferenceModels_pb2.CloseRequestedByClient

    def test_close_inference_no_session_id(self) -> None:
        """Test close inference without session ID raises error."""
        config = DaisiConfig(client_key="test-key")
        client = InferenceClient(config)
        client.inference_id = "inference-456"

        with pytest.raises(DaisiError, match="No session ID provided or available"):
            client.close_inference()

    def test_close_inference_no_inference_id(self) -> None:
        """Test close inference without inference ID raises error."""
        config = DaisiConfig(client_key="test-key")
        client = InferenceClient(config, session_id="session-123")

        with pytest.raises(DaisiError, match="No inference ID provided or available"):
            client.close_inference()

    def test_close_inference_grpc_error(self) -> None:
        """Test close inference with gRPC error."""
        config = DaisiConfig(client_key="test-key")
        client = InferenceClient(config, session_id="session-123")
        client.inference_id = "inference-456"

        mock_stub = MagicMock()
        mock_stub.Close.side_effect = create_grpc_error()

        client._stub = mock_stub
        with pytest.raises(DaisiError, match="Failed to close inference"):
            client.close_inference()
