"""
Tests for AsyncInferenceClient.
"""

from unittest.mock import AsyncMock, MagicMock, patch

import grpc
import pytest


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


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


class TestAsyncInferenceClientInitialization:
    """Test AsyncInferenceClient initialization."""

    def test_init_with_config(self) -> None:
        """Test initialization with config."""
        config = DaisiConfig(client_key="test-key")
        client = AsyncInferenceClient(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 = AsyncInferenceClient(config, session_id="session-123")
        assert client.session_id == "session-123"


class TestAsyncInferenceClientCreate:
    """Test AsyncInferenceClient create method."""

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

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

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

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

    @pytest.mark.asyncio
    async def test_create_inference_with_params(self) -> None:
        """Test creating inference with full parameters."""
        config = DaisiConfig(client_key="test-key")
        client = AsyncInferenceClient(config)

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

        client._stub = mock_stub
        session_id, inference_id = await 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

    @pytest.mark.asyncio
    async def test_create_inference_no_session_id(self) -> None:
        """Test creating inference without session ID raises error."""
        config = DaisiConfig(client_key="test-key")
        client = AsyncInferenceClient(config)

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

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

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

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


class TestAsyncInferenceClientSend:
    """Test AsyncInferenceClient send method."""

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

        mock_stub = MagicMock()
        mock_response1 = MagicMock()
        mock_response2 = MagicMock()

        async def async_generator() -> None:
            yield mock_response1
            yield mock_response2

        mock_stub.Send = MagicMock(return_value=async_generator())

        client._stub = mock_stub
        responses = [r async for r in client.send("Hello")]

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

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

        mock_stub = MagicMock()

        async def async_generator() -> None:
            return
            yield  # Make it a generator

        mock_stub.Send = MagicMock(return_value=async_generator())

        client._stub = mock_stub
        _ = [
            r
            async for r in 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

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

        with pytest.raises(DaisiError, match="No session ID provided or available"):
            _ = [r async for r in client.send("Hello")]

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

        with pytest.raises(DaisiError, match="No inference ID provided or available"):
            _ = [r async for r in client.send("Hello")]

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

        mock_stub = MagicMock()

        async def async_generator_error() -> None:
            raise create_grpc_error()
            yield  # Make it a generator

        mock_stub.Send = MagicMock(return_value=async_generator_error())

        client._stub = mock_stub
        with pytest.raises(DaisiError, match="Failed to send inference"):
            _ = [r async for r in client.send("Hello")]


class TestAsyncInferenceClientStats:
    """Test AsyncInferenceClient stats method."""

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

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

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

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

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

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

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

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

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

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

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


class TestAsyncInferenceClientCloseInference:
    """Test AsyncInferenceClient close_inference method."""

    @pytest.mark.asyncio
    async def test_close_inference_success(self) -> None:
        """Test closing inference."""
        config = DaisiConfig(client_key="test-key")
        client = AsyncInferenceClient(config, session_id="session-123")
        client.inference_id = "inference-456"

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

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

        assert result is True
        assert client.inference_id is None

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

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

        client._stub = mock_stub
        await 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

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

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

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

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

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

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

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