"""Tests for the Anthropic wrapper."""

from __future__ import annotations

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

import pytest

from tokengrip import Tokengrip, wrap_anthropic


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


class TestNonStreaming:
    def test_messages_create_tracks_usage(self):
        tg = _make_tg()
        mock_client = MagicMock()
        mock_client.messages.create.return_value = SimpleNamespace(
            model="claude-sonnet-4-5-20250929",
            content=[SimpleNamespace(type="text", text="Hello!")],
            usage=SimpleNamespace(input_tokens=25, output_tokens=50),
        )

        wrapped = wrap_anthropic(mock_client, tg)
        result = wrapped.messages.create(
            model="claude-sonnet-4-5-20250929",
            max_tokens=1024,
            messages=[{"role": "user", "content": "Hi"}],
        )

        assert result.model == "claude-sonnet-4-5-20250929"
        assert len(tg._buffer) == 1
        rec = tg._buffer[0]
        assert rec["provider"] == "anthropic"
        assert rec["model"] == "claude-sonnet-4-5-20250929"
        assert rec["inputTokens"] == 25
        assert rec["outputTokens"] == 50
        assert "latencyMs" in rec
        tg.shutdown()

    def test_passthrough_attributes(self):
        tg = _make_tg()
        mock_client = MagicMock()
        mock_client.api_key = "sk-ant-test"

        wrapped = wrap_anthropic(mock_client, tg)
        assert wrapped.api_key == "sk-ant-test"
        tg.shutdown()

    def test_no_usage_no_tracking(self):
        tg = _make_tg()
        mock_client = MagicMock()
        mock_client.messages.create.return_value = SimpleNamespace(
            model="claude-sonnet-4-5-20250929",
            content=[],
            usage=None,
        )

        wrapped = wrap_anthropic(mock_client, tg)
        wrapped.messages.create(
            model="claude-sonnet-4-5-20250929", max_tokens=1024, messages=[]
        )
        assert len(tg._buffer) == 0
        tg.shutdown()

    def test_wrap_options_applied(self):
        tg = _make_tg()
        mock_client = MagicMock()
        mock_client.messages.create.return_value = SimpleNamespace(
            model="claude-sonnet-4-5-20250929",
            content=[],
            usage=SimpleNamespace(input_tokens=10, output_tokens=5),
        )

        wrapped = wrap_anthropic(
            mock_client,
            tg,
            {"project_id": "proj-1", "agent_name": "bot", "task_type": "test"},
        )
        wrapped.messages.create(
            model="claude-sonnet-4-5-20250929", max_tokens=1024, 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.messages.create.return_value = SimpleNamespace(
            model="claude-sonnet-4-5-20250929",
            content=[SimpleNamespace(type="text", text="Reply!")],
            usage=SimpleNamespace(input_tokens=10, output_tokens=5),
        )

        wrapped = wrap_anthropic(mock_client, tg, {"capture_content": True})
        wrapped.messages.create(
            model="claude-sonnet-4-5-20250929",
            max_tokens=1024,
            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.messages.create.side_effect = RuntimeError("API error")

        wrapped = wrap_anthropic(mock_client, tg)
        with pytest.raises(RuntimeError, match="API error"):
            wrapped.messages.create(
                model="claude-sonnet-4-5-20250929", max_tokens=1024, messages=[]
            )
        tg.shutdown()


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

        events = [
            SimpleNamespace(
                type="message_start",
                message=SimpleNamespace(
                    model="claude-sonnet-4-5-20250929",
                    usage=SimpleNamespace(input_tokens=25),
                ),
            ),
            SimpleNamespace(
                type="content_block_delta",
                delta=SimpleNamespace(type="text_delta", text="Hello"),
            ),
            SimpleNamespace(
                type="message_delta",
                usage=SimpleNamespace(output_tokens=50),
            ),
            SimpleNamespace(type="message_stop"),
        ]
        mock_client.messages.create.return_value = iter(events)

        wrapped = wrap_anthropic(mock_client, tg)
        stream = wrapped.messages.create(
            model="claude-sonnet-4-5-20250929",
            max_tokens=1024,
            messages=[{"role": "user", "content": "Hi"}],
            stream=True,
        )

        collected = list(stream)
        assert len(collected) == 4
        assert len(tg._buffer) == 1
        rec = tg._buffer[0]
        assert rec["provider"] == "anthropic"
        assert rec["model"] == "claude-sonnet-4-5-20250929"
        assert rec["inputTokens"] == 25
        assert rec["outputTokens"] == 50
        tg.shutdown()

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

        events = [
            SimpleNamespace(
                type="message_start",
                message=SimpleNamespace(
                    model="claude-sonnet-4-5-20250929",
                    usage=SimpleNamespace(input_tokens=10),
                ),
            ),
            SimpleNamespace(
                type="message_delta",
                usage=SimpleNamespace(output_tokens=20),
            ),
        ]

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

            def __enter__(self):
                return self

            def __exit__(self, *args):
                pass

        mock_client.messages.create.return_value = FakeStream()

        wrapped = wrap_anthropic(mock_client, tg)
        with wrapped.messages.create(
            model="claude-sonnet-4-5-20250929",
            max_tokens=1024,
            messages=[],
            stream=True,
        ) as stream:
            for _ in stream:
                pass

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

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

        class FakeStream:
            def __iter__(self):
                return iter([])

            def get_final_message(self):
                return SimpleNamespace(
                    model="claude-sonnet-4-5-20250929",
                    usage=SimpleNamespace(input_tokens=30, output_tokens=60),
                )

        mock_client.messages.create.return_value = FakeStream()

        wrapped = wrap_anthropic(mock_client, tg)
        stream = wrapped.messages.create(
            model="claude-sonnet-4-5-20250929",
            max_tokens=1024,
            messages=[],
            stream=True,
        )

        msg = stream.get_final_message()
        assert msg.model == "claude-sonnet-4-5-20250929"
        assert len(tg._buffer) == 1
        rec = tg._buffer[0]
        assert rec["inputTokens"] == 30
        assert rec["outputTokens"] == 60
        tg.shutdown()


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

        async def mock_create(**kwargs):
            return SimpleNamespace(
                model="claude-sonnet-4-5-20250929",
                content=[],
                usage=SimpleNamespace(input_tokens=100, output_tokens=200),
            )

        mock_client.messages.create = mock_create

        wrapped = wrap_anthropic(mock_client, tg)

        async def run():
            return await wrapped.messages.create(
                model="claude-sonnet-4-5-20250929", max_tokens=1024, messages=[]
            )

        result = asyncio.get_event_loop().run_until_complete(run())
        assert result.model == "claude-sonnet-4-5-20250929"
        assert len(tg._buffer) == 1
        assert tg._buffer[0]["inputTokens"] == 100
        tg.shutdown()


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

        events = [
            SimpleNamespace(
                type="message_start",
                message=SimpleNamespace(
                    model="claude-sonnet-4-5-20250929",
                    usage=SimpleNamespace(input_tokens=50),
                ),
            ),
            SimpleNamespace(
                type="message_delta",
                usage=SimpleNamespace(output_tokens=100),
            ),
        ]

        class AsyncEvents:
            def __init__(self):
                self._iter = iter(events)

            def __aiter__(self):
                return self

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

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

        mock_client.messages.create = mock_create

        wrapped = wrap_anthropic(mock_client, tg)

        async def run():
            stream = await wrapped.messages.create(
                model="claude-sonnet-4-5-20250929",
                max_tokens=1024,
                messages=[],
                stream=True,
            )
            collected = []
            async for event in stream:
                collected.append(event)
            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()
