"""Tests for staso.integrations.anthropic."""
from __future__ import annotations

import asyncio
import json
from types import SimpleNamespace
from unittest.mock import patch

import pytest


@pytest.fixture(autouse=True)
def _reset_patch_flag():
    """Reset the _patched flag between tests."""
    import staso.integrations.anthropic as mod
    original = mod._patched
    mod._patched = False
    yield
    mod._patched = original


class TestPatchAnthropic:
    def test_raises_without_anthropic_installed(self):
        with patch.dict("sys.modules", {"anthropic": None, "anthropic.resources": None}):
            import staso.integrations.anthropic as mod
            mod._patched = False
            with pytest.raises(ImportError, match="anthropic package is required"):
                mod.patch_anthropic()

    def test_idempotent(self):
        import staso.integrations.anthropic as mod
        mod._patched = True
        mod.patch_anthropic()  # Should be a no-op


class TestSyncWrapper:
    def test_emits_span(self, capturing):
        class FakeMessages:
            def create(self, *args, **kwargs):
                return SimpleNamespace(
                    content=[],
                    id="msg_test",
                    usage=SimpleNamespace(input_tokens=100, output_tokens=50),
                    stop_reason="end_turn",
                )

        from staso.integrations.anthropic import _wrap_sync

        _wrap_sync(FakeMessages)
        result = FakeMessages().create(model="claude-test", max_tokens=1024)

        assert result.usage.input_tokens == 100
        assert len(capturing) == 1

        span = capturing[0]
        assert span["name"] == "anthropic.messages.create"
        assert span["kind"] == "llm"
        assert span["model"] == "claude-test"
        assert span["input_tokens"] == 100
        assert span["output_tokens"] == 50
        assert span["total_tokens"] == 150
        assert span["status"] == "ok"

    def test_captures_error(self, capturing):
        class FakeMessages:
            def create(self, *args, **kwargs):
                raise ConnectionError("API down")

        from staso.integrations.anthropic import _wrap_sync

        _wrap_sync(FakeMessages)

        with pytest.raises(ConnectionError, match="API down"):
            FakeMessages().create(model="claude-test")

        assert capturing[0]["status"] == "error"
        assert capturing[0]["error_message"] == "API down"

    def test_captures_input_and_metadata(self, capturing):
        class FakeMessages:
            def create(self, *args, **kwargs):
                return SimpleNamespace(
                    content=[],
                    id="msg_test",
                    usage=SimpleNamespace(input_tokens=10, output_tokens=5),
                    stop_reason="end_turn",
                )

        from staso.integrations.anthropic import _wrap_sync

        _wrap_sync(FakeMessages)
        FakeMessages().create(
            model="claude-test",
            max_tokens=512,
            temperature=0.7,
            system="You are a helpful assistant.",
        )

        input_val = json.loads(capturing[0]["input"])
        assert input_val["model"] == "claude-test"
        assert "You are a helpful assistant" in input_val["system"]

        # max_tokens and temperature are in metadata, not input
        attrs = json.loads(capturing[0]["attributes"])
        assert attrs["max_tokens"] == 512
        assert attrs["temperature"] == 0.7

    def test_captures_output(self, capturing):
        class FakeMessages:
            def create(self, *args, **kwargs):
                return SimpleNamespace(
                    content=[],
                    id="msg_test",
                    usage=SimpleNamespace(input_tokens=10, output_tokens=5),
                    stop_reason="max_tokens",
                )

        from staso.integrations.anthropic import _wrap_sync

        _wrap_sync(FakeMessages)
        FakeMessages().create(model="claude-test")

        # stop_reason is in metadata (attributes), not output
        attrs = json.loads(capturing[0]["attributes"])
        assert attrs["stop_reason"] == "max_tokens"


class TestAsyncWrapper:
    def test_emits_span(self, capturing):
        class FakeAsyncMessages:
            async def create(self, *args, **kwargs):
                return SimpleNamespace(
                    content=[],
                    id="msg_test",
                    usage=SimpleNamespace(input_tokens=200, output_tokens=75),
                    stop_reason="end_turn",
                )

        from staso.integrations.anthropic import _wrap_async

        _wrap_async(FakeAsyncMessages)
        result = asyncio.run(FakeAsyncMessages().create(model="claude-async"))

        assert result.usage.input_tokens == 200
        assert len(capturing) == 1

        span = capturing[0]
        assert span["kind"] == "llm"
        assert span["model"] == "claude-async"
        assert span["input_tokens"] == 200
        assert span["output_tokens"] == 75
        assert span["total_tokens"] == 275

    def test_captures_error(self, capturing):
        class FakeAsyncMessages:
            async def create(self, *args, **kwargs):
                raise ConnectionError("async fail")

        from staso.integrations.anthropic import _wrap_async

        _wrap_async(FakeAsyncMessages)

        with pytest.raises(ConnectionError, match="async fail"):
            asyncio.run(FakeAsyncMessages().create(model="claude-async"))

        assert capturing[0]["status"] == "error"


class TestExtractModel:
    def test_from_kwargs(self):
        from staso.integrations.anthropic import _extract_model
        assert _extract_model((), {"model": "claude-test"}) == "claude-test"

    def test_from_args(self):
        from staso.integrations.anthropic import _extract_model
        assert _extract_model(("claude-test",), {}) == "claude-test"

    def test_empty(self):
        from staso.integrations.anthropic import _extract_model
        assert _extract_model((), {}) == ""
