"""Tool span decorator."""

from __future__ import annotations

import asyncio
import contextlib
import functools
import json
from collections.abc import Callable
from typing import Any, TypeVar

from coalex.ext._base import _get_tracer

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


def tool_span(
    name: str | None = None,
    description: str | None = None,
) -> Callable[[F], F]:
    """Decorator that wraps a tool function with an OpenInference TOOL span.

    Supports both sync and async functions. The span captures:
    - ``openinference.span.kind`` = ``TOOL``
    - ``tool.name`` — span name
    - ``tool.description`` — if provided
    - ``tool.parameters`` — JSON-encoded kwargs passed to the function

    Args:
        name: Tool name and span name. Defaults to the function name.
        description: Human-readable tool description.
    """

    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 = _get_tracer()
                with tracer.start_as_current_span(span_name) as span:
                    span.set_attribute("openinference.span.kind", "TOOL")
                    span.set_attribute("tool.name", span_name)
                    if description:
                        span.set_attribute("tool.description", description)
                    if kwargs:
                        with contextlib.suppress(TypeError, ValueError):
                            span.set_attribute("tool.parameters", json.dumps(kwargs))
                    return await fn(*args, **kwargs)

            return async_wrapper  # type: ignore[return-value]

        @functools.wraps(fn)
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            tracer = _get_tracer()
            with tracer.start_as_current_span(span_name) as span:
                span.set_attribute("openinference.span.kind", "TOOL")
                span.set_attribute("tool.name", span_name)
                if description:
                    span.set_attribute("tool.description", description)
                if kwargs:
                    with contextlib.suppress(TypeError, ValueError):
                        span.set_attribute("tool.parameters", json.dumps(kwargs))
                return fn(*args, **kwargs)

        return wrapper  # type: ignore[return-value]

    return decorator
