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

import types

import pytest

from agentbyte.llm import OpenAIEmbeddingClient
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 _FakeOpenAIClient:
    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_openai_embedding_create_single_returns_result():
    fake_response = _FakeEmbeddingResponse(
        data=[_FakeEmbeddingItem([0.1, 0.2, 0.3], index=0)]
    )
    fake_client = _FakeOpenAIClient(response=fake_response)
    client = OpenAIEmbeddingClient(model="text-embedding-3-small", client=fake_client)

    result = await client.create("hello")

    assert result.embedding == [0.1, 0.2, 0.3]
    assert result.embeddings == [[0.1, 0.2, 0.3]]
    assert result.metadata["input_count"] == 1
    assert fake_client.last_kwargs["model"] == "text-embedding-3-small"
    assert fake_client.last_kwargs["input"] == ["hello"]


@pytest.mark.asyncio
async def test_openai_embedding_create_batch_preserves_order_using_index():
    fake_response = _FakeEmbeddingResponse(
        data=[
            _FakeEmbeddingItem([2.0, 2.1], index=1),
            _FakeEmbeddingItem([1.0, 1.1], index=0),
        ]
    )
    fake_client = _FakeOpenAIClient(response=fake_response)
    client = OpenAIEmbeddingClient(model="text-embedding-3-small", client=fake_client)

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

    assert result.embedding == [1.0, 1.1]
    assert result.embeddings == [[1.0, 1.1], [2.0, 2.1]]
    assert result.metadata["input_count"] == 2


@pytest.mark.asyncio
async def test_openai_embedding_create_batch_rejects_empty_input():
    fake_client = _FakeOpenAIClient(response=_FakeEmbeddingResponse(data=[]))
    client = OpenAIEmbeddingClient(model="text-embedding-3-small", client=fake_client)

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


def test_openai_embedding_from_api_key_loads_env(monkeypatch):
    import agentbyte.llm.openai.embedding as module

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

    monkeypatch.setenv("OPENAI_API_KEY", "sk-test")
    monkeypatch.setenv("OPENAI_ORG", "org-test")
    monkeypatch.setattr(module, "AsyncOpenAI", _CtorClient)

    client = OpenAIEmbeddingClient.from_api_key(model="text-embedding-3-small")

    assert isinstance(client, OpenAIEmbeddingClient)
    assert client.client.api_key == "sk-test"
    assert client.client.organization == "org-test"


def test_openai_embedding_estimate_cost_known_model():
    fake_client = _FakeOpenAIClient(response=_FakeEmbeddingResponse(data=[]))
    client = OpenAIEmbeddingClient(model="text-embedding-3-large", client=fake_client)

    # 100K input tokens at $0.13 / 1M = $0.013
    assert client._estimate_cost(prompt_tokens=100_000) == 0.013


def test_openai_embedding_estimate_cost_fallback_unknown_model():
    fake_client = _FakeOpenAIClient(response=_FakeEmbeddingResponse(data=[]))
    client = OpenAIEmbeddingClient(model="custom-embed-deployment", client=fake_client)

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


def test_openai_embedding_estimate_cost_warns_for_unknown_model(caplog):
    fake_client = _FakeOpenAIClient(response=_FakeEmbeddingResponse(data=[]))
    client = OpenAIEmbeddingClient(model="custom-embed-deployment", client=fake_client)

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

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


def test_openai_embedding_estimate_cost_uses_custom_model_pricing_override():
    fake_client = _FakeOpenAIClient(response=_FakeEmbeddingResponse(data=[]))
    client = OpenAIEmbeddingClient(
        model="text-embedding-3-large",
        client=fake_client,
        model_pricing=ModelPricing(input_per_1m=1.0, output_per_1m=0.0),
    )

    assert client._estimate_cost(prompt_tokens=100_000) == 0.1


@pytest.mark.asyncio
async def test_openai_embedding_create_batch_includes_usage_cost():
    fake_response = _FakeEmbeddingResponse(
        data=[_FakeEmbeddingItem([0.1, 0.2], index=0)],
        usage=_FakeUsage(prompt_tokens=100_000, total_tokens=100_000),
        model="text-embedding-3-large",
    )
    fake_client = _FakeOpenAIClient(response=fake_response)
    client = OpenAIEmbeddingClient(model="text-embedding-3-large", client=fake_client)

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

    assert result.embedding == [0.1, 0.2]
    assert result.embeddings == [[0.1, 0.2]]
    assert result.usage.tokens_input == 100_000
    assert result.usage.tokens_output == 0
    assert result.usage.cost_estimate == 0.013
    assert result.metadata["input_count"] == 1
    assert result.metadata["total_tokens"] == 100_000
