"""Tests for KoanTracer OpenTelemetry span creation."""

from __future__ import annotations

from opentelemetry.sdk.trace import TracerProvider  # type: ignore[import-untyped]
from opentelemetry.sdk.trace.export import (  # type: ignore[import-untyped]
    ReadableSpan,
    SimpleSpanProcessor,
    SpanExporter,
    SpanExportResult,
)

from koan.telemetry.tracing import KoanTracer


class _ListExporter(SpanExporter):  # type: ignore[misc]
    """Simple in-memory exporter that collects spans in a list."""

    def __init__(self) -> None:
        self.spans: list[ReadableSpan] = []

    def export(self, spans: list[ReadableSpan]) -> SpanExportResult:  # type: ignore[override]
        self.spans.extend(spans)
        return SpanExportResult.SUCCESS

    def shutdown(self) -> None:
        pass


def _make_tracer() -> tuple[KoanTracer, _ListExporter]:
    """Create a KoanTracer with its own TracerProvider and list exporter."""
    exporter = _ListExporter()
    provider = TracerProvider()
    provider.add_span_processor(SimpleSpanProcessor(exporter))
    # Directly inject the provider's tracer into KoanTracer
    tracer = KoanTracer.__new__(KoanTracer)
    tracer._tracer = provider.get_tracer("koan-test")
    return tracer, exporter


class TestKoanTracer:
    def test_agent_span_creates_span_with_attributes(self) -> None:
        tracer, exporter = _make_tracer()

        with tracer.agent_span("researcher", run_id="r-1", thread_id="t-1"):
            pass

        assert len(exporter.spans) == 1
        span = exporter.spans[0]
        assert span.name == "agent.researcher"
        attrs = dict(span.attributes or {})
        assert attrs["koan.agent.name"] == "researcher"
        assert attrs["koan.run_id"] == "r-1"
        assert attrs["koan.thread_id"] == "t-1"

    def test_llm_span_with_provider(self) -> None:
        tracer, exporter = _make_tracer()

        with tracer.llm_span("gpt-4o", provider="openai"):
            pass

        assert len(exporter.spans) == 1
        span = exporter.spans[0]
        assert span.name == "llm.openai.gpt-4o"
        attrs = dict(span.attributes or {})
        assert attrs["koan.llm.model"] == "gpt-4o"
        assert attrs["koan.llm.provider"] == "openai"

    def test_llm_span_without_provider(self) -> None:
        tracer, exporter = _make_tracer()

        with tracer.llm_span("claude-3-haiku"):
            pass

        assert exporter.spans[0].name == "llm.claude-3-haiku"

    def test_tool_span_creates_span(self) -> None:
        tracer, exporter = _make_tracer()

        with tracer.tool_span("web_search", arguments='{"query": "test"}'):
            pass

        assert len(exporter.spans) == 1
        attrs = dict(exporter.spans[0].attributes or {})
        assert attrs["koan.tool.name"] == "web_search"
        assert attrs["koan.tool.arguments"] == '{"query": "test"}'

    def test_nested_spans(self) -> None:
        tracer, exporter = _make_tracer()

        with tracer.agent_span("researcher", run_id="r-1"):
            with tracer.llm_span("gpt-4o", provider="openai"):
                pass
            with tracer.tool_span("search"):
                pass

        assert len(exporter.spans) == 3

        agent_span = next(s for s in exporter.spans if s.name == "agent.researcher")
        llm_span = next(s for s in exporter.spans if s.name.startswith("llm."))
        tool_span = next(s for s in exporter.spans if s.name.startswith("tool."))

        assert llm_span.parent is not None
        assert llm_span.parent.span_id == agent_span.context.span_id
        assert tool_span.parent is not None
        assert tool_span.parent.span_id == agent_span.context.span_id

    def test_record_llm_usage(self) -> None:
        tracer, exporter = _make_tracer()

        with tracer.llm_span("gpt-4o", provider="openai") as span:
            tracer.record_llm_usage(span, input_tokens=150, output_tokens=50)

        attrs = dict(exporter.spans[0].attributes or {})
        assert attrs["koan.tokens.input"] == 150
        assert attrs["koan.tokens.output"] == 50
        assert attrs["koan.tokens.total"] == 200

    def test_record_llm_usage_explicit_total(self) -> None:
        tracer, exporter = _make_tracer()

        with tracer.llm_span("gpt-4o") as span:
            tracer.record_llm_usage(
                span, input_tokens=100, output_tokens=50, total_tokens=175
            )

        attrs = dict(exporter.spans[0].attributes or {})
        assert attrs["koan.tokens.total"] == 175

    def test_record_error(self) -> None:
        tracer, exporter = _make_tracer()

        with tracer.agent_span("researcher") as span:
            tracer.record_error(span, RuntimeError("something broke"))

        attrs = dict(exporter.spans[0].attributes or {})
        assert attrs["koan.error"] is True
        assert "something broke" in attrs["koan.error.message"]
        assert len(exporter.spans[0].events) >= 1
