#
# Copyright 2026 The Dapr Authors
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#     http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import asyncio
import logging
from typing import Any

from ..constants import (
    AGENT,
    CHAIN,
    INPUT_MIME_TYPE,
    INPUT_VALUE,
    OUTPUT_MIME_TYPE,
    OUTPUT_VALUE,
    OPENINFERENCE_SPAN_KIND,
    GEN_AI_OPERATION_NAME,
    GEN_AI_AGENT_NAME,
    GenAiOperationNameValues,
    Status,
    StatusCode,
    context_api,
)
from ..utils import bind_arguments, extract_content_from_result, get_input_value
from openinference.instrumentation import get_attributes_from_context

logger = logging.getLogger(__name__)


# ============================================================================
# Agent Execution Wrapper
# ============================================================================


class AgentRunWrapper:
    """
    Wrapper for `Agent.run()` to create the top-level AGENT span for standalone runs.

    This only targets the in-process `Agent` (non-durable) implementation. Durable
    agents execute inside Dapr Workflows and are therefore covered by the workflow
    wrappers in this package.

    Key features:
    - Extracts comprehensive agent metadata (name, role, goal, tools, max_iterations)
    - Handles input_data processing with automatic user query extraction
    - Supports both async and sync agent execution patterns
    - Proper error handling with span status management and exception recording
    - Phoenix UI compatibility with OpenInference attribute standards
    - GenAI semconv dual-emit (gen_ai.operation.name, gen_ai.agent.name)

    The wrapper creates spans with AGENT span kind and captures the complete
    agent execution flow from input processing to final result generation.
    """

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

        Args:
            tracer: OpenTelemetry tracer instance
        """
        self._tracer = tracer

    def __call__(self, wrapped: Any, instance: Any, args: Any, kwargs: Any) -> Any:
        """
        Wrap Agent.run() method with comprehensive AGENT span tracing.

        Creates a root AGENT span that captures the complete agent execution
        lifecycle, including metadata extraction, input processing, and result
        handling with proper OpenTelemetry context propagation.

        Args:
            wrapped (callable): Original Agent.run method to be instrumented
            instance (Agent): Agent instance containing metadata (name, role, goal, tools)
            args (tuple): Positional arguments passed to Agent.run, typically (input_data,)
            kwargs (dict): Keyword arguments passed to the original method

        Returns:
            Any: Result from wrapped method execution with span attributes capturing
                 agent metadata, input queries, and execution outcomes
        """
        # Check for instrumentation suppression
        if context_api and context_api.get_value(
            context_api._SUPPRESS_INSTRUMENTATION_KEY
        ):
            return wrapped(*args, **kwargs)

        # Skip DurableAgent instances - they use WorkflowMonitorWrapper instead
        if instance.__class__.__name__ == "DurableAgent":
            return wrapped(*args, **kwargs)

        # Extract agent information for span naming and attributes
        agent_name = getattr(instance, "name", instance.__class__.__name__)
        span_name = f"invoke_agent {agent_name}"

        # Process input arguments to extract user query
        arguments = bind_arguments(wrapped, *args, **kwargs)
        input_data = arguments.get("input_data")

        # Build comprehensive span attributes
        attributes = self._build_agent_attributes(instance, agent_name, input_data)

        # 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 _build_agent_attributes(
        self, instance: Any, agent_name: str, input_data: Any
    ) -> dict:
        """
        Build comprehensive span attributes for agent execution tracing.

        Extracts agent metadata (profile role/goal, execution settings, tools, etc.)
        and processes user input so spans render meaningful details in Phoenix.

        Args:
            instance (Agent): Agent instance with metadata attributes (role, goal, tools, etc.)
            agent_name (str): Agent name extracted for span identification
            input_data (Any): Raw input data from Agent.run arguments, processed for user query extraction

        Returns:
            dict: Comprehensive span attributes including agent metadata, tool information,
                  processed input value, and OpenInference-compatible attribute structure
        """
        attributes = {
            OPENINFERENCE_SPAN_KIND: AGENT,
            INPUT_MIME_TYPE: "application/json",
            OUTPUT_MIME_TYPE: "application/json",
            "agent.name": agent_name,
            # GenAI semconv
            GEN_AI_OPERATION_NAME: GenAiOperationNameValues.INVOKE_AGENT,
            GEN_AI_AGENT_NAME: agent_name,
        }

        profile = getattr(instance, "profile", None)
        role = getattr(profile, "role", getattr(instance, "role", None))
        goal = getattr(profile, "goal", getattr(instance, "goal", None))
        execution = getattr(instance, "execution", None)

        if role is not None:
            attributes["agent.role"] = role
        if goal is not None:
            attributes["agent.goal"] = goal

        tools = [
            getattr(tool, "name", getattr(tool, "__name__", str(tool)))
            for tool in getattr(instance, "tools", [])
        ]
        attributes["agent.tools"] = tools
        attributes["agent.tools.count"] = len(tools)

        max_iterations = getattr(execution, "max_iterations", None)
        if max_iterations is not None:
            attributes["agent.execution.max_iterations"] = max_iterations

        # Extract actual input value - the user's query/request
        if input_data:
            if isinstance(input_data, str):
                attributes[INPUT_VALUE] = input_data
            elif isinstance(input_data, dict):
                # Extract the actual user input from input_data
                user_input = input_data.get("input_data") or str(input_data)
                attributes[INPUT_VALUE] = user_input
            else:
                attributes[INPUT_VALUE] = str(input_data)
        else:
            attributes[INPUT_VALUE] = ""

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

        return attributes

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

        Args:
            wrapped (callable): Original async Agent.run 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 (e.g., "invoke_agent MyAgent")
            attributes (dict): Comprehensive span attributes including agent metadata

        Returns:
            Any: Coroutine that executes the wrapped method with proper span instrumentation,
                 result extraction, and status management
        """

        async def async_wrapper():
            with self._tracer.start_as_current_span(
                span_name, attributes=attributes
            ) as span:
                # Store span ID for cross-process parent-child relationship
                try:
                    from ..context_propagation import extract_otel_context
                    from ..context_storage import store_workflow_context

                    # Extract current context from the AGENT span
                    current_context = extract_otel_context()
                    if current_context.get("traceparent"):
                        # Store the span ID for workflow tasks to use as parent
                        span_id = span.get_span_context().span_id
                        trace_id = span.get_span_context().trace_id

                        # Store span context with trace and span IDs
                        span_context = {
                            "trace_id": format(
                                trace_id, "032x"
                            ),  # Convert to 32-char hex string
                            "span_id": format(
                                span_id, "016x"
                            ),  # Convert to 16-char hex string
                            "traceparent": current_context.get("traceparent"),
                            "tracestate": current_context.get("tracestate", ""),
                        }

                        # Store under multiple keys for different lookup patterns
                        store_workflow_context(
                            "__current_workflow_context__", span_context
                        )
                        logger.debug(
                            f"Stored Agent span context: trace_id={format(trace_id, '032x')}, span_id={format(span_id, '016x')}"
                        )
                    else:
                        logger.warning("No traceparent found in AGENT span context")
                except Exception as e:
                    logger.warning(f"Failed to store span context: {e}")

                try:
                    result = await wrapped(*args, **kwargs)

                    # Extract and set output value
                    output_content = extract_content_from_result(result)
                    span.set_attribute(OUTPUT_VALUE, output_content)

                    span.set_status(Status(StatusCode.OK))
                    return result
                except Exception as e:
                    span.set_status(Status(StatusCode.ERROR, str(e)))
                    span.set_attribute("error.type", type(e).__qualname__)
                    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 agent execution with comprehensive span tracing.

        Args:
            wrapped (callable): Original sync Agent.run 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 (e.g., "invoke_agent MyAgent")
            attributes (dict): Comprehensive span attributes including agent metadata

        Returns:
            Any: Result from wrapped method execution with proper span instrumentation,
                 result extraction, and status management
        """
        with self._tracer.start_as_current_span(
            span_name, attributes=attributes
        ) as span:
            # Store span ID for cross-process parent-child relationship
            try:
                from ..context_propagation import extract_otel_context

                # Extract current context from the AGENT span
                current_context = extract_otel_context()
                if current_context.get("traceparent"):
                    # Store the span ID for workflow tasks to use as parent
                    span_id = span.get_span_context().span_id
                    trace_id = span.get_span_context().trace_id

                    # Context will be stored per-instance by WorkflowMonitorWrapper
                    # to avoid cross-instance contamination
                    logger.debug(
                        f"Stored Agent span context: trace_id={format(trace_id, '032x')}, span_id={format(span_id, '016x')}"
                    )
                else:
                    logger.warning("No traceparent found in AGENT span context")
            except Exception as e:
                logger.warning(f"Failed to store span context: {e}")

            try:
                result = wrapped(*args, **kwargs)

                # Extract and set output value
                output_content = extract_content_from_result(result)
                span.set_attribute(OUTPUT_VALUE, output_content)

                span.set_status(Status(StatusCode.OK))
                return result
            except Exception as e:
                span.set_status(Status(StatusCode.ERROR, str(e)))
                span.set_attribute("error.type", type(e).__qualname__)
                span.record_exception(e)
                raise


# ============================================================================
# Process Iterations Wrapper
# ============================================================================


class ProcessIterationsWrapper:
    """
    Wrapper for `Agent._conversation_loop()` to create CHAIN spans per reasoning turn.

    The standalone agent now performs its entire loop inside `_conversation_loop`;
    this wrapper surfaces each iteration so Phoenix shows the step-by-step progress.

    Key features:
    - Captures iterative processing flow and reasoning steps
    - Handles both async and sync processing patterns
    - Proper input/output value extraction and serialization
    - Error handling with span status management and exception recording
    - Maintains consistent tracing hierarchy as child of AGENT spans

    The wrapper creates spans with CHAIN span kind to represent the sequential
    processing logic that occurs during agent reasoning and decision-making.
    """

    def __init__(self, tracer: Any) -> None:
        """
        Initialize the process iterations wrapper.

        Args:
            tracer: OpenTelemetry tracer instance
        """
        self._tracer = tracer

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

        Args:
            wrapped (callable): `_conversation_loop`
            instance (Agent): Agent instance containing metadata and configuration
            args (tuple): Positional arguments passed to the method
            kwargs (dict): Keyword arguments passed to the original method

        Returns:
            Any: Result from wrapped method execution with span attributes capturing
                 processing flow, input/output values, and execution context
        """
        # Check for instrumentation suppression
        if context_api and context_api.get_value(
            context_api._SUPPRESS_INSTRUMENTATION_KEY
        ):
            return wrapped(*args, **kwargs)

        # Extract agent information for span naming
        agent_name = getattr(instance, "name", instance.__class__.__name__)
        span_name = f"{agent_name}._conversation_loop"

        # Build span attributes
        attributes = {
            OPENINFERENCE_SPAN_KIND: CHAIN,
            INPUT_VALUE: get_input_value(wrapped, *args, **kwargs),
            INPUT_MIME_TYPE: "application/json",
            OUTPUT_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 process iterations execution with comprehensive span tracing.

        Args:
            wrapped (callable): `_conversation_loop`
            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._conversation_loop")
            attributes (dict): Span attributes including agent name and processing context

        Returns:
            Any: Coroutine that executes the wrapped method with proper span instrumentation,
                 capturing processing iterations and reasoning flow
        """

        async def async_wrapper():
            with self._tracer.start_as_current_span(
                span_name, attributes=attributes
            ) as span:
                try:
                    result = await wrapped(*args, **kwargs)

                    # Extract and set output value
                    output_content = extract_content_from_result(result)
                    span.set_attribute(OUTPUT_VALUE, output_content)

                    span.set_status(Status(StatusCode.OK))
                    return result
                except Exception as e:
                    span.set_status(Status(StatusCode.ERROR, str(e)))
                    span.set_attribute("error.type", type(e).__qualname__)
                    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 conversation execution with comprehensive span tracing.

        Args:
            wrapped (callable): `_conversation_loop`
            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._conversation_loop")
            attributes (dict): Span attributes including agent name and processing context

        Returns:
            Any: Result from wrapped method execution with proper span instrumentation,
                 capturing processing iterations and reasoning flow
        """
        with self._tracer.start_as_current_span(
            span_name, attributes=attributes
        ) as span:
            try:
                result = wrapped(*args, **kwargs)

                # Extract and set output value
                output_content = extract_content_from_result(result)
                span.set_attribute(OUTPUT_VALUE, output_content)

                span.set_status(Status(StatusCode.OK))
                return result
            except Exception as e:
                span.set_status(Status(StatusCode.ERROR, str(e)))
                span.set_attribute("error.type", type(e).__qualname__)
                span.record_exception(e)
                raise


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

__all__ = [
    "AgentRunWrapper",
    "ProcessIterationsWrapper",
]
