"""Tests for the OpenAI wrapper."""

from __future__ import annotations

import asyncio
from types import SimpleNamespace
from unittest.mock import MagicMock

import pytest

from tokengrip import Tokengrip, wrap_openai


def _make_tg() -> Tokengrip:
    return Tokengrip(api_key="tq_test_123", flush_interval_ms=0)


class TestNonStreaming:
    def test_chat_completions_tracks_usage(self):
        tg = _make_tg()
        mock_client = MagicMock()
        mock_client.chat.completions.create.return_value = SimpleNamespace(
            model="gpt-4.1",
            choices=[SimpleNamespace(message=SimpleNamespace(content="Hello!"))],
            usage=SimpleNamespace(prompt_tokens=25, completion_tokens=50),
        )

        wrapped = wrap_openai(mock_client, tg)
        result = wrapped.chat.completions.create(
            model="gpt-4.1",
            messages=[{"role": "user", "content": "Hi"}],
        )

        assert result.model == "gpt-4.1"
        assert len(tg._buffer) == 1
        rec = tg._buffer[0]
        assert rec["provider"] == "openai"
        assert rec["model"] == "gpt-4.1"
        assert rec["inputTokens"] == 25
        assert rec["outputTokens"] == 50
        assert "latencyMs" in rec
        tg.shutdown()

    def test_responses_api_tracks_usage(self):
        tg = _make_tg()
        mock_client = MagicMock()
        mock_client.responses.create.return_value = SimpleNamespace(
            model="gpt-4.1",
            output_text="Response text",
            usage=SimpleNamespace(input_tokens=30, output_tokens=60),
        )

        wrapped = wrap_openai(mock_client, tg)
        result = wrapped.responses.create(
            model="gpt-4.1",
            input="Hello",
        )

        assert result.output_text == "Response text"
        assert len(tg._buffer) == 1
        rec = tg._buffer[0]
        assert rec["inputTokens"] == 30
        assert rec["outputTokens"] == 60
        tg.shutdown()

    def test_passthrough_attributes(self):
        tg = _make_tg()
        mock_client = MagicMock()
        mock_client.api_key = "sk-test"
        mock_client.base_url = "https://api.openai.com"

        wrapped = wrap_openai(mock_client, tg)
        assert wrapped.api_key == "sk-test"
        assert wrapped.base_url == "https://api.openai.com"
        tg.shutdown()

    def test_no_usage_no_tracking(self):
        tg = _make_tg()
        mock_client = MagicMock()
        mock_client.chat.completions.create.return_value = SimpleNamespace(
            model="gpt-4.1",
            choices=[],
            usage=None,
        )

        wrapped = wrap_openai(mock_client, tg)
        wrapped.chat.completions.create(model="gpt-4.1", messages=[])
        assert len(tg._buffer) == 0
        tg.shutdown()

    def test_wrap_options_applied(self):
        tg = _make_tg()
        mock_client = MagicMock()
        mock_client.chat.completions.create.return_value = SimpleNamespace(
            model="gpt-4.1",
            choices=[],
            usage=SimpleNamespace(prompt_tokens=10, completion_tokens=5),
        )

        wrapped = wrap_openai(
            mock_client,
            tg,
            {"project_id": "proj-1", "agent_name": "bot", "task_type": "test"},
        )
        wrapped.chat.completions.create(model="gpt-4.1", messages=[])
        rec = tg._buffer[0]
        assert rec["projectId"] == "proj-1"
        assert rec["agentName"] == "bot"
        assert rec["taskType"] == "test"
        tg.shutdown()

    def test_content_capture(self):
        tg = _make_tg()
        mock_client = MagicMock()
        mock_client.chat.completions.create.return_value = SimpleNamespace(
            model="gpt-4.1",
            choices=[SimpleNamespace(message=SimpleNamespace(content="Reply!"))],
            usage=SimpleNamespace(prompt_tokens=10, completion_tokens=5),
        )

        wrapped = wrap_openai(mock_client, tg, {"capture_content": True})
        wrapped.chat.completions.create(
            model="gpt-4.1",
            messages=[{"role": "user", "content": "Hello world"}],
        )
        rec = tg._buffer[0]
        assert rec["promptContent"] == "Hello world"
        assert rec["responseContent"] == "Reply!"
        tg.shutdown()

    def test_provider_error_not_swallowed(self):
        tg = _make_tg()
        mock_client = MagicMock()
        mock_client.chat.completions.create.side_effect = RuntimeError("API error")

        wrapped = wrap_openai(mock_client, tg)
        with pytest.raises(RuntimeError, match="API error"):
            wrapped.chat.completions.create(model="gpt-4.1", messages=[])
        tg.shutdown()


class TestSyncStreaming:
    def test_stream_tracks_on_completion(self):
        tg = _make_tg()
        mock_client = MagicMock()

        chunks = [
            SimpleNamespace(model="gpt-4.1", usage=None),
            SimpleNamespace(model="gpt-4.1", usage=None),
            SimpleNamespace(
                model="gpt-4.1",
                usage=SimpleNamespace(prompt_tokens=20, completion_tokens=40),
            ),
        ]
        mock_client.chat.completions.create.return_value = iter(chunks)

        wrapped = wrap_openai(mock_client, tg)
        stream = wrapped.chat.completions.create(
            model="gpt-4.1", messages=[], stream=True
        )

        collected = list(stream)
        assert len(collected) == 3
        assert len(tg._buffer) == 1
        rec = tg._buffer[0]
        assert rec["inputTokens"] == 20
        assert rec["outputTokens"] == 40
        tg.shutdown()

    def test_stream_context_manager(self):
        tg = _make_tg()
        mock_client = MagicMock()

        chunks = [
            SimpleNamespace(
                model="gpt-4.1",
                usage=SimpleNamespace(prompt_tokens=10, completion_tokens=20),
            ),
        ]

        class FakeStream:
            def __iter__(self):
                return iter(chunks)

            def __enter__(self):
                return self

            def __exit__(self, *args):
                pass

        mock_client.chat.completions.create.return_value = FakeStream()

        wrapped = wrap_openai(mock_client, tg)
        with wrapped.chat.completions.create(
            model="gpt-4.1", messages=[], stream=True
        ) as stream:
            for _ in stream:
                pass

        assert len(tg._buffer) == 1
        tg.shutdown()

    def test_responses_api_streaming(self):
        tg = _make_tg()
        mock_client = MagicMock()

        chunks = [
            SimpleNamespace(model="gpt-4.1", usage=None),
            SimpleNamespace(
                model="gpt-4.1",
                usage=SimpleNamespace(input_tokens=15, output_tokens=30),
            ),
        ]
        mock_client.responses.create.return_value = iter(chunks)

        wrapped = wrap_openai(mock_client, tg)
        stream = wrapped.responses.create(
            model="gpt-4.1", input="test", stream=True
        )
        list(stream)

        rec = tg._buffer[0]
        assert rec["inputTokens"] == 15
        assert rec["outputTokens"] == 30
        tg.shutdown()


class TestAsyncNonStreaming:
    def test_async_chat_completions(self):
        tg = _make_tg()
        mock_client = MagicMock()

        async def mock_create(**kwargs):
            return SimpleNamespace(
                model="gpt-4.1",
                choices=[],
                usage=SimpleNamespace(prompt_tokens=100, completion_tokens=200),
            )

        mock_client.chat.completions.create = mock_create

        wrapped = wrap_openai(mock_client, tg)

        async def run():
            result = await wrapped.chat.completions.create(
                model="gpt-4.1", messages=[]
            )
            return result

        result = asyncio.get_event_loop().run_until_complete(run())
        assert result.model == "gpt-4.1"
        assert len(tg._buffer) == 1
        assert tg._buffer[0]["inputTokens"] == 100
        tg.shutdown()


class TestAsyncStreaming:
    def test_async_stream_tracks(self):
        tg = _make_tg()
        mock_client = MagicMock()

        chunks = [
            SimpleNamespace(model="gpt-4.1", usage=None),
            SimpleNamespace(
                model="gpt-4.1",
                usage=SimpleNamespace(prompt_tokens=50, completion_tokens=100),
            ),
        ]

        class AsyncChunks:
            def __init__(self):
                self._iter = iter(chunks)

            def __aiter__(self):
                return self

            async def __anext__(self):
                try:
                    return next(self._iter)
                except StopIteration:
                    raise StopAsyncIteration

        async def mock_create(**kwargs):
            return AsyncChunks()

        mock_client.chat.completions.create = mock_create

        wrapped = wrap_openai(mock_client, tg)

        async def run():
            stream = await wrapped.chat.completions.create(
                model="gpt-4.1", messages=[], stream=True
            )
            collected = []
            async for chunk in stream:
                collected.append(chunk)
            return collected

        result = asyncio.get_event_loop().run_until_complete(run())
        assert len(result) == 2
        assert len(tg._buffer) == 1
        assert tg._buffer[0]["inputTokens"] == 50
        assert tg._buffer[0]["outputTokens"] == 100
        tg.shutdown()
