"""Tests for AsyncOpenAI and AsyncAzureOpenAI wrappers."""

from unittest.mock import AsyncMock, MagicMock, patch

import pytest
import openai as _openai

from llm_observatory.client import AsyncOpenAI, AsyncAzureOpenAI, _AsyncBaseClient


def test_async_openai_inherits_async_base():
    """AsyncOpenAI should be a subclass of _AsyncBaseClient."""
    assert issubclass(AsyncOpenAI, _AsyncBaseClient)


def test_async_azure_inherits_async_base():
    """AsyncAzureOpenAI should be a subclass of _AsyncBaseClient."""
    assert issubclass(AsyncAzureOpenAI, _AsyncBaseClient)


@patch("llm_observatory.client.openai.AsyncOpenAI")
def test_async_openai_wraps_client(mock_cls):
    """AsyncOpenAI should instantiate openai.AsyncOpenAI under the hood."""
    mock_instance = MagicMock()
    mock_cls.return_value = mock_instance

    client = AsyncOpenAI(api_key="sk-test")
    mock_cls.assert_called_once_with(api_key="sk-test")
    assert client._client is mock_instance


@patch("llm_observatory.client.openai.AsyncAzureOpenAI")
def test_async_azure_wraps_client(mock_cls):
    """AsyncAzureOpenAI should instantiate openai.AsyncAzureOpenAI under the hood."""
    mock_instance = MagicMock()
    mock_cls.return_value = mock_instance

    client = AsyncAzureOpenAI(
        azure_endpoint="https://x.openai.azure.com",
        api_key="k",
        api_version="2024-02-01",
    )
    mock_cls.assert_called_once_with(
        azure_endpoint="https://x.openai.azure.com",
        api_key="k",
        api_version="2024-02-01",
    )
    assert client._client is mock_instance


@patch("llm_observatory.client.openai.AsyncOpenAI")
def test_async_openai_has_chat_completions(mock_cls):
    """AsyncOpenAI should expose chat.completions.create."""
    mock_cls.return_value = MagicMock()
    client = AsyncOpenAI(api_key="sk-test")
    assert hasattr(client, "chat")
    assert hasattr(client.chat, "completions")


@patch("llm_observatory.client.openai.AsyncOpenAI")
def test_async_openai_has_responses(mock_cls):
    """AsyncOpenAI should expose responses.create."""
    mock_cls.return_value = MagicMock()
    client = AsyncOpenAI(api_key="sk-test")
    assert hasattr(client, "responses")


@patch("llm_observatory.client.openai.AsyncOpenAI")
def test_async_openai_set_metadata(mock_cls):
    """set_metadata should persist metadata."""
    mock_cls.return_value = MagicMock()
    client = AsyncOpenAI(api_key="sk-test")
    client.set_metadata(user_id="u1")
    assert client._metadata == {"user_id": "u1"}


@patch("llm_observatory.client.openai.AsyncOpenAI")
def test_async_openai_getattr_proxies(mock_cls):
    """Unknown attributes should proxy to the underlying client."""
    mock_instance = MagicMock()
    mock_instance.embeddings = "embed-ns"
    mock_cls.return_value = mock_instance

    client = AsyncOpenAI(api_key="sk-test")
    assert client.embeddings == "embed-ns"


@pytest.mark.asyncio
@patch("llm_observatory.client.get_transport")
@patch("llm_observatory.client.get_config")
@patch("llm_observatory.client.openai.AsyncOpenAI")
async def test_async_openai_create_captures_event(mock_cls, mock_get_config, mock_get_transport):
    """chat.completions.create should capture an event via transport."""
    mock_get_config.return_value = {"enabled": True, "redact_messages": False}

    mock_response = MagicMock()
    mock_response.id = "chatcmpl-async"
    mock_response.model = "gpt-4o"
    mock_response.choices = [
        MagicMock(index=0, finish_reason="stop", message=MagicMock(role="assistant", content="Hi!", tool_calls=None))
    ]
    mock_response.usage = MagicMock(prompt_tokens=10, completion_tokens=5, total_tokens=15)

    mock_underlying = MagicMock()
    mock_underlying.chat.completions.create = AsyncMock(return_value=mock_response)
    mock_cls.return_value = mock_underlying

    mock_transport = MagicMock()
    mock_get_transport.return_value = mock_transport

    client = AsyncOpenAI(api_key="sk-test")
    result = await client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": "Hi"}],
    )

    assert result is mock_response
    mock_transport.enqueue.assert_called_once()

    event = mock_transport.enqueue.call_args[0][0]
    assert event["request"]["model"] == "gpt-4o"
    assert event["response"]["id"] == "chatcmpl-async"
    assert event["usage"]["total_tokens"] == 15


@pytest.mark.asyncio
@patch("llm_observatory.client.get_config")
@patch("llm_observatory.client.openai.AsyncOpenAI")
async def test_async_openai_disabled_passthrough(mock_cls, mock_get_config):
    """When disabled, create should pass through without capturing."""
    mock_get_config.return_value = {"enabled": False}

    mock_response = MagicMock()
    mock_underlying = MagicMock()
    mock_underlying.chat.completions.create = AsyncMock(return_value=mock_response)
    mock_cls.return_value = mock_underlying

    client = AsyncOpenAI(api_key="sk-test")
    result = await client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": "Hi"}],
    )

    assert result is mock_response


@pytest.mark.asyncio
@patch("llm_observatory.client.get_transport")
@patch("llm_observatory.client.get_config")
@patch("llm_observatory.client.openai.AsyncOpenAI")
async def test_async_openai_error_captures_event(mock_cls, mock_get_config, mock_get_transport):
    """On API error, the event should still be captured."""
    mock_get_config.return_value = {"enabled": True, "redact_messages": False}

    mock_underlying = MagicMock()
    mock_request = MagicMock()
    api_error = _openai.APIError("Rate limit exceeded", request=mock_request, body=None)
    mock_underlying.chat.completions.create = AsyncMock(side_effect=api_error)
    mock_cls.return_value = mock_underlying

    mock_transport = MagicMock()
    mock_get_transport.return_value = mock_transport

    client = AsyncOpenAI(api_key="sk-test")
    with pytest.raises(_openai.APIError):
        await client.chat.completions.create(
            model="gpt-4o",
            messages=[{"role": "user", "content": "Hi"}],
        )

    mock_transport.enqueue.assert_called_once()
    event = mock_transport.enqueue.call_args[0][0]
    assert event["error"]["type"] == "APIError"


@pytest.mark.asyncio
@patch("llm_observatory.client.get_transport")
@patch("llm_observatory.client.get_config")
@patch("llm_observatory.client.openai.AsyncOpenAI")
async def test_async_openai_per_request_metadata(mock_cls, mock_get_config, mock_get_transport):
    """Per-request metadata should override client-level metadata."""
    mock_get_config.return_value = {"enabled": True, "redact_messages": False}

    mock_response = MagicMock()
    mock_response.id = "chatcmpl-meta"
    mock_response.model = "gpt-4o"
    mock_response.choices = [
        MagicMock(index=0, finish_reason="stop", message=MagicMock(role="assistant", content="Ok", tool_calls=None))
    ]
    mock_response.usage = MagicMock(prompt_tokens=10, completion_tokens=5, total_tokens=15)

    mock_underlying = MagicMock()
    mock_underlying.chat.completions.create = AsyncMock(return_value=mock_response)
    mock_cls.return_value = mock_underlying

    mock_transport = MagicMock()
    mock_get_transport.return_value = mock_transport

    client = AsyncOpenAI(api_key="sk-test")
    client.set_metadata(user_id="global-user")
    await client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": "Hi"}],
        metadata={"user_id": "request-user"},
    )

    event = mock_transport.enqueue.call_args[0][0]
    assert event["metadata"]["user_id"] == "request-user"


@pytest.mark.asyncio
@patch("llm_observatory.client.get_transport")
@patch("llm_observatory.client.get_config")
@patch("llm_observatory.client.openai.AsyncOpenAI")
async def test_async_openai_streaming(mock_cls, mock_get_config, mock_get_transport):
    """Async streaming should return a StreamWrapper."""
    mock_get_config.return_value = {"enabled": True, "redact_messages": False}

    from llm_observatory.streaming import StreamWrapper

    mock_stream = MagicMock()
    mock_underlying = MagicMock()
    mock_underlying.chat.completions.create = AsyncMock(return_value=mock_stream)
    mock_cls.return_value = mock_underlying

    mock_transport = MagicMock()
    mock_get_transport.return_value = mock_transport

    client = AsyncOpenAI(api_key="sk-test")
    result = await client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": "Hi"}],
        stream=True,
    )

    assert isinstance(result, StreamWrapper)
