import sys
from functools import wraps

import sentry_sdk
from sentry_sdk.consts import SPANDATA
from sentry_sdk.integrations import DidNotEnable
from sentry_sdk.utils import capture_internal_exceptions, reraise

from ..spans import agent_workflow_span, update_invoke_agent_span
from ..utils import _capture_exception

try:
    from agents.exceptions import AgentsException
except ImportError:
    raise DidNotEnable("OpenAI Agents not installed")

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from typing import Any, AsyncIterator, Callable


def _create_run_wrapper(original_func: "Callable[..., Any]") -> "Callable[..., Any]":
    """
    Wraps the agents.Runner.run methods to
    - create and manage a root span for the agent workflow runs.
    - end the agent invocation span if an `AgentsException` is raised in `run()`.

    Note agents.Runner.run_sync() is a wrapper around agents.Runner.run(),
    so it does not need to be wrapped separately.
    """

    @wraps(original_func)
    async def wrapper(*args: "Any", **kwargs: "Any") -> "Any":
        # Isolate each workflow so that when agents are run in asyncio tasks they
        # don't touch each other's scopes
        with sentry_sdk.isolation_scope():
            # Clone agent because agent invocation spans are attached per run.
            if "starting_agent" in kwargs:
                agent = kwargs["starting_agent"].clone()
            else:
                agent = args[0].clone()

            with agent_workflow_span(agent) as workflow_span:
                # Set conversation ID on workflow span early so it's captured even on errors
                conversation_id = kwargs.get("conversation_id")
                if conversation_id:
                    agent._sentry_conversation_id = conversation_id
                    workflow_span.set_data(
                        SPANDATA.GEN_AI_CONVERSATION_ID, conversation_id
                    )

                if "starting_agent" in kwargs:
                    kwargs["starting_agent"] = agent
                else:
                    args = (agent, *args[1:])

                try:
                    run_result = await original_func(*args, **kwargs)
                except AgentsException as exc:
                    exc_info = sys.exc_info()
                    with capture_internal_exceptions():
                        _capture_exception(exc)

                        context_wrapper = getattr(exc.run_data, "context_wrapper", None)
                        if context_wrapper is not None:
                            invoke_agent_span = getattr(
                                context_wrapper, "_sentry_agent_span", None
                            )

                            if (
                                invoke_agent_span is not None
                                and invoke_agent_span.timestamp is None
                            ):
                                update_invoke_agent_span(
                                    span=invoke_agent_span,
                                    context=context_wrapper,
                                    agent=agent,
                                )

                                invoke_agent_span.__exit__(*exc_info)
                                delattr(context_wrapper, "_sentry_agent_span")
                    reraise(*exc_info)
                except Exception as exc:
                    exc_info = sys.exc_info()
                    with capture_internal_exceptions():
                        # Invoke agent span is not finished in this case.
                        # This is much less likely to occur than other cases because
                        # AgentRunner.run() is "just" a while loop around _run_single_turn.
                        _capture_exception(exc)
                    reraise(*exc_info)

                invoke_agent_span = getattr(
                    run_result.context_wrapper, "_sentry_agent_span", None
                )
                if not invoke_agent_span:
                    return run_result

                update_invoke_agent_span(
                    span=invoke_agent_span,
                    context=run_result.context_wrapper,
                    agent=agent,
                )

                invoke_agent_span.__exit__(None, None, None)
                delattr(run_result.context_wrapper, "_sentry_agent_span")
                return run_result

    return wrapper


def _create_run_streamed_wrapper(
    original_func: "Callable[..., Any]",
) -> "Callable[..., Any]":
    """
    Wraps the agents.Runner.run_streamed method to
    - create a root span for streaming agent workflow runs.
    - end the workflow span if and only if the response stream is consumed or cancelled.

    Unlike run(), run_streamed() returns immediately with a RunResultStreaming object
    while execution continues in a background task. The workflow span must stay open
    throughout the streaming operation and close when streaming completes or is abandoned.

    Note: We don't use isolation_scope() here because it uses context variables that
    cannot span async boundaries (the __enter__ and __exit__ would be called from
    different async contexts, causing ValueError).
    """

    @wraps(original_func)
    def wrapper(*args: "Any", **kwargs: "Any") -> "Any":
        # Clone agent because agent invocation spans are attached per run.
        if "starting_agent" in kwargs:
            agent = kwargs["starting_agent"].clone()
        else:
            agent = args[0].clone()

        # Capture conversation_id from kwargs if provided
        conversation_id = kwargs.get("conversation_id")
        if conversation_id:
            agent._sentry_conversation_id = conversation_id

        # Start workflow span immediately (before run_streamed returns)
        workflow_span = agent_workflow_span(agent)
        workflow_span.__enter__()

        # Set conversation ID on workflow span early so it's captured even on errors
        if conversation_id:
            workflow_span.set_data(SPANDATA.GEN_AI_CONVERSATION_ID, conversation_id)

        # Store span on agent for cleanup
        agent._sentry_workflow_span = workflow_span

        if "starting_agent" in kwargs:
            kwargs["starting_agent"] = agent
        else:
            args = (agent, *args[1:])

        try:
            # Call original function to get RunResultStreaming
            run_result = original_func(*args, **kwargs)
        except Exception as exc:
            # If run_streamed itself fails (not the background task), clean up immediately
            workflow_span.__exit__(*sys.exc_info())
            _capture_exception(exc)
            raise

        def _close_workflow_span() -> None:
            if hasattr(agent, "_sentry_workflow_span"):
                workflow_span.__exit__(*sys.exc_info())
                delattr(agent, "_sentry_workflow_span")

        if hasattr(run_result, "stream_events"):
            original_stream_events = run_result.stream_events

            @wraps(original_stream_events)
            async def wrapped_stream_events(
                *stream_args: "Any", **stream_kwargs: "Any"
            ) -> "AsyncIterator[Any]":
                try:
                    async for event in original_stream_events(
                        *stream_args, **stream_kwargs
                    ):
                        yield event
                finally:
                    _close_workflow_span()

            run_result.stream_events = wrapped_stream_events

        if hasattr(run_result, "cancel"):
            original_cancel = run_result.cancel

            @wraps(original_cancel)
            def wrapped_cancel(*cancel_args: "Any", **cancel_kwargs: "Any") -> "Any":
                try:
                    return original_cancel(*cancel_args, **cancel_kwargs)
                finally:
                    _close_workflow_span()

            run_result.cancel = wrapped_cancel

        return run_result

    return wrapper