"""
OpenAI Agents SDK patches for Runner execution tracing.

Wraps Runner.run() and Runner.run_sync() to create AGENT spans around
the entire agent execution. Does NOT suppress provider instrumentation —
lets the OpenAI provider patch handle LLM calls naturally.
"""

from __future__ import annotations

import logging
import time
from typing import Any, Callable, Dict, TypeVar

import wrapt

from risicare.integrations._base import (
    get_tracer,
    is_tracing_enabled,
    record_error,
    safe_set_attribute,
    should_trace_content,
    truncate_content,
    create_framework_attributes,
    get_framework_version,
)

logger = logging.getLogger(__name__)

T = TypeVar("T")


def _extract_agent_name(agent: Any) -> str:
    """Extract agent name from an OpenAI Agents SDK Agent instance."""
    return (
        getattr(agent, "name", None)
        or getattr(agent, "model", None)
        or "agent"
    )


def _extract_agent_attrs(agent: Any) -> Dict[str, Any]:
    """Extract agent attributes for span."""
    attrs: Dict[str, Any] = {}
    name = _extract_agent_name(agent)
    attrs["agent.name"] = name

    # Extract tools
    tools = getattr(agent, "tools", [])
    if tools:
        tool_names = []
        for tool in tools[:10]:
            tool_name = getattr(tool, "name", None) or getattr(tool, "__name__", str(tool))
            tool_names.append(str(tool_name))
        attrs["framework.openai_agents.tools"] = tool_names

    # Extract handoffs
    handoffs = getattr(agent, "handoffs", [])
    if handoffs:
        handoff_names = []
        for h in handoffs[:10]:
            h_name = _extract_agent_name(h) if hasattr(h, "name") else str(h)
            handoff_names.append(h_name)
        attrs["framework.openai_agents.handoffs"] = handoff_names

    # Model
    model = getattr(agent, "model", None)
    if model:
        attrs["framework.openai_agents.model"] = model

    return attrs


async def _wrap_runner_run(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for Runner.run (async)."""
    if not is_tracing_enabled():
        return await wrapped(*args, **kwargs)

    tracer = get_tracer()
    if tracer is None:
        return await wrapped(*args, **kwargs)

    from risicare_core import SpanKind

    # Extract agent from first arg or kwargs
    agent = args[0] if args else kwargs.get("starting_agent") or kwargs.get("agent")
    agent_name = _extract_agent_name(agent) if agent else "agent"
    version = get_framework_version("openai-agents")

    attrs = create_framework_attributes(
        "openai_agents", version, agent_name=agent_name
    )
    attrs["agent.role"] = "orchestrator"
    attrs["agent.type"] = "openai_agents"

    if agent:
        attrs.update(_extract_agent_attrs(agent))

    # Capture input
    input_val = args[1] if len(args) > 1 else kwargs.get("input")
    if input_val and should_trace_content():
        if isinstance(input_val, str):
            attrs["framework.openai_agents.input"] = truncate_content(input_val, 5000)

    with tracer.start_span(
        name=f"openai_agents.run/{agent_name}",
        kind=SpanKind.AGENT,
        attributes=attrs,
    ) as span:
        try:
            start_time = time.perf_counter()
            result = await wrapped(*args, **kwargs)
            latency_ms = (time.perf_counter() - start_time) * 1000
            span.set_attribute("gen_ai.latency_ms", latency_ms)

            # Capture result metadata
            final_output = getattr(result, "final_output", None)
            if final_output and should_trace_content():
                if isinstance(final_output, str):
                    safe_set_attribute(
                        span,
                        "framework.openai_agents.output",
                        truncate_content(final_output, 5000),
                    )

            # Capture handoff info
            last_agent = getattr(result, "last_agent", None)
            if last_agent:
                last_name = _extract_agent_name(last_agent)
                if last_name != agent_name:
                    span.set_attribute(
                        "framework.openai_agents.final_agent", last_name
                    )

            # Capture step count
            new_items = getattr(result, "new_items", [])
            if new_items:
                span.set_attribute(
                    "framework.openai_agents.step_count", len(new_items)
                )

            return result
        except Exception as e:
            record_error(span, e)
            raise


def _wrap_runner_run_sync(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for Runner.run_sync (sync)."""
    if not is_tracing_enabled():
        return wrapped(*args, **kwargs)

    tracer = get_tracer()
    if tracer is None:
        return wrapped(*args, **kwargs)

    from risicare_core import SpanKind

    agent = args[0] if args else kwargs.get("starting_agent") or kwargs.get("agent")
    agent_name = _extract_agent_name(agent) if agent else "agent"
    version = get_framework_version("openai-agents")

    attrs = create_framework_attributes(
        "openai_agents", version, agent_name=agent_name
    )
    attrs["agent.role"] = "orchestrator"
    attrs["agent.type"] = "openai_agents"

    if agent:
        attrs.update(_extract_agent_attrs(agent))

    input_val = args[1] if len(args) > 1 else kwargs.get("input")
    if input_val and should_trace_content():
        if isinstance(input_val, str):
            attrs["framework.openai_agents.input"] = truncate_content(input_val, 5000)

    with tracer.start_span(
        name=f"openai_agents.run/{agent_name}",
        kind=SpanKind.AGENT,
        attributes=attrs,
    ) as span:
        try:
            start_time = time.perf_counter()
            result = wrapped(*args, **kwargs)
            latency_ms = (time.perf_counter() - start_time) * 1000
            span.set_attribute("gen_ai.latency_ms", latency_ms)

            final_output = getattr(result, "final_output", None)
            if final_output and should_trace_content():
                if isinstance(final_output, str):
                    safe_set_attribute(
                        span,
                        "framework.openai_agents.output",
                        truncate_content(final_output, 5000),
                    )

            last_agent = getattr(result, "last_agent", None)
            if last_agent:
                last_name = _extract_agent_name(last_agent)
                if last_name != agent_name:
                    span.set_attribute(
                        "framework.openai_agents.final_agent", last_name
                    )

            new_items = getattr(result, "new_items", [])
            if new_items:
                span.set_attribute(
                    "framework.openai_agents.step_count", len(new_items)
                )

            return result
        except Exception as e:
            record_error(span, e)
            raise


def patch_openai_agents(module: Any) -> None:
    """Apply wrapt patches to OpenAI Agents SDK Runner."""
    try:
        wrapt.wrap_function_wrapper(
            "agents",
            "Runner.run",
            _wrap_runner_run,
        )
    except Exception as e:
        logger.debug(f"Failed to patch Runner.run: {e}")

    try:
        wrapt.wrap_function_wrapper(
            "agents",
            "Runner.run_sync",
            _wrap_runner_run_sync,
        )
    except Exception as e:
        logger.debug(f"Failed to patch Runner.run_sync: {e}")
