"""Tests for staso.integrations.openai."""
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.openai as mod

    original = mod._patched
    mod._patched = False
    yield
    mod._patched = original


def _fake_completion(
    content="Hello!",
    model="gpt-4",
    prompt_tokens=10,
    completion_tokens=20,
    finish_reason="stop",
    tool_calls=None,
):
    """Build a fake ChatCompletion response."""
    return SimpleNamespace(
        id="chatcmpl-test",
        choices=[
            SimpleNamespace(
                message=SimpleNamespace(
                    content=content,
                    role="assistant",
                    tool_calls=tool_calls,
                ),
                finish_reason=finish_reason,
            )
        ],
        usage=SimpleNamespace(
            prompt_tokens=prompt_tokens,
            completion_tokens=completion_tokens,
            total_tokens=prompt_tokens + completion_tokens,
            completion_tokens_details=None,
        ),
        model=model,
    )


class TestPatchOpenAI:
    def test_raises_without_openai_installed(self):
        with patch.dict(
            "sys.modules",
            {
                "openai": None,
                "openai.resources": None,
                "openai.resources.chat": None,
            },
        ):
            import staso.integrations.openai as mod

            mod._patched = False
            with pytest.raises(ImportError, match="openai package is required"):
                mod.patch_openai()

    def test_idempotent(self):
        import staso.integrations.openai as mod

        mod._patched = True
        mod.patch_openai()  # Should be a no-op


class TestSyncWrapper:
    def test_emits_span(self, capturing):
        class FakeCompletions:
            def create(self, *args, **kwargs):
                return _fake_completion()

        from staso.integrations.openai import _wrap_sync

        _wrap_sync(FakeCompletions)
        result = FakeCompletions().create(model="gpt-4", max_tokens=1024)

        assert result.usage.prompt_tokens == 10
        assert len(capturing) == 1

        span = capturing[0]
        assert span["name"] == "openai.chat.completions.create"
        assert span["kind"] == "llm"
        assert span["model"] == "gpt-4"
        assert span["input_tokens"] == 10
        assert span["output_tokens"] == 20
        assert span["total_tokens"] == 30
        assert span["status"] == "ok"

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

        from staso.integrations.openai import _wrap_sync

        _wrap_sync(FakeCompletions)

        with pytest.raises(ConnectionError, match="API down"):
            FakeCompletions().create(model="gpt-4")

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

    def test_captures_input_and_metadata(self, capturing):
        class FakeCompletions:
            def create(self, *args, **kwargs):
                return _fake_completion()

        from staso.integrations.openai import _wrap_sync

        _wrap_sync(FakeCompletions)
        FakeCompletions().create(
            model="gpt-4",
            max_tokens=512,
            temperature=0.7,
        )

        input_val = json.loads(capturing[0]["input"])
        assert input_val["model"] == "gpt-4"

        attrs = json.loads(capturing[0]["attributes"])
        assert attrs["max_tokens"] == 512
        assert attrs["temperature"] == 0.7
        assert attrs["provider"] == "openai"

    def test_captures_tool_calls(self, capturing):
        tool_call = SimpleNamespace(
            id="call_123",
            type="function",
            function=SimpleNamespace(name="get_weather", arguments='{"city":"SF"}'),
        )
        tool_call.model_dump = lambda: {
            "id": "call_123",
            "type": "function",
            "function": {"name": "get_weather", "arguments": '{"city":"SF"}'},
        }

        class FakeCompletions:
            def create(self, *args, **kwargs):
                return _fake_completion(
                    content=None,
                    tool_calls=[tool_call],
                    finish_reason="tool_calls",
                )

        from staso.integrations.openai import _wrap_sync

        _wrap_sync(FakeCompletions)
        FakeCompletions().create(model="gpt-4")

        output = json.loads(capturing[0]["output"])
        assert "tool_calls" in output
        assert output["tool_calls"][0]["id"] == "call_123"

    def test_captures_output_content(self, capturing):
        class FakeCompletions:
            def create(self, *args, **kwargs):
                return _fake_completion(content="Hello world!")

        from staso.integrations.openai import _wrap_sync

        _wrap_sync(FakeCompletions)
        FakeCompletions().create(model="gpt-4")

        output = json.loads(capturing[0]["output"])
        assert output["content"] == "Hello world!"
        assert output["role"] == "assistant"

    def test_captures_stop_reason(self, capturing):
        class FakeCompletions:
            def create(self, *args, **kwargs):
                return _fake_completion(finish_reason="length")

        from staso.integrations.openai import _wrap_sync

        _wrap_sync(FakeCompletions)
        FakeCompletions().create(model="gpt-4")

        attrs = json.loads(capturing[0]["attributes"])
        assert attrs["stop_reason"] == "length"


class TestAsyncWrapper:
    def test_emits_span(self, capturing):
        class FakeAsyncCompletions:
            async def create(self, *args, **kwargs):
                return _fake_completion(
                    model="gpt-4-turbo", prompt_tokens=200, completion_tokens=75
                )

        from staso.integrations.openai import _wrap_async

        _wrap_async(FakeAsyncCompletions)
        result = asyncio.run(FakeAsyncCompletions().create(model="gpt-4-turbo"))

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

        span = capturing[0]
        assert span["kind"] == "llm"
        assert span["model"] == "gpt-4-turbo"
        assert span["input_tokens"] == 200
        assert span["output_tokens"] == 75
        assert span["total_tokens"] == 275

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

        from staso.integrations.openai import _wrap_async

        _wrap_async(FakeAsyncCompletions)

        with pytest.raises(ConnectionError, match="async fail"):
            asyncio.run(FakeAsyncCompletions().create(model="gpt-4"))

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


class TestExtractModel:
    def test_from_kwargs(self):
        from staso.integrations.openai import _extract_model

        assert _extract_model((), {"model": "gpt-4"}) == "gpt-4"

    def test_from_args(self):
        from staso.integrations.openai import _extract_model

        assert _extract_model(("gpt-4",), {}) == "gpt-4"

    def test_empty(self):
        from staso.integrations.openai import _extract_model

        assert _extract_model((), {}) == ""
