import asyncio
from typing import Any

from ..constants import (
    INPUT_MIME_TYPE,
    INPUT_VALUE,
    OPENINFERENCE_SPAN_KIND,
    OUTPUT_MIME_TYPE,
    OUTPUT_VALUE,
    Status,
    StatusCode,
    TOOL,
    context_api,
)
from ..utils import bind_arguments, get_input_value

from openinference.instrumentation import get_attributes_from_context

# ============================================================================
# Batch Tool Execution Wrapper
# ============================================================================


class ExecuteToolsWrapper:
    """
    Wrapper for `Agent._execute_tool_calls()` to emit TOOL spans for batched calls.

    This wrapper instruments the `_execute_tool_calls()` helper, creating
    TOOL spans that capture batch tool execution operations within an agent's
    workflow with comprehensive tracing and error handling.

    Key features:
    - Captures batch tool execution flow and coordination
    - Handles both async and sync tool execution patterns
    - Proper input/output value extraction and JSON serialization
    - Error handling with span status management and exception recording
    - Agent context preservation and attribute extraction
    - Maintains compatibility with OpenInference TOOL span standards

    The wrapper creates spans with TOOL span kind to represent the batch
    coordination of multiple tool executions within agent workflows.
    """

    def __init__(self, tracer: Any) -> None:
        """
        Initialize the execute tools wrapper with OpenTelemetry tracer.

        Args:
            tracer (Any): OpenTelemetry tracer instance for creating TOOL spans
        """
        self._tracer = tracer

    def __call__(self, wrapped: Any, instance: Any, args: Any, kwargs: Any) -> Any:
        """
        Wrap `_execute_tool_calls()` with comprehensive TOOL span tracing.

        Args:
            wrapped (callable): `_execute_tool_calls` coroutine to be instrumented
            instance (Agent): Agent instance issuing the tool calls
            args (tuple): Positional arguments passed to the helper
            kwargs (dict): Keyword arguments forwarded to the helper

        Returns:
            Any: Result from wrapped method execution with span attributes capturing
                 batch tool execution context, inputs, and aggregated results
        """
        # Check for instrumentation suppression
        if context_api and context_api.get_value(
            context_api._SUPPRESS_INSTRUMENTATION_KEY
        ):
            return wrapped(*args, **kwargs)

        agent_name = getattr(instance, "name", instance.__class__.__name__)
        span_name = f"{agent_name}.tool_batch"

        # Build span attributes
        attributes = {
            OPENINFERENCE_SPAN_KIND: TOOL,
            INPUT_VALUE: get_input_value(wrapped, *args, **kwargs),
            INPUT_MIME_TYPE: "application/json",
            "agent.name": agent_name,
        }

        # Add context attributes
        attributes.update(get_attributes_from_context())

        # Handle async vs sync execution
        if asyncio.iscoroutinefunction(wrapped):
            return self._handle_async_execution(
                wrapped, args, kwargs, span_name, attributes
            )
        else:
            return self._handle_sync_execution(
                wrapped, args, kwargs, span_name, attributes
            )

    def _handle_async_execution(
        self, wrapped: Any, args: Any, kwargs: Any, span_name: str, attributes: dict
    ) -> Any:
        """
        Handle asynchronous batch tool execution with comprehensive span tracing.

        Args:
            wrapped (callable): `_execute_tool_calls`
            args (tuple): Positional arguments for the wrapped method
            kwargs (dict): Keyword arguments for the wrapped method
            span_name (str): Name for the created span (e.g., "MyAgent.tool_batch")
            attributes (dict): Span attributes including agent context and tool information

        Returns:
            Any: Coroutine that executes the wrapped method with proper span instrumentation,
                 capturing batch tool execution and result aggregation
        """

        async def async_wrapper():
            with self._tracer.start_as_current_span(
                span_name, attributes=attributes
            ) as span:
                try:
                    result = await wrapped(*args, **kwargs)
                    span.set_attribute(OUTPUT_VALUE, str(result))
                    span.set_attribute(OUTPUT_MIME_TYPE, "application/json")
                    span.set_status(Status(StatusCode.OK))
                    return result
                except Exception as e:
                    span.set_status(Status(StatusCode.ERROR, str(e)))
                    span.record_exception(e)
                    raise

        return async_wrapper()

    def _handle_sync_execution(
        self, wrapped: Any, args: Any, kwargs: Any, span_name: str, attributes: dict
    ) -> Any:
        """
        Handle synchronous batch tool execution with comprehensive span tracing.

        Args:
            wrapped (callable): `_execute_tool_calls`
            args (tuple): Positional arguments for the wrapped method
            kwargs (dict): Keyword arguments for the wrapped method
            span_name (str): Name for the created span (e.g., "MyAgent.tool_batch")
            attributes (dict): Span attributes including agent context and tool information

        Returns:
            Any: Result from wrapped method execution with proper span instrumentation,
                 capturing batch tool execution and result aggregation
        """
        with self._tracer.start_as_current_span(
            span_name, attributes=attributes
        ) as span:
            try:
                result = wrapped(*args, **kwargs)
                span.set_attribute(OUTPUT_VALUE, str(result))
                span.set_attribute(OUTPUT_MIME_TYPE, "application/json")
                span.set_status(Status(StatusCode.OK))
                return result
            except Exception as e:
                span.set_status(Status(StatusCode.ERROR, str(e)))
                span.record_exception(e)
                raise


# ============================================================================
# Individual Tool Execution Wrapper
# ============================================================================


class RunToolWrapper:
    """
    Wrapper for `AgentToolExecutor.run_tool()` to emit TOOL spans per tool call.

    This wrapper instruments the executor layer, creating a span for each
    tool function invocation regardless of which agent issued it.

    Key features:
    - Captures individual tool execution with tool name extraction and identification
    - Handles both async and sync tool execution patterns
    - Proper input/output value extraction and JSON serialization
    - Error handling with span status management and exception recording
    - Tool name-based span naming for clear hierarchy and identification
    - Comprehensive attribute extraction including tool metadata

    The wrapper creates spans with TOOL span kind and uses the actual tool
    name as the span name for clear identification in trace hierarchies.
    """

    def __init__(self, tracer: Any) -> None:
        """
        Initialize the run tool wrapper with OpenTelemetry tracer.

        Args:
            tracer (Any): OpenTelemetry tracer instance for creating individual TOOL spans
        """
        self._tracer = tracer

    def __call__(self, wrapped: Any, instance: Any, args: Any, kwargs: Any) -> Any:
        """
        Wrap the executor method with TOOL span tracing.

        Args:
            wrapped (callable): Original `AgentToolExecutor.run_tool`
            instance (AgentToolExecutor): Executor instance
            args (tuple): Positional arguments - typically (tool_name, *tool_args)
            kwargs (dict): Keyword arguments passed to the original method

        Returns:
            Any: Result from wrapped method execution with span attributes capturing
                 individual tool execution context, inputs, outputs, and tool metadata
        """
        # Check for instrumentation suppression
        if context_api and context_api.get_value(
            context_api._SUPPRESS_INSTRUMENTATION_KEY
        ):
            return wrapped(*args, **kwargs)

        # Extract agent and tool information for span naming
        agent_name = getattr(instance, "name", instance.__class__.__name__)
        arguments = bind_arguments(wrapped, *args, **kwargs)
        tool_name = arguments.get("tool_name", args[0] if args else "unknown_tool")
        span_name = f"{tool_name}"

        # Build span attributes
        attributes = {
            OPENINFERENCE_SPAN_KIND: TOOL,
            INPUT_VALUE: get_input_value(wrapped, *args, **kwargs),
            INPUT_MIME_TYPE: "application/json",
            "agent.name": agent_name,
            "tool.name": tool_name,
        }

        # Add context attributes
        attributes.update(get_attributes_from_context())

        # Handle async vs sync execution
        if asyncio.iscoroutinefunction(wrapped):
            return self._handle_async_execution(
                wrapped, args, kwargs, span_name, attributes
            )
        else:
            return self._handle_sync_execution(
                wrapped, args, kwargs, span_name, attributes
            )

    def _handle_async_execution(
        self, wrapped: Any, args: Any, kwargs: Any, span_name: str, attributes: dict
    ) -> Any:
        """
        Handle asynchronous individual tool execution with comprehensive span tracing.

        Manages async run_tool execution by creating spans with proper
        attribute handling, tool result processing, and error management for
        individual tool execution workflows.

        Args:
            wrapped (callable): Original async run_tool method to execute
            args (tuple): Positional arguments for the wrapped method
            kwargs (dict): Keyword arguments for the wrapped method
            span_name (str): Name for the created span (typically the tool name)
            attributes (dict): Span attributes including agent context and tool metadata

        Returns:
            Any: Coroutine that executes the wrapped method with proper span instrumentation,
                 capturing individual tool execution and result processing
        """

        async def async_wrapper():
            with self._tracer.start_as_current_span(
                span_name, attributes=attributes
            ) as span:
                try:
                    result = await wrapped(*args, **kwargs)
                    span.set_attribute(OUTPUT_VALUE, str(result))
                    span.set_attribute(OUTPUT_MIME_TYPE, "application/json")
                    span.set_status(Status(StatusCode.OK))
                    return result
                except Exception as e:
                    span.set_status(Status(StatusCode.ERROR, str(e)))
                    span.record_exception(e)
                    raise

        return async_wrapper()

    def _handle_sync_execution(
        self, wrapped: Any, args: Any, kwargs: Any, span_name: str, attributes: dict
    ) -> Any:
        """
        Handle synchronous individual tool execution with comprehensive span tracing.

        Manages sync run_tool execution by creating spans with proper
        attribute handling, tool result processing, and error management for
        individual tool execution workflows.

        Args:
            wrapped (callable): Original sync run_tool method to execute
            args (tuple): Positional arguments for the wrapped method
            kwargs (dict): Keyword arguments for the wrapped method
            span_name (str): Name for the created span (typically the tool name)
            attributes (dict): Span attributes including agent context and tool metadata

        Returns:
            Any: Result from wrapped method execution with proper span instrumentation,
                 capturing individual tool execution and result processing
        """
        with self._tracer.start_as_current_span(
            span_name, attributes=attributes
        ) as span:
            try:
                result = wrapped(*args, **kwargs)
                span.set_attribute(OUTPUT_VALUE, str(result))
                span.set_attribute(OUTPUT_MIME_TYPE, "application/json")
                span.set_status(Status(StatusCode.OK))
                return result
            except Exception as e:
                span.set_status(Status(StatusCode.ERROR, str(e)))
                span.record_exception(e)
                raise


# ============================================================================
# Exported Classes
# ============================================================================

__all__ = [
    "ExecuteToolsWrapper",
    "RunToolWrapper",
]
