"""@observe decorator for tracing functions as spans or generations."""

import inspect
from functools import wraps
from typing import Any, Callable, Optional, TypeVar

from mingx._client.get_client import get_client

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


def observe(
    name: Optional[str] = None,
    *,
    as_type: str = "span",
    capture_input: bool = True,
    capture_output: bool = True,
    **kwargs: Any,
) -> Callable[[F], F]:
    """Decorate a function to run inside a Mingx observation (span or generation).

    name: Span name; defaults to the function's __name__.
    as_type: "span" or "generation".
    capture_input: If True, record the first argument or kwargs as input.
    capture_output: If True, record the return value as output.
    **kwargs: Passed to start_as_current_observation (e.g. model=..., metadata=...).
    """

    def decorator(fn: F) -> F:
        _name = name or getattr(fn, "__name__", "observed")

        @wraps(fn)
        def sync_wrapper(*args: Any, **fn_kwargs: Any) -> Any:
            client = get_client()
            input_val: Any = None
            if capture_input:
                input_val = fn_kwargs if fn_kwargs else (args[0] if len(args) == 1 else args)
            with client.start_as_current_observation(
                name=_name, as_type=as_type, input=input_val, **kwargs
            ) as obs:
                try:
                    result = fn(*args, **fn_kwargs)
                    if capture_output:
                        obs.update(output=result)
                    return result
                except Exception as e:
                    obs.update(level="ERROR", status_message=str(e))
                    raise

        @wraps(fn)
        async def async_wrapper(*args: Any, **fn_kwargs: Any) -> Any:
            client = get_client()
            input_val = None
            if capture_input:
                input_val = fn_kwargs if fn_kwargs else (args[0] if len(args) == 1 else args)
            with client.start_as_current_observation(
                name=_name, as_type=as_type, input=input_val, **kwargs
            ) as obs:
                try:
                    result = await fn(*args, **fn_kwargs)
                    if capture_output:
                        obs.update(output=result)
                    return result
                except Exception as e:
                    obs.update(level="ERROR", status_message=str(e))
                    raise

        if inspect.iscoroutinefunction(fn):
            return async_wrapper  # type: ignore[return-value]
        return sync_wrapper  # type: ignore[return-value]

    return decorator
