"""Tracing decorators for agent operations and tool executions."""

import functools
import inspect
from collections.abc import Callable
from inspect import BoundArguments
from typing import ParamSpec, TypeVar, cast

from opentelemetry.trace import Span, SpanKind, StatusCode

from .telemetry import get_tracer

P = ParamSpec("P")
R = TypeVar("R")


def _create_tracing_wrapper(
    func: Callable[P, R],
    span_name: str,
    span_kind: SpanKind,
    attribute_extractor: Callable[[BoundArguments, Span], None],
) -> Callable[P, R]:
    """Internal function to create tracing wrappers for both sync and async functions.

    Args:
        func: The function to wrap
        span_name: Name for the OpenTelemetry span
        span_kind: Kind of span (INTERNAL, CLIENT, etc.)
        attribute_extractor: Callable to extract and set span attributes from function arguments

    Returns:
        Wrapped function with tracing
    """

    @functools.wraps(func)
    def sync_wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
        tracer = get_tracer()

        with tracer.start_as_current_span(span_name, kind=span_kind) as span:
            try:
                # Extract and set attributes
                sig = inspect.signature(func)
                bound_args = sig.bind(*args, **kwargs)
                bound_args.apply_defaults()
                attribute_extractor(bound_args, span)

                # Execute the function
                result = func(*args, **kwargs)
                return result

            except Exception as e:
                # Record exception and set error status
                span.record_exception(e)
                span.set_status(StatusCode.ERROR, str(e))
                raise

    @functools.wraps(func)
    async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
        tracer = get_tracer()

        with tracer.start_as_current_span(span_name, kind=span_kind) as span:
            try:
                # Extract and set attributes
                sig = inspect.signature(func)
                bound_args = sig.bind(*args, **kwargs)
                bound_args.apply_defaults()
                attribute_extractor(bound_args, span)

                # Execute the async function
                result = await func(*args, **kwargs)  # type: ignore[misc]  # Generic async wrapper limitation
                return result  # type: ignore[no-any-return]  # Generic async wrapper limitation

            except Exception as e:
                # Record exception and set error status
                span.record_exception(e)
                span.set_status(StatusCode.ERROR, str(e))
                raise

    # Return appropriate wrapper based on whether function is async
    if inspect.iscoroutinefunction(func):
        return cast(Callable[P, R], async_wrapper)
    else:
        return cast(Callable[P, R], sync_wrapper)


def trace_agent_operation(
    operation_name: str,
) -> Callable[[Callable[P, R]], Callable[P, R]]:
    """Decorator that adds OpenTelemetry tracing to agent operations.

    Automatically creates spans with coordinator-specific attributes.

    Args:
        operation_name: Name of the agent operation being traced

    Returns:
        Decorated function with tracing

    Example:
        >>> @trace_agent_operation(operation_name="assign_agent")
        >>> async def assign_agent(issue_id: int, agent_type: str) -> bool:
        >>>     # Implementation
        >>>     return True
    """

    def attribute_extractor(bound_args: BoundArguments, span: Span) -> None:
        """Extract agent-specific attributes from function arguments."""
        if "issue_id" in bound_args.arguments:
            span.set_attribute("agent.issue_id", bound_args.arguments["issue_id"])
        if "issue_number" in bound_args.arguments:
            span.set_attribute("agent.issue_number", bound_args.arguments["issue_number"])
        if "agent_type" in bound_args.arguments:
            span.set_attribute("agent.agent_type", bound_args.arguments["agent_type"])
        if "agent_id" in bound_args.arguments:
            span.set_attribute("agent.agent_id", bound_args.arguments["agent_id"])

    def decorator(func: Callable[P, R]) -> Callable[P, R]:
        span_name = f"agent.{operation_name}"
        return _create_tracing_wrapper(func, span_name, SpanKind.INTERNAL, attribute_extractor)

    return decorator


def trace_tool_execution(
    tool_name: str,
) -> Callable[[Callable[P, R]], Callable[P, R]]:
    """Decorator that adds OpenTelemetry tracing to tool executions.

    Automatically creates spans with tool-specific attributes.

    Args:
        tool_name: Name of the tool being executed

    Returns:
        Decorated function with tracing

    Example:
        >>> @trace_tool_execution(tool_name="quality_gate")
        >>> async def run_quality_gate(issue_id: int) -> bool:
        >>>     # Implementation
        >>>     return True
    """

    def attribute_extractor(bound_args: BoundArguments, span: Span) -> None:
        """Extract tool-specific attributes from function arguments."""
        span.set_attribute("tool.name", tool_name)

        if "issue_id" in bound_args.arguments:
            span.set_attribute("tool.issue_id", bound_args.arguments["issue_id"])
        if "issue_number" in bound_args.arguments:
            span.set_attribute("tool.issue_number", bound_args.arguments["issue_number"])

    def decorator(func: Callable[P, R]) -> Callable[P, R]:
        span_name = f"tool.{tool_name}"
        return _create_tracing_wrapper(func, span_name, SpanKind.CLIENT, attribute_extractor)

    return decorator
