"""KoanTracer — wrap OpenTelemetry spans around LLM calls, tool calls, agent runs."""

from __future__ import annotations

from contextlib import contextmanager
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from collections.abc import Iterator

    from opentelemetry.trace import Span, Tracer


def _get_tracer(name: str = "koan") -> Tracer:
    """Get an OTel tracer, importing lazily."""
    from opentelemetry import trace  # type: ignore[import-untyped]

    return trace.get_tracer(name)


class KoanTracer:
    """Wraps OpenTelemetry spans around KOAN operations.

    Provides convenience methods for creating spans around LLM calls,
    tool calls, and agent runs with correct attributes.

    Usage::

        tracer = KoanTracer()

        with tracer.agent_span("researcher", run_id="r-1") as span:
            with tracer.llm_span("gpt-4o", provider="openai") as llm:
                # ... LLM call ...
                llm.set_attribute("koan.tokens.input", 150)
            with tracer.tool_span("web_search") as tool:
                # ... tool call ...
                pass
    """

    def __init__(self, tracer_name: str = "koan") -> None:
        self._tracer = _get_tracer(tracer_name)

    @contextmanager
    def agent_span(
        self,
        agent_name: str,
        *,
        run_id: str = "",
        thread_id: str = "",
    ) -> Iterator[Span]:
        """Create a span for an agent run."""
        with self._tracer.start_as_current_span(
            f"agent.{agent_name}",
            attributes={
                "koan.agent.name": agent_name,
                "koan.run_id": run_id,
                "koan.thread_id": thread_id,
            },
        ) as span:
            yield span

    @contextmanager
    def llm_span(
        self,
        model: str,
        *,
        provider: str = "",
    ) -> Iterator[Span]:
        """Create a span for an LLM call."""
        with self._tracer.start_as_current_span(
            f"llm.{provider}.{model}" if provider else f"llm.{model}",
            attributes={
                "koan.llm.model": model,
                "koan.llm.provider": provider,
            },
        ) as span:
            yield span

    @contextmanager
    def tool_span(
        self,
        tool_name: str,
        *,
        arguments: str = "",
    ) -> Iterator[Span]:
        """Create a span for a tool call."""
        with self._tracer.start_as_current_span(
            f"tool.{tool_name}",
            attributes={
                "koan.tool.name": tool_name,
                "koan.tool.arguments": arguments,
            },
        ) as span:
            yield span

    def record_llm_usage(
        self,
        span: Span,
        *,
        input_tokens: int = 0,
        output_tokens: int = 0,
        total_tokens: int = 0,
    ) -> None:
        """Record token usage on an LLM span."""
        span.set_attribute("koan.tokens.input", input_tokens)
        span.set_attribute("koan.tokens.output", output_tokens)
        span.set_attribute("koan.tokens.total", total_tokens or (input_tokens + output_tokens))

    def record_error(self, span: Span, error: Exception) -> None:
        """Record an error on a span."""
        span.set_attribute("koan.error", True)
        span.set_attribute("koan.error.message", str(error))
        span.record_exception(error)
