"""Tests for the AzureOpenAI wrapper."""

from unittest.mock import MagicMock, patch

import openai as _openai

from llm_observatory.client import AzureOpenAI, _BaseClient


def test_azure_client_inherits_base():
    """AzureOpenAI should be a subclass of _BaseClient."""
    assert issubclass(AzureOpenAI, _BaseClient)


@patch("llm_observatory.client.openai.AzureOpenAI")
def test_azure_client_wraps_azure_openai(mock_azure_cls):
    """AzureOpenAI should instantiate openai.AzureOpenAI under the hood."""
    mock_instance = MagicMock()
    mock_azure_cls.return_value = mock_instance

    client = AzureOpenAI(
        azure_endpoint="https://my-resource.openai.azure.com",
        api_key="azure-key",
        api_version="2024-02-01",
    )

    mock_azure_cls.assert_called_once_with(
        azure_endpoint="https://my-resource.openai.azure.com",
        api_key="azure-key",
        api_version="2024-02-01",
    )
    assert client._client is mock_instance


@patch("llm_observatory.client.openai.AzureOpenAI")
def test_azure_client_has_chat_completions(mock_azure_cls):
    """AzureOpenAI should expose chat.completions.create."""
    mock_azure_cls.return_value = MagicMock()
    client = AzureOpenAI(api_key="k", azure_endpoint="https://x.openai.azure.com", api_version="2024-02-01")

    assert hasattr(client, "chat")
    assert hasattr(client.chat, "completions")
    assert callable(client.chat.completions.create)


@patch("llm_observatory.client.openai.AzureOpenAI")
def test_azure_set_metadata(mock_azure_cls):
    """set_metadata should persist metadata on the client."""
    mock_azure_cls.return_value = MagicMock()
    client = AzureOpenAI(api_key="k", azure_endpoint="https://x.openai.azure.com", api_version="2024-02-01")

    client.set_metadata(user_id="u1", tags=["prod"])
    assert client._metadata == {"user_id": "u1", "tags": ["prod"]}


@patch("llm_observatory.client.openai.AzureOpenAI")
def test_azure_getattr_proxies(mock_azure_cls):
    """Attributes not on the wrapper should proxy to the underlying client."""
    mock_instance = MagicMock()
    mock_instance.models = "models-namespace"
    mock_azure_cls.return_value = mock_instance

    client = AzureOpenAI(api_key="k", azure_endpoint="https://x.openai.azure.com", api_version="2024-02-01")
    assert client.models == "models-namespace"


@patch("llm_observatory.client.get_transport")
@patch("llm_observatory.client.get_config")
@patch("llm_observatory.client.openai.AzureOpenAI")
def test_azure_create_captures_event(mock_azure_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-abc"
    mock_response.model = "gpt-4o"
    mock_response.choices = [
        MagicMock(index=0, finish_reason="stop", message=MagicMock(role="assistant", content="Hello", tool_calls=None))
    ]
    mock_response.usage = MagicMock(prompt_tokens=10, completion_tokens=5, total_tokens=15)

    mock_underlying = MagicMock()
    mock_underlying.chat.completions.create.return_value = mock_response
    mock_azure_cls.return_value = mock_underlying

    mock_transport = MagicMock()
    mock_get_transport.return_value = mock_transport

    client = AzureOpenAI(api_key="k", azure_endpoint="https://x.openai.azure.com", api_version="2024-02-01")
    result = 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-abc"
    assert event["usage"]["total_tokens"] == 15


@patch("llm_observatory.client.get_config")
@patch("llm_observatory.client.openai.AzureOpenAI")
def test_azure_disabled_passthrough(mock_azure_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.return_value = mock_response
    mock_azure_cls.return_value = mock_underlying

    client = AzureOpenAI(api_key="k", azure_endpoint="https://x.openai.azure.com", api_version="2024-02-01")
    result = client.chat.completions.create(model="gpt-4o", messages=[{"role": "user", "content": "Hi"}])

    assert result is mock_response


@patch("llm_observatory.client.get_transport")
@patch("llm_observatory.client.get_config")
@patch("llm_observatory.client.openai.AzureOpenAI")
def test_azure_error_captures_event(mock_azure_cls, mock_get_config, mock_get_transport):
    """On API error, the event should still be captured with error info."""
    mock_get_config.return_value = {"enabled": True, "redact_messages": False}

    mock_underlying = MagicMock()
    mock_request = MagicMock()
    api_error = _openai.APIError("Deployment not found", request=mock_request, body=None)
    mock_underlying.chat.completions.create.side_effect = api_error
    mock_azure_cls.return_value = mock_underlying

    mock_transport = MagicMock()
    mock_get_transport.return_value = mock_transport

    client = AzureOpenAI(api_key="k", azure_endpoint="https://x.openai.azure.com", api_version="2024-02-01")
    try:
        client.chat.completions.create(model="gpt-4o", messages=[{"role": "user", "content": "Hi"}])
    except _openai.APIError:
        pass

    mock_transport.enqueue.assert_called_once()
    event = mock_transport.enqueue.call_args[0][0]
    assert event["error"]["type"] == "APIError"
    assert "Deployment not found" in event["error"]["message"]


@patch("llm_observatory.client.get_transport")
@patch("llm_observatory.client.get_config")
@patch("llm_observatory.client.openai.AzureOpenAI")
def test_azure_auth_error_captures_event(mock_azure_cls, mock_get_config, mock_get_transport):
    """On 401 authentication error, the event should be captured with error info."""
    mock_get_config.return_value = {"enabled": True, "redact_messages": False}

    mock_underlying = MagicMock()
    mock_http_response = MagicMock()
    mock_http_response.status_code = 401
    mock_http_response.headers = {}
    auth_error = _openai.AuthenticationError(
        "Invalid API key", response=mock_http_response, body=None
    )
    mock_underlying.chat.completions.create.side_effect = auth_error
    mock_azure_cls.return_value = mock_underlying

    mock_transport = MagicMock()
    mock_get_transport.return_value = mock_transport

    client = AzureOpenAI(api_key="bad-key", azure_endpoint="https://x.openai.azure.com", api_version="2024-02-01")
    try:
        client.chat.completions.create(model="gpt-4o", messages=[{"role": "user", "content": "Hi"}])
    except _openai.AuthenticationError:
        pass

    mock_transport.enqueue.assert_called_once()
    event = mock_transport.enqueue.call_args[0][0]
    assert event["error"]["type"] == "AuthenticationError"
    assert "Invalid API key" in event["error"]["message"]


@patch("llm_observatory.client.get_transport")
@patch("llm_observatory.client.get_config")
@patch("llm_observatory.client.openai.AzureOpenAI")
def test_azure_rate_limit_error_captures_event(mock_azure_cls, mock_get_config, mock_get_transport):
    """On 429 rate limit error, the event should be captured with error info."""
    mock_get_config.return_value = {"enabled": True, "redact_messages": False}

    mock_underlying = MagicMock()
    mock_http_response = MagicMock()
    mock_http_response.status_code = 429
    mock_http_response.headers = {}
    rate_error = _openai.RateLimitError(
        "Rate limit exceeded", response=mock_http_response, body=None
    )
    mock_underlying.chat.completions.create.side_effect = rate_error
    mock_azure_cls.return_value = mock_underlying

    mock_transport = MagicMock()
    mock_get_transport.return_value = mock_transport

    client = AzureOpenAI(api_key="k", azure_endpoint="https://x.openai.azure.com", api_version="2024-02-01")
    try:
        client.chat.completions.create(model="gpt-4o", messages=[{"role": "user", "content": "Hi"}])
    except _openai.RateLimitError:
        pass

    mock_transport.enqueue.assert_called_once()
    event = mock_transport.enqueue.call_args[0][0]
    assert event["error"]["type"] == "RateLimitError"
    assert "Rate limit" in event["error"]["message"]
