"""
Decorators for tracing LLM and tool calls as typed spans.

Usage:
    @trace_llm(name="classify_intent")
    def classify(text: str) -> Message:
        return client.messages.create(model="claude-sonnet-4-20250514", ...)

    @trace_tool(name="search_db")
    def search(query: str) -> list[dict]:
        return db.search(query)
"""

import asyncio
import functools
from typing import Optional, Callable, TypeVar, Any

from opentelemetry import trace
from opentelemetry.trace import StatusCode

_SDK_NAME = "pendo_sdk"

F = TypeVar("F", bound=Callable[..., Any])


def trace_llm(name: Optional[str] = None) -> Callable[[F], F]:
    """Decorator that wraps a function in a GENERATION span.

    Args:
        name: Span name. Defaults to the decorated function's __name__.
    """

    def decorator(fn: F) -> F:
        span_name = name or fn.__name__

        if asyncio.iscoroutinefunction(fn):

            @functools.wraps(fn)
            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
                tracer = trace.get_tracer(_SDK_NAME)
                with tracer.start_as_current_span(span_name) as span:
                    span.set_attribute("pendo.span.type", "GENERATION")
                    try:
                        result = await fn(*args, **kwargs)
                    except Exception as exc:
                        span.set_status(StatusCode.ERROR, str(exc))
                        span.record_exception(exc)
                        raise
                    _extract_llm_attributes(span, result)
                    return result

            return async_wrapper  # type: ignore[return-value]

        else:

            @functools.wraps(fn)
            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
                tracer = trace.get_tracer(_SDK_NAME)
                with tracer.start_as_current_span(span_name) as span:
                    span.set_attribute("pendo.span.type", "GENERATION")
                    try:
                        result = fn(*args, **kwargs)
                    except Exception as exc:
                        span.set_status(StatusCode.ERROR, str(exc))
                        span.record_exception(exc)
                        raise
                    _extract_llm_attributes(span, result)
                    return result

            return sync_wrapper  # type: ignore[return-value]

    return decorator


def trace_tool(name: Optional[str] = None) -> Callable[[F], F]:
    """Decorator that wraps a function in a TOOL span.

    Captures the function's input arguments and output as span attributes,
    letting customers trace tool/function calls within their agent pipelines.

    Args:
        name: Span name. Defaults to the decorated function's __name__.
    """

    def decorator(fn: F) -> F:
        span_name = name or fn.__name__

        if asyncio.iscoroutinefunction(fn):

            @functools.wraps(fn)
            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
                tracer = trace.get_tracer(_SDK_NAME)
                with tracer.start_as_current_span(span_name) as span:
                    span.set_attribute("pendo.span.type", "TOOL")
                    _set_tool_input(span, fn, args, kwargs)
                    try:
                        result = await fn(*args, **kwargs)
                    except Exception as exc:
                        span.set_status(StatusCode.ERROR, str(exc))
                        span.record_exception(exc)
                        raise
                    _set_tool_output(span, result)
                    return result

            return async_wrapper  # type: ignore[return-value]

        else:

            @functools.wraps(fn)
            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
                tracer = trace.get_tracer(_SDK_NAME)
                with tracer.start_as_current_span(span_name) as span:
                    span.set_attribute("pendo.span.type", "TOOL")
                    _set_tool_input(span, fn, args, kwargs)
                    try:
                        result = fn(*args, **kwargs)
                    except Exception as exc:
                        span.set_status(StatusCode.ERROR, str(exc))
                        span.record_exception(exc)
                        raise
                    _set_tool_output(span, result)
                    return result

            return sync_wrapper  # type: ignore[return-value]

    return decorator


def _set_tool_input(span: trace.Span, fn: Callable, args: tuple, kwargs: dict) -> None:
    """Record tool input arguments as a span attribute."""
    import inspect

    try:
        sig = inspect.signature(fn)
        bound = sig.bind(*args, **kwargs)
        bound.apply_defaults()
        span.set_attribute("tool.input", str(dict(bound.arguments)))
    except Exception:
        # Don't let input capture break the actual function call
        pass


def _set_tool_output(span: trace.Span, result: Any) -> None:
    """Record tool output as a span attribute."""
    try:
        span.set_attribute("tool.output", str(result))
    except Exception:
        pass


def _extract_llm_attributes(span: trace.Span, result: Any) -> None:
    """Extract token usage and model name from an LLM response object."""
    usage = getattr(result, "usage", None)
    if usage is not None:
        input_tokens = getattr(usage, "input_tokens", None)
        if input_tokens is not None:
            span.set_attribute("llm.token_count.prompt", input_tokens)
        output_tokens = getattr(usage, "output_tokens", None)
        if output_tokens is not None:
            span.set_attribute("llm.token_count.completion", output_tokens)

    model = getattr(result, "model", None)
    if model is not None:
        span.set_attribute("llm.model_name", model)
