"""@observe decorator for automatic span creation."""
from __future__ import annotations

import asyncio
import functools
from typing import Any, Callable

from fibinaci.tracing.context import span


def observe(
    name: str | None = None,
    *,
    span_kind: str = "custom",
    **kwargs: Any,
) -> Callable:
    """Decorator that wraps a function in a span.

    Works with both sync and async functions.

    Usage::

        @observe("my-step", span_kind="tool_call")
        def my_function():
            ...

        @observe()
        async def my_async_function():
            ...
    """

    def decorator(fn: Callable) -> Callable:
        span_name = name or fn.__qualname__

        if asyncio.iscoroutinefunction(fn):
            @functools.wraps(fn)
            async def async_wrapper(*args: Any, **kw: Any) -> Any:
                with span(span_name, span_kind=span_kind, **kwargs) as _sd:
                    result = await fn(*args, **kw)
                return result

            return async_wrapper
        else:
            @functools.wraps(fn)
            def sync_wrapper(*args: Any, **kw: Any) -> Any:
                with span(span_name, span_kind=span_kind, **kwargs) as _sd:
                    result = fn(*args, **kw)
                return result

            return sync_wrapper

    return decorator
