"""
Tests for OpenAIChatCompletionClient with mocked API calls.

Run with: uv run pytest tests/llm/test_openai_client.py -v
"""
import types

import pytest
from pydantic import BaseModel

from agentbyte.llm import OpenAIChatCompletionClient
from agentbyte.llm.types import ModelPricing
from agentbyte.messages import MultiModalMessage, UserMessage


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


class _FakeFunction:
    def __init__(self, name, arguments):
        self.name = name
        self.arguments = arguments


class _FakeToolCall:
    def __init__(self, call_id, function):
        self.id = call_id
        self.function = function


class _FakeMessage:
    def __init__(self, content, tool_calls=None):
        self.content = content
        self.tool_calls = tool_calls


class _FakeChoice:
    def __init__(self, message, finish_reason="stop"):
        self.message = message
        self.finish_reason = finish_reason


class _FakeResponse:
    def __init__(self, model, choices, usage):
        self.model = model
        self.choices = choices
        self.usage = usage


class _FakeDelta:
    def __init__(self, content=None, tool_calls=None):
        self.content = content
        self.tool_calls = tool_calls


class _FakeToolCallDeltaFunction:
    def __init__(self, name=None, arguments=None):
        self.name = name
        self.arguments = arguments


class _FakeToolCallDelta:
    def __init__(self, call_id=None, index=None, function=None):
        self.id = call_id
        self.index = index
        self.function = function


class _FakeChunkChoice:
    def __init__(self, delta, finish_reason=None):
        self.delta = delta
        self.finish_reason = finish_reason


class _FakeChunk:
    def __init__(self, model, choices=None, usage=None):
        self.model = model
        self.choices = choices or []
        self.usage = usage


class _FakeStream:
    def __init__(self, chunks):
        self._chunks = chunks

    def __aiter__(self):
        return self._iterate()

    async def _iterate(self):
        for chunk in self._chunks:
            yield chunk


class _FakeOpenAIClient:
    def __init__(self, response=None, stream=None):
        self.response = response
        self.stream = stream
        self.last_kwargs = None
        self.chat = types.SimpleNamespace(
            completions=types.SimpleNamespace(create=self._create)
        )

    async def _create(self, **kwargs):
        self.last_kwargs = kwargs
        if kwargs.get("stream"):
            return self.stream
        return self.response


class _OutputModel(BaseModel):
    answer: str


@pytest.mark.asyncio
async def test_openai_create_parses_tool_calls_and_usage():
    fake_response = _FakeResponse(
        model="gpt-4.1-mini",
        choices=[
            _FakeChoice(
                _FakeMessage(
                    content="result",
                    tool_calls=[
                        _FakeToolCall(
                            "call_1",
                            _FakeFunction("calculator", "{\"x\": 1}"),
                        )
                    ],
                ),
                finish_reason="stop",
            )
        ],
        usage=_FakeUsage(prompt_tokens=12, completion_tokens=5),
    )
    fake_client = _FakeOpenAIClient(response=fake_response)
    client = OpenAIChatCompletionClient(model="gpt-4.1-mini", client=fake_client)

    result = await client.create(messages=[UserMessage(content="hi", source="user")])

    assert result.message.content == "result"
    assert result.finish_reason == "stop"
    assert result.usage.tokens_input == 12
    assert result.usage.tokens_output == 5
    assert result.usage.llm_calls == 1
    assert result.usage.tool_calls == 1
    assert result.usage.cost_estimate is not None


@pytest.mark.asyncio
async def test_openai_create_sets_response_format_when_output_format_used():
    fake_response = _FakeResponse(
        model="gpt-4.1-mini",
        choices=[_FakeChoice(_FakeMessage(content="{\"answer\": \"ok\"}"))],
        usage=_FakeUsage(prompt_tokens=1, completion_tokens=1),
    )
    fake_client = _FakeOpenAIClient(response=fake_response)
    client = OpenAIChatCompletionClient(model="gpt-4.1-mini", client=fake_client)

    await client.create(
        messages=[UserMessage(content="hi", source="user")],
        output_format=_OutputModel,
    )

    assert "response_format" in fake_client.last_kwargs


@pytest.mark.asyncio
async def test_openai_create_stream_yields_usage_chunk():
    chunks = [
        _FakeChunk(
            model="gpt-4.1-mini",
            choices=[_FakeChunkChoice(_FakeDelta(content="Hello"))],
        ),
        _FakeChunk(
            model="gpt-4.1-mini",
            choices=[],
            usage=_FakeUsage(prompt_tokens=4, completion_tokens=2),
        ),
    ]
    fake_stream = _FakeStream(chunks)
    fake_client = _FakeOpenAIClient(stream=fake_stream)
    client = OpenAIChatCompletionClient(model="gpt-4.1-mini", client=fake_client)

    received = []
    async for chunk in client.create_stream(
        messages=[UserMessage(content="hi", source="user")],
        stream_options={"include_usage": True},
    ):
        received.append(chunk)

    assert received[0].content == "Hello"
    assert received[-1].is_complete is True
    assert received[-1].usage is not None


@pytest.mark.asyncio
async def test_openai_create_handles_multimodal_messages():
    fake_response = _FakeResponse(
        model="gpt-4.1-mini",
        choices=[_FakeChoice(_FakeMessage(content="ok"))],
        usage=_FakeUsage(prompt_tokens=1, completion_tokens=1),
    )
    fake_client = _FakeOpenAIClient(response=fake_response)
    client = OpenAIChatCompletionClient(model="gpt-4.1-mini", client=fake_client)

    multimodal = MultiModalMessage(
        content="See this",
        source="user",
        role="user",
        mime_type="image/png",
        media_url="https://example.com/image.png",
    )
    await client.create(messages=[multimodal])

    sent = fake_client.last_kwargs["messages"][0]
    assert isinstance(sent["content"], list)
    assert sent["content"][0]["type"] == "text"
    assert sent["content"][1]["type"] == "image_url"


def test_openai_estimate_cost_warns_for_unknown_model(caplog):
    fake_client = _FakeOpenAIClient(response=None)
    client = OpenAIChatCompletionClient(model="custom-openai-model", client=fake_client)

    with caplog.at_level("WARNING", logger="agentbyte.llm.openai.chat"):
        cost = client._estimate_cost(prompt_tokens=100_000, completion_tokens=0)

    assert cost == 3.0
    assert "Model pricing not available for model 'custom-openai-model'" in caplog.text


def test_openai_estimate_cost_uses_custom_model_pricing_override():
    fake_client = _FakeOpenAIClient(response=None)
    client = OpenAIChatCompletionClient(
        model="gpt-4.1-mini",
        client=fake_client,
        model_pricing=ModelPricing(input_per_1m=1.0, output_per_1m=2.0),
    )

    cost = client._estimate_cost(prompt_tokens=100_000, completion_tokens=50_000)

    assert cost == 0.2
