"""
Tests for provider-owned retry observability in chat and embedding clients.

Covers:
- Observability absent when no retries occur
- Correct structure after one retry
- Accumulates correctly across multiple retries
- Both OpenAI and Azure providers
- Both chat and embedding clients
"""

from __future__ import annotations

import types
from typing import Any
from unittest.mock import AsyncMock, patch

import pytest

from agentbyte.llm.azure.chat import AzureOpenAIChatCompletionClient
from agentbyte.llm.azure.embedding import AzureOpenAIEmbeddingClient
from agentbyte.llm.base import RateLimitError
from agentbyte.llm.openai.chat import OpenAIChatCompletionClient
from agentbyte.llm.openai.embedding import OpenAIEmbeddingClient
from agentbyte.messages import UserMessage


# ---------------------------------------------------------------------------
# Shared fake API response builders
# ---------------------------------------------------------------------------

def _fake_chat_response(model: str = "fake") -> Any:
    """Build a minimal fake OpenAI chat response object."""
    choice = types.SimpleNamespace(
        message=types.SimpleNamespace(content="hello", tool_calls=None),
        finish_reason="stop",
    )
    return types.SimpleNamespace(
        choices=[choice],
        model=model,
        usage=types.SimpleNamespace(prompt_tokens=10, completion_tokens=5),
    )


def _fake_embedding_response(model: str = "fake-embed") -> Any:
    """Build a minimal fake OpenAI embedding response object."""
    item = types.SimpleNamespace(embedding=[0.1, 0.2], index=0)
    return types.SimpleNamespace(
        data=[item],
        model=model,
        usage=types.SimpleNamespace(prompt_tokens=5, total_tokens=5),
    )


def _fake_openai_client(chat_response=None, embedding_response=None) -> Any:
    """Build a fake AsyncOpenAI client."""
    client = types.SimpleNamespace()
    if chat_response is not None:
        client.chat = types.SimpleNamespace(
            completions=types.SimpleNamespace(
                create=AsyncMock(return_value=chat_response)
            )
        )
    if embedding_response is not None:
        client.embeddings = types.SimpleNamespace(
            create=AsyncMock(return_value=embedding_response)
        )
    return client


# ---------------------------------------------------------------------------
# Helper: build a client that raises RateLimitError N times then succeeds
# ---------------------------------------------------------------------------

def _flaky_create(success_response: Any, fail_times: int):
    """Return an AsyncMock that raises RateLimitError fail_times then returns success."""
    call_count = 0

    async def _side_effect(*args, **kwargs):
        nonlocal call_count
        call_count += 1
        if call_count <= fail_times:
            raise RateLimitError(f"rate limited #{call_count}")
        return success_response

    return AsyncMock(side_effect=_side_effect)


# ---------------------------------------------------------------------------
# OpenAI Chat client
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_openai_chat_no_retry_observability_absent():
    """retry_observability key must not appear when no retries happened."""
    fake_client = _fake_openai_client(chat_response=_fake_chat_response())
    client = OpenAIChatCompletionClient(model="gpt-4o-mini", client=fake_client)

    with patch("asyncio.sleep", new_callable=AsyncMock):
        result = await client.create([UserMessage(content="hi", source="user")])

    assert "retry_observability" not in result.metadata


@pytest.mark.asyncio
async def test_openai_chat_one_retry_observability_present():
    """After one retry, metadata contains the correct observability structure."""
    chat_resp = _fake_chat_response()
    fake_client = types.SimpleNamespace(
        chat=types.SimpleNamespace(
            completions=types.SimpleNamespace(
                create=_flaky_create(chat_resp, fail_times=1)
            )
        )
    )
    client = OpenAIChatCompletionClient(
        model="gpt-4o-mini", client=fake_client, max_retries=3, initial_retry_delay=0.01
    )

    with patch("asyncio.sleep", new_callable=AsyncMock):
        result = await client.create([UserMessage(content="hi", source="user")])

    obs = result.metadata.get("retry_observability")
    assert obs is not None
    assert len(obs["history"]) == 1
    assert obs["history"][0]["attempt"] == 1
    assert obs["history"][0]["exception_type"] == "RateLimitError"
    assert obs["history"][0]["outcome"] == "retrying"
    assert obs["latest"]["retries"] == 1
    assert obs["latest"]["total_attempts"] == 2


@pytest.mark.asyncio
async def test_openai_chat_multiple_retries_accumulate():
    """After two retries, history has two records."""
    chat_resp = _fake_chat_response()
    fake_client = types.SimpleNamespace(
        chat=types.SimpleNamespace(
            completions=types.SimpleNamespace(
                create=_flaky_create(chat_resp, fail_times=2)
            )
        )
    )
    client = OpenAIChatCompletionClient(
        model="gpt-4o-mini", client=fake_client, max_retries=3, initial_retry_delay=0.01
    )

    with patch("asyncio.sleep", new_callable=AsyncMock):
        result = await client.create([UserMessage(content="hi", source="user")])

    obs = result.metadata["retry_observability"]
    assert len(obs["history"]) == 2
    assert obs["latest"]["retries"] == 2
    assert obs["latest"]["total_attempts"] == 3


# ---------------------------------------------------------------------------
# Azure Chat client
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_azure_chat_no_retry_observability_absent():
    fake_client = _fake_openai_client(chat_response=_fake_chat_response())
    client = AzureOpenAIChatCompletionClient(model="gpt-4o-mini", client=fake_client)

    with patch("asyncio.sleep", new_callable=AsyncMock):
        result = await client.create([UserMessage(content="hi", source="user")])

    assert "retry_observability" not in result.metadata


@pytest.mark.asyncio
async def test_azure_chat_one_retry_observability_present():
    chat_resp = _fake_chat_response()
    fake_client = types.SimpleNamespace(
        chat=types.SimpleNamespace(
            completions=types.SimpleNamespace(
                create=_flaky_create(chat_resp, fail_times=1)
            )
        )
    )
    client = AzureOpenAIChatCompletionClient(
        model="gpt-4o-mini", client=fake_client, max_retries=3, initial_retry_delay=0.01
    )

    with patch("asyncio.sleep", new_callable=AsyncMock):
        result = await client.create([UserMessage(content="hi", source="user")])

    obs = result.metadata.get("retry_observability")
    assert obs is not None
    assert len(obs["history"]) == 1
    assert obs["latest"]["retries"] == 1


# ---------------------------------------------------------------------------
# OpenAI Embedding client
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_openai_embedding_no_retry_observability_absent():
    fake_client = _fake_openai_client(embedding_response=_fake_embedding_response())
    client = OpenAIEmbeddingClient(model="text-embedding-3-small", client=fake_client)

    with patch("asyncio.sleep", new_callable=AsyncMock):
        result = await client.create("hello world")

    assert "retry_observability" not in result.metadata


@pytest.mark.asyncio
async def test_openai_embedding_one_retry_observability_present():
    embed_resp = _fake_embedding_response()
    fake_client = types.SimpleNamespace(
        embeddings=types.SimpleNamespace(
            create=_flaky_create(embed_resp, fail_times=1)
        )
    )
    client = OpenAIEmbeddingClient(
        model="text-embedding-3-small", client=fake_client,
        max_retries=3, initial_retry_delay=0.01
    )

    with patch("asyncio.sleep", new_callable=AsyncMock):
        result = await client.create("hello world")

    obs = result.metadata.get("retry_observability")
    assert obs is not None
    assert len(obs["history"]) == 1
    assert obs["history"][0]["exception_type"] == "RateLimitError"
    assert obs["latest"]["retries"] == 1
    assert obs["latest"]["total_attempts"] == 2


@pytest.mark.asyncio
async def test_openai_embedding_multiple_retries_accumulate():
    embed_resp = _fake_embedding_response()
    fake_client = types.SimpleNamespace(
        embeddings=types.SimpleNamespace(
            create=_flaky_create(embed_resp, fail_times=2)
        )
    )
    client = OpenAIEmbeddingClient(
        model="text-embedding-3-small", client=fake_client,
        max_retries=3, initial_retry_delay=0.01
    )

    with patch("asyncio.sleep", new_callable=AsyncMock):
        result = await client.create("hello world")

    obs = result.metadata["retry_observability"]
    assert len(obs["history"]) == 2
    assert obs["latest"]["retries"] == 2
    assert obs["latest"]["total_attempts"] == 3


# ---------------------------------------------------------------------------
# Azure Embedding client
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_azure_embedding_no_retry_observability_absent():
    fake_client = _fake_openai_client(embedding_response=_fake_embedding_response())
    client = AzureOpenAIEmbeddingClient(model="embed-prod", client=fake_client)

    with patch("asyncio.sleep", new_callable=AsyncMock):
        result = await client.create("hello world")

    assert "retry_observability" not in result.metadata


@pytest.mark.asyncio
async def test_azure_embedding_one_retry_observability_present():
    embed_resp = _fake_embedding_response()
    fake_client = types.SimpleNamespace(
        embeddings=types.SimpleNamespace(
            create=_flaky_create(embed_resp, fail_times=1)
        )
    )
    client = AzureOpenAIEmbeddingClient(
        model="embed-prod", client=fake_client,
        max_retries=3, initial_retry_delay=0.01
    )

    with patch("asyncio.sleep", new_callable=AsyncMock):
        result = await client.create("hello world")

    obs = result.metadata.get("retry_observability")
    assert obs is not None
    assert len(obs["history"]) == 1
    assert obs["latest"]["retries"] == 1


@pytest.mark.asyncio
async def test_azure_embedding_multiple_retries_accumulate():
    embed_resp = _fake_embedding_response()
    fake_client = types.SimpleNamespace(
        embeddings=types.SimpleNamespace(
            create=_flaky_create(embed_resp, fail_times=2)
        )
    )
    client = AzureOpenAIEmbeddingClient(
        model="embed-prod", client=fake_client,
        max_retries=3, initial_retry_delay=0.01
    )

    with patch("asyncio.sleep", new_callable=AsyncMock):
        result = await client.create("hello world")

    obs = result.metadata["retry_observability"]
    assert len(obs["history"]) == 2
    assert obs["latest"]["retries"] == 2
    assert obs["latest"]["total_attempts"] == 3
