"""Tests for the @trace_llm() decorator."""

import asyncio
from types import SimpleNamespace
from typing import List, Sequence

import pytest
from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider, ReadableSpan
from opentelemetry.sdk.trace.export import SimpleSpanProcessor, SpanExporter, SpanExportResult
from opentelemetry.trace import StatusCode

from pendo_sdk.decorators import trace_llm


class _InMemoryExporter(SpanExporter):
    """Minimal in-memory exporter for test assertions."""

    def __init__(self) -> None:
        self._spans: List[ReadableSpan] = []

    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
        self._spans.extend(spans)
        return SpanExportResult.SUCCESS

    def get_finished_spans(self) -> List[ReadableSpan]:
        return list(self._spans)

    def shutdown(self) -> None:
        pass


@pytest.fixture(autouse=True)
def _otel_setup():
    """Set up an in-memory exporter so we can inspect finished spans."""
    exporter = _InMemoryExporter()
    provider = TracerProvider()
    provider.add_span_processor(SimpleSpanProcessor(exporter))
    trace.set_tracer_provider(provider)
    yield exporter
    provider.shutdown()
    trace._TRACER_PROVIDER = None
    trace._TRACER_PROVIDER_SET_ONCE._done = False


class TestTraceLlmSync:
    def test_creates_generation_span(self, _otel_setup):
        exporter = _otel_setup

        @trace_llm(name="my_span")
        def my_func():
            return "ok"

        result = my_func()
        assert result == "ok"

        spans = exporter.get_finished_spans()
        assert len(spans) == 1
        span = spans[0]
        assert span.name == "my_span"
        assert span.attributes["pendo.span.type"] == "GENERATION"

    def test_default_name_uses_function_name(self, _otel_setup):
        exporter = _otel_setup

        @trace_llm()
        def classify_intent():
            return "ok"

        classify_intent()

        spans = exporter.get_finished_spans()
        assert spans[0].name == "classify_intent"

    def test_extracts_token_usage(self, _otel_setup):
        exporter = _otel_setup

        @trace_llm(name="with_usage")
        def llm_call():
            return SimpleNamespace(
                model="claude-sonnet-4-20250514",
                usage=SimpleNamespace(input_tokens=100, output_tokens=50),
            )

        result = llm_call()
        assert result.model == "claude-sonnet-4-20250514"

        span = exporter.get_finished_spans()[0]
        assert span.attributes["llm.token_count.prompt"] == 100
        assert span.attributes["llm.token_count.completion"] == 50
        assert span.attributes["llm.model_name"] == "claude-sonnet-4-20250514"

    def test_handles_missing_usage(self, _otel_setup):
        exporter = _otel_setup

        @trace_llm(name="no_usage")
        def llm_call():
            return {"text": "hello"}

        result = llm_call()
        assert result == {"text": "hello"}

        span = exporter.get_finished_spans()[0]
        assert span.attributes["pendo.span.type"] == "GENERATION"
        assert "llm.token_count.prompt" not in span.attributes
        assert "llm.model_name" not in span.attributes

    def test_exception_sets_error_status(self, _otel_setup):
        exporter = _otel_setup

        @trace_llm(name="failing")
        def bad_call():
            raise ValueError("boom")

        with pytest.raises(ValueError, match="boom"):
            bad_call()

        span = exporter.get_finished_spans()[0]
        assert span.status.status_code == StatusCode.ERROR
        assert "boom" in span.status.description

    def test_preserves_function_metadata(self):
        @trace_llm(name="custom")
        def my_documented_func():
            """My docstring."""
            pass

        assert my_documented_func.__name__ == "my_documented_func"
        assert my_documented_func.__doc__ == "My docstring."


class TestTraceLlmAsync:
    def test_async_creates_generation_span(self, _otel_setup):
        exporter = _otel_setup

        @trace_llm(name="async_span")
        async def async_llm():
            return SimpleNamespace(
                model="gpt-4",
                usage=SimpleNamespace(input_tokens=200, output_tokens=80),
            )

        result = asyncio.run(async_llm())
        assert result.model == "gpt-4"

        spans = exporter.get_finished_spans()
        assert len(spans) == 1
        span = spans[0]
        assert span.name == "async_span"
        assert span.attributes["pendo.span.type"] == "GENERATION"
        assert span.attributes["llm.token_count.prompt"] == 200
        assert span.attributes["llm.token_count.completion"] == 80
        assert span.attributes["llm.model_name"] == "gpt-4"

    def test_async_exception_sets_error_status(self, _otel_setup):
        exporter = _otel_setup

        @trace_llm(name="async_fail")
        async def bad_async():
            raise RuntimeError("async boom")

        with pytest.raises(RuntimeError, match="async boom"):
            asyncio.run(bad_async())

        span = exporter.get_finished_spans()[0]
        assert span.status.status_code == StatusCode.ERROR
