"""Tests for AzureOpenAIEmbeddingClient with mocked API calls."""

import types

import pytest

from agentbyte.llm import AzureOpenAIEmbeddingClient
from agentbyte.llm.types import ModelPricing


class _FakeEmbeddingItem:
    def __init__(self, embedding, index):
        self.embedding = embedding
        self.index = index


class _FakeEmbeddingResponse:
    def __init__(self, data, usage=None, model="text-embedding-3-small"):
        self.data = data
        self.usage = usage
        self.model = model


class _FakeUsage:
    def __init__(self, prompt_tokens=0, total_tokens=0):
        self.prompt_tokens = prompt_tokens
        self.total_tokens = total_tokens


class _FakeAzureClient:
    def __init__(self, response=None):
        self.response = response
        self.last_kwargs = None
        self.embeddings = types.SimpleNamespace(create=self._create)

    async def _create(self, **kwargs):
        self.last_kwargs = kwargs
        return self.response


@pytest.mark.asyncio
async def test_azure_embedding_create_single_returns_result():
    fake_response = _FakeEmbeddingResponse(
        data=[_FakeEmbeddingItem([0.7, 0.8], index=0)]
    )
    fake_client = _FakeAzureClient(response=fake_response)
    client = AzureOpenAIEmbeddingClient(model="embed-prod", client=fake_client)

    result = await client.create("hello")

    assert result.embedding == [0.7, 0.8]
    assert result.embeddings == [[0.7, 0.8]]
    assert result.metadata["input_count"] == 1
    assert fake_client.last_kwargs["model"] == "embed-prod"
    assert fake_client.last_kwargs["input"] == ["hello"]


@pytest.mark.asyncio
async def test_azure_embedding_create_batch_preserves_order_using_index():
    fake_response = _FakeEmbeddingResponse(
        data=[
            _FakeEmbeddingItem([8.0, 8.1], index=1),
            _FakeEmbeddingItem([7.0, 7.1], index=0),
        ]
    )
    fake_client = _FakeAzureClient(response=fake_response)
    client = AzureOpenAIEmbeddingClient(model="embed-prod", client=fake_client)

    result = await client.create_batch(["first", "second"])

    assert result.embedding == [7.0, 7.1]
    assert result.embeddings == [[7.0, 7.1], [8.0, 8.1]]
    assert result.metadata["input_count"] == 2


@pytest.mark.asyncio
async def test_azure_embedding_create_batch_rejects_empty_input():
    fake_client = _FakeAzureClient(response=_FakeEmbeddingResponse(data=[]))
    client = AzureOpenAIEmbeddingClient(model="embed-prod", client=fake_client)

    with pytest.raises(ValueError, match="cannot be empty"):
        await client.create_batch([])


def test_azure_embedding_from_api_key_loads_env(monkeypatch):
    import agentbyte.llm.azure.embedding as module

    class _CtorClient:
        def __init__(self, api_key, azure_endpoint, api_version):
            self.api_key = api_key
            self.azure_endpoint = azure_endpoint
            self.api_version = api_version
            self.embeddings = types.SimpleNamespace(create=lambda **_: None)

    monkeypatch.setenv("AZURE_OPENAI_ENDPOINT", "https://example.openai.azure.com")
    monkeypatch.setenv("AZURE_OPENAI_API_VERSION", "2024-10-21")
    monkeypatch.setenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME", "embed-prod")
    monkeypatch.setenv("AZURE_OPENAI_API_KEY", "azure-key")
    monkeypatch.setattr(module, "AsyncAzureOpenAI", _CtorClient)

    client = AzureOpenAIEmbeddingClient.from_api_key()

    assert isinstance(client, AzureOpenAIEmbeddingClient)
    assert client.model == "embed-prod"
    assert client.client.api_key == "azure-key"


def test_azure_embedding_estimate_cost_known_model():
    fake_client = _FakeAzureClient(response=_FakeEmbeddingResponse(data=[]))
    client = AzureOpenAIEmbeddingClient(model="text-embedding-3-small", client=fake_client)

    # 250K input tokens at $0.02 / 1M = $0.005
    assert client._estimate_cost(prompt_tokens=250_000) == 0.005


def test_azure_embedding_estimate_cost_fallback_unknown_model():
    fake_client = _FakeAzureClient(response=_FakeEmbeddingResponse(data=[]))
    client = AzureOpenAIEmbeddingClient(model="embed-prod", client=fake_client)

    # Fallback to ada pricing: 200K input tokens at $0.10 / 1M = $0.02
    assert client._estimate_cost(prompt_tokens=200_000) == 0.02


def test_azure_embedding_estimate_cost_warns_for_unknown_model(caplog):
    fake_client = _FakeAzureClient(response=_FakeEmbeddingResponse(data=[]))
    client = AzureOpenAIEmbeddingClient(model="embed-prod", client=fake_client)

    with caplog.at_level("WARNING", logger="agentbyte.llm.azure.embedding"):
        _ = client._estimate_cost(prompt_tokens=200_000)

    assert "Model pricing not available for model 'embed-prod'" in caplog.text


def test_azure_embedding_estimate_cost_uses_custom_model_pricing_override():
    fake_client = _FakeAzureClient(response=_FakeEmbeddingResponse(data=[]))
    client = AzureOpenAIEmbeddingClient(
        model="text-embedding-3-small",
        client=fake_client,
        model_pricing=ModelPricing(input_per_1m=1.0, output_per_1m=0.0),
    )

    assert client._estimate_cost(prompt_tokens=200_000) == 0.2


@pytest.mark.asyncio
async def test_azure_embedding_create_batch_includes_usage_cost():
    fake_response = _FakeEmbeddingResponse(
        data=[_FakeEmbeddingItem([0.7, 0.8], index=0)],
        usage=_FakeUsage(prompt_tokens=250_000, total_tokens=250_000),
        model="text-embedding-3-small",
    )
    fake_client = _FakeAzureClient(response=fake_response)
    client = AzureOpenAIEmbeddingClient(model="text-embedding-3-small", client=fake_client)

    result = await client.create_batch(["hello"])

    assert result.embedding == [0.7, 0.8]
    assert result.embeddings == [[0.7, 0.8]]
    assert result.usage.tokens_input == 250_000
    assert result.usage.tokens_output == 0
    assert result.usage.cost_estimate == 0.005
    assert result.metadata["input_count"] == 1
    assert result.metadata["total_tokens"] == 250_000
