"""Guardrail span decorator."""

from __future__ import annotations

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

from opentelemetry.trace import Status, StatusCode

from coalex.ext._base import _get_tracer

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


def guardrail_span(
    name: str | None = None,
    guardrail_name: str | None = None,
) -> Callable[[F], F]:
    """Decorator that wraps a guardrail function with an OpenInference GUARDRAIL span.

    Supports both sync and async functions. The span captures:
    - ``openinference.span.kind`` = ``GUARDRAIL``
    - ``guardrail.name`` — from ``guardrail_name`` or the function name
    - ``guardrail.result`` — ``"pass"`` if no exception raised, ``"fail"`` otherwise
    - ``guardrail.score`` — set this via the return value if it is a ``float``

    Args:
        name: Span name. Defaults to the function name.
        guardrail_name: Logical guardrail name recorded as an attribute.
    """

    def decorator(fn: F) -> F:
        span_name = name or fn.__name__
        logical_name = guardrail_name or span_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", "GUARDRAIL")
                    span.set_attribute("guardrail.name", logical_name)
                    try:
                        result = await fn(*args, **kwargs)
                        span.set_attribute("guardrail.result", "pass")
                        if isinstance(result, float):
                            span.set_attribute("guardrail.score", result)
                        span.set_status(Status(StatusCode.OK))
                        return result
                    except Exception:
                        span.set_attribute("guardrail.result", "fail")
                        raise

            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", "GUARDRAIL")
                span.set_attribute("guardrail.name", logical_name)
                try:
                    result = fn(*args, **kwargs)
                    span.set_attribute("guardrail.result", "pass")
                    if isinstance(result, float):
                        span.set_attribute("guardrail.score", result)
                    span.set_status(Status(StatusCode.OK))
                    return result
                except Exception:
                    span.set_attribute("guardrail.result", "fail")
                    raise

        return wrapper  # type: ignore[return-value]

    return decorator
