import logging
from functools import wraps
from inspect import BoundArguments, iscoroutinefunction, signature
from typing import Any, Awaitable, Callable, Mapping, Optional, Sequence, TypeVar, cast, overload

from openinference.semconv.trace import OpenInferenceSpanKindValues, SpanAttributes
from opentelemetry import trace as trace_api
from opentelemetry.trace import NoOpTracer, Status, StatusCode, Tracer
from opentelemetry.util.types import AttributeValue
from typing_extensions import ParamSpec

logger = logging.getLogger(__name__)


def get_tracer(tracer: Optional[Tracer] = None) -> Tracer:
    """Get a tracer instance for tracing operations.

    This function follows a priority order:
    1. Use the provided Tracer if given
    2. Otherwise, pull from the global tracer provider
    3. Fall back to NoOpTracer if all else fails

    Args:
        tracer (Optional[Tracer]): Optional tracer to use. If None, will use global provider.

    Returns:
        Tracer: A tracer instance.
    """

    try:
        if tracer is not None:
            return tracer

        global_tracer_provider = trace_api.get_tracer_provider()
        return global_tracer_provider.get_tracer(__name__)

    except Exception:
        logger.debug("Failed to get tracer, falling back to NoOpTracer")
        return NoOpTracer()


FnParams = ParamSpec("FnParams")
ReturnValue = TypeVar("ReturnValue")


@overload
def trace(
    *,
    span_name: Optional[str] = None,
    span_kind: Optional[OpenInferenceSpanKindValues] = None,
    tracer: Optional[Tracer] = None,
    process_input: Optional[Mapping[str, Callable[[BoundArguments], Any]]] = None,
    process_output: Optional[Mapping[str, Callable[[ReturnValue], Any]]] = None,
) -> Callable[[Callable[FnParams, ReturnValue]], Callable[FnParams, ReturnValue]]: ...


@overload
def trace(
    *,
    span_name: Optional[str] = None,
    span_kind: Optional[OpenInferenceSpanKindValues] = None,
    tracer: Optional[Tracer] = None,
    process_input: Optional[Mapping[str, Callable[[BoundArguments], Any]]] = None,
    process_output: Optional[Mapping[str, Callable[[ReturnValue], Any]]] = None,
) -> Callable[
    [Callable[FnParams, Awaitable[ReturnValue]]], Callable[FnParams, Awaitable[ReturnValue]]
]: ...


def trace(
    *,
    span_name: Optional[str] = None,
    span_kind: Optional[OpenInferenceSpanKindValues] = None,
    tracer: Optional[Tracer] = None,
    process_input: Optional[Mapping[str, Callable[[BoundArguments], Any]]] = None,
    process_output: Optional[Mapping[str, Callable[[ReturnValue], Any]]] = None,
) -> Callable[[Callable[FnParams, Any]], Callable[FnParams, Any]]:
    """Trace the decorated function.

    If the decorated function has a `tracer` argument, it will be used to trace the function.
    Otherwise the global TracerProvider will be used.

    Args:
        span_name (Optional[str]): The name of the span to trace. If not provided, the function's
            __qualname__ will be used.
        span_kind (Optional[OpenInferenceSpanKindValues]): The kind of span to create.
        tracer (Optional[Tracer]): The tracer to use to trace the function. If not provided, the
            global TracerProvider will be used.
        process_input (Optional[Mapping[str, Callable[[BoundArguments], Any]]]): A mapping of
            attribute names to callables that will be called with the bound arguments to process the
            input arguments.
        process_output (Optional[Mapping[str, Callable[[ReturnValue], Any]]]): A mapping of
            attribute names to callables that will be called with the return value to process the
            output.

    Returns:
        Callable[[Callable[FnParams, Any]], Callable[FnParams, Any]]: A decorator function that
            wraps the input function with tracing capabilities.
    """

    def _decorator(
        func: Callable[FnParams, Any],
    ) -> Callable[FnParams, Any]:
        @wraps(func)
        def _wrapper_sync(*args: FnParams.args, **kwargs: FnParams.kwargs) -> ReturnValue:
            span_label = (
                span_name
                if span_name is not None
                else cast(str, getattr(func, "__qualname__", func.__name__))
            )
            _span_name: str = span_label

            bound: Optional[BoundArguments]
            try:
                sig = signature(func)
                bound = sig.bind_partial(*args, **kwargs)
                try:
                    bound.apply_defaults()
                except Exception:
                    pass
            except Exception:
                bound = None

            tracer_from_args: Optional[Tracer] = None
            if bound is not None:
                maybe_tracer = bound.arguments.get("tracer")
                if isinstance(maybe_tracer, Tracer):
                    tracer_from_args = maybe_tracer

            _tracer = get_tracer(tracer_from_args or tracer)
            function_raised: bool = False
            function_executed: bool = False
            result: Any = None
            try:
                with _tracer.start_as_current_span(_span_name) as span:
                    if span_kind is not None:
                        try:
                            span.set_attribute(
                                SpanAttributes.OPENINFERENCE_SPAN_KIND, span_kind.value
                            )
                        except Exception:
                            pass
                    if process_input is not None and bound is not None:
                        for attr_key, input_mapper_fn in process_input.items():
                            try:
                                value = input_mapper_fn(bound)
                                span.set_attribute(attr_key, _otel_attribute_value(value))
                            except Exception:
                                continue
                    try:
                        result = func(*args, **kwargs)
                        function_executed = True
                    except Exception as exc:
                        try:
                            span.record_exception(exc)
                            span.set_status(Status(StatusCode.ERROR))
                        except Exception:
                            pass
                        function_raised = True
                        raise

                    if process_output is not None:
                        for attr_key, output_mapper_fn in process_output.items():
                            try:
                                value = output_mapper_fn(result)
                                span.set_attribute(attr_key, _otel_attribute_value(value))
                            except Exception:
                                continue
                    return cast(ReturnValue, result)
            except Exception:
                if function_raised:
                    # Propagate the original function exception unchanged
                    raise
                if function_executed:
                    # If function ran successfully but tracing teardown failed, return its result
                    return cast(ReturnValue, result)
                try:
                    # Tracing failures should not affect business logic; run without tracing
                    return cast(ReturnValue, func(*args, **kwargs))
                except Exception:
                    raise

        @wraps(func)
        async def _wrapper_async(*args: FnParams.args, **kwargs: FnParams.kwargs) -> ReturnValue:
            span_label = (
                span_name
                if span_name is not None
                else cast(str, getattr(func, "__qualname__", func.__name__))
            )
            _span_name: str = span_label

            bound: Optional[BoundArguments]
            try:
                sig = signature(func)
                bound = sig.bind_partial(*args, **kwargs)
                try:
                    bound.apply_defaults()
                except Exception:
                    pass
            except Exception:
                bound = None

            tracer_from_args: Optional[Tracer] = None
            if bound is not None:
                maybe_tracer = bound.arguments.get("tracer")
                if isinstance(maybe_tracer, Tracer):
                    tracer_from_args = maybe_tracer

            _tracer = get_tracer(tracer_from_args or tracer)
            function_raised: bool = False
            function_executed: bool = False
            result: Any = None
            try:
                with _tracer.start_as_current_span(_span_name) as span:
                    if span_kind is not None:
                        try:
                            span.set_attribute(
                                SpanAttributes.OPENINFERENCE_SPAN_KIND, span_kind.value
                            )
                        except Exception:
                            pass
                    if process_input is not None and bound is not None:
                        for attr_key, input_mapper_fn in process_input.items():
                            try:
                                value = input_mapper_fn(bound)
                                span.set_attribute(attr_key, _otel_attribute_value(value))
                            except Exception:
                                continue
                    try:
                        result = await func(*args, **kwargs)
                        function_executed = True
                    except Exception as exc:
                        try:
                            span.record_exception(exc)
                            span.set_status(Status(StatusCode.ERROR))
                        except Exception:
                            pass
                        function_raised = True
                        raise

                    if process_output is not None:
                        for attr_key, output_mapper_fn in process_output.items():
                            try:
                                value = output_mapper_fn(result)
                                span.set_attribute(attr_key, _otel_attribute_value(value))
                            except Exception:
                                continue
                    return cast(ReturnValue, result)
            except Exception:
                if function_raised:
                    # Propagate the original function exception unchanged
                    raise
                if function_executed:
                    # If function ran successfully but tracing teardown failed, return its result
                    return cast(ReturnValue, result)
                try:
                    # Tracing failures should not affect business logic; run without tracing
                    return cast(ReturnValue, await cast(Callable[..., Any], func)(*args, **kwargs))
                except Exception:
                    raise

        if iscoroutinefunction(func):
            return cast(Callable[FnParams, Any], _wrapper_async)
        return cast(Callable[FnParams, Any], _wrapper_sync)

    return _decorator


def _otel_attribute_value(value: Any) -> AttributeValue:
    """Convert a value to an OpenTelemetry-compatible attribute value.

    Args:
        value (Any): The value to convert.

    Returns:
        AttributeValue: An OpenTelemetry-compatible attribute value.
    """
    if isinstance(value, (bool, int, float, str)):
        return value
    if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)):
        items = cast(Sequence[object], value)
        return [str(item) for item in items]
    return str(value)
