import logging
import weakref
from datetime import datetime
from typing import Optional

from langsmith import run_trees as rt
from langsmith._internal import _context
from langsmith._internal._ls_agent_type import (
    NON_ROOT_LS_AGENT_TYPES,
    apply_default_ls_agent_type,
)
from langsmith._internal._package_version import get_package_version
from langsmith.run_helpers import get_current_run_tree

try:
    from agents import tracing  # type: ignore[import]

    required = (
        "TracingProcessor",
        "Trace",
        "Span",
        "ResponseSpanData",
    )
    if not all(hasattr(tracing, name) for name in required):
        raise ImportError("The `agents` package is not installed.")

    from langsmith.integrations.openai_agents_sdk import (
        _openai_agent_utils as agent_utils,
    )

    HAVE_AGENTS = True
except ImportError:
    HAVE_AGENTS = False

    class OpenAIAgentsTracingProcessor:
        """Tracing processor for the [OpenAI Agents SDK](https://openai.github.io/openai-agents-python/).

        Traces all intermediate steps of your OpenAI Agent to LangSmith.

        Requirements: Make sure to install `pip install -U langsmith[openai-agents]`.

        Args:
            client: An instance of `langsmith.client.Client`. If not provided, a default
                client is created.

        Example:
            ```python
            from agents import (
                Agent,
                FileSearchTool,
                Runner,
                WebSearchTool,
                function_tool,
                set_trace_processors,
            )

            from langsmith.wrappers import OpenAIAgentsTracingProcessor

            set_trace_processors([OpenAIAgentsTracingProcessor()])


            @function_tool
            def get_weather(city: str) -> str:
                return f"The weather in {city} is sunny"


            haiku_agent = Agent(
                name="Haiku agent",
                instructions="Always respond in haiku form",
                model="o3-mini",
                tools=[get_weather],
            )
            agent = Agent(
                name="Assistant",
                tools=[WebSearchTool()],
                instructions="speak in spanish. use Haiku agent if they ask for a haiku or for the weather",
                handoffs=[haiku_agent],
            )

            result = await Runner.run(
                agent,
                "write a haiku about the weather today and tell me a recent news story about new york",
            )
            print(result.final_output)
            ```
        """  # noqa: E501

        def __init__(self, *args, **kwargs):
            raise ImportError(
                "The `agents` package is not installed. "
                "Please install it with `pip install langsmith[openai-agents]`."
            )


from langsmith import client as ls_client

logger = logging.getLogger(__name__)


if HAVE_AGENTS:

    class OpenAIAgentsTracingProcessor(tracing.TracingProcessor):  # type: ignore[no-redef]
        """Tracing processor for the [OpenAI Agents SDK](https://openai.github.io/openai-agents-python/).

        Traces all intermediate steps of your OpenAI Agent to LangSmith.

        Requirements: Make sure to install `pip install -U langsmith[openai-agents]`.

        Args:
            client: An instance of `langsmith.client.Client`. If not provided,
                a default client is created.
            metadata: Metadata to associate with all traces.
            tags: Tags to associate with all traces.
            project_name: LangSmith project to trace to.
            name: Name of the root trace.

        Example:
            ```python
            from agents import (
                Agent,
                FileSearchTool,
                Runner,
                WebSearchTool,
                function_tool,
                set_trace_processors,
            )

            from langsmith.wrappers import OpenAIAgentsTracingProcessor

            set_trace_processors([OpenAIAgentsTracingProcessor()])


            @function_tool
            def get_weather(city: str) -> str:
                return f"The weather in {city} is sunny"


            haiku_agent = Agent(
                name="Haiku agent",
                instructions="Always respond in haiku form",
                model="o3-mini",
                tools=[get_weather],
            )
            agent = Agent(
                name="Assistant",
                tools=[WebSearchTool()],
                instructions="speak in spanish. use Haiku agent if they ask for a haiku or for the weather",
                handoffs=[haiku_agent],
            )

            result = await Runner.run(
                agent,
                "write a haiku about the weather today and tell me a recent news story about new york",
            )
            print(result.final_output)
            ```
        """  # noqa: E501

        def __init__(
            self,
            client: Optional[ls_client.Client] = None,
            *,
            metadata: Optional[dict] = None,
            tags: Optional[list[str]] = None,
            project_name: Optional[str] = None,
            name: Optional[str] = None,
        ):
            self.client = client or rt.get_cached_client()
            self._metadata = metadata
            self._tags = tags
            self._project_name = project_name
            self._name = name
            self._first_response_inputs: dict = {}
            self._last_response_outputs: dict = {}

            self._runs: dict[str, rt.RunTree] = {}
            self._span_data_types: dict[
                str, type
            ] = {}  # Track span data types by span_id
            self._unposted_traces: set[str] = set()
            self._unposted_spans: set[str] = set()

        def on_trace_start(self, trace: tracing.Trace) -> None:
            current_run_tree = get_current_run_tree()

            # Determine run name
            if self._name:
                run_name = self._name
            elif trace.name:
                run_name = trace.name
            else:
                run_name = "Agent workflow"

            # Merge three sources: processor defaults, trace.metadata (user's
            # per-run ls_agent_type), and force-set ls_integration.
            merged_metadata = {
                **(self._metadata or {}),
                **(getattr(trace, "metadata", None) or {}),
                "ls_integration": "openai-agents-sdk",
                "ls_integration_version": get_package_version("openai-agents"),
            }
            apply_default_ls_agent_type(merged_metadata, current_run_tree)
            run_extra = {"metadata": merged_metadata}
            trace_dict = trace.export() or {}
            if trace_dict.get("group_id") is not None:
                run_extra["metadata"]["thread_id"] = trace_dict["group_id"]

            try:
                if current_run_tree is not None:
                    # Nest under existing trace
                    new_run = current_run_tree.create_child(
                        name=run_name,
                        run_type="chain",
                        inputs={},
                        extra=run_extra,
                        tags=self._tags,
                    )
                else:
                    # Create new root trace
                    run_kwargs = {
                        "name": run_name,
                        "run_type": "chain",
                        "inputs": {},
                        "extra": run_extra,
                        "tags": self._tags,
                        "client": self.client,
                    }
                    if self._project_name is not None:
                        run_kwargs["project_name"] = self._project_name
                    new_run = rt.RunTree(**run_kwargs)  # type: ignore[arg-type]

                # Delay posting until first response/generation span ends
                # so inputs can be included in the POST.
                self._unposted_traces.add(trace.trace_id)
                if new_run is not None:
                    _context._PARENT_RUN_TREE_REF.set(weakref.ref(new_run))
                self._runs[trace.trace_id] = new_run
            except Exception as e:
                logger.exception(f"Error creating trace run: {e}")

        def on_trace_end(self, trace: tracing.Trace) -> None:
            run = self._runs.pop(trace.trace_id, None)
            if not run:
                return

            trace_dict = trace.export() or {}
            # trace.metadata may have new keys since trace-start; pull them in.
            metadata = {
                **(self._metadata or {}),
                **(trace_dict.get("metadata") or {}),
            }

            try:
                # Update run with final inputs/outputs
                run.outputs = self._last_response_outputs.pop(trace.trace_id, {})

                # Update metadata
                if "metadata" not in run.extra:
                    run.extra["metadata"] = {}
                run.extra["metadata"].update(metadata)

                # End and patch
                run.end()

                if trace.trace_id in self._unposted_traces:
                    # No response/generation spans ended, post now
                    run.inputs = self._first_response_inputs.pop(trace.trace_id, {})
                    self._unposted_traces.discard(trace.trace_id)
                    run.post()
                else:
                    self._first_response_inputs.pop(trace.trace_id, None)
                    run.patch(exclude_inputs=True)

                # Restore parent context
                if run.parent_run is not None:
                    _context._PARENT_RUN_TREE_REF.set(weakref.ref(run.parent_run))
                else:
                    _context._PARENT_RUN_TREE_REF.set(None)
            except Exception as e:
                logger.exception(f"Error updating trace run: {e}")

        def on_span_start(self, span: tracing.Span) -> None:
            # Find parent run
            parent_run = (
                self._runs.get(span.parent_id)
                if span.parent_id
                else self._runs.get(span.trace_id)
            )

            if parent_run is None:
                logger.warning(
                    f"No trace info found for span, skipping: {span.span_id}"
                )
                return

            # Extract span data
            run_name = agent_utils.get_run_name(span)
            if isinstance(span.span_data, tracing.ResponseSpanData):
                parent_name = parent_run.name
                raw_span_name = getattr(span, "name", None) or getattr(
                    span.span_data, "name", None
                )
                span_name = str(raw_span_name) if raw_span_name else run_name
                if parent_name:
                    run_name = f"{parent_name} {span_name}".strip()
                else:
                    run_name = span_name

            run_type = agent_utils.get_run_type(span)
            extracted = agent_utils.extract_span_data(span)

            try:
                # Create child run
                child_run = parent_run.create_child(
                    name=run_name,
                    run_type=run_type,
                    inputs=extracted.get("inputs", {}),
                    extra=extracted,
                    start_time=datetime.fromisoformat(span.started_at)
                    if span.started_at
                    else None,
                )

                # Add ls_agent_type metadata for agent spans that are children of
                # function spans (i.e., agents used as tools via as_tool()).
                # Note: Handoff agents are considered root agents, not subagents,
                # since they take over the conversation rather than being called
                # as tools.
                if isinstance(span.span_data, tracing.AgentSpanData):
                    # Check if parent span is a function span (agent used as tool)
                    parent_span_data_type = (
                        self._span_data_types.get(span.parent_id)
                        if span.parent_id
                        else None
                    )
                    if parent_span_data_type is tracing.FunctionSpanData:
                        metadata = child_run.extra.setdefault("metadata", {})
                        if metadata.get("ls_agent_type") not in NON_ROOT_LS_AGENT_TYPES:
                            metadata["ls_agent_type"] = "subagent"

                # Track span data type for parent lookups
                self._span_data_types[span.span_id] = type(span.span_data)

                # Delay posting for spans whose inputs aren't available at start
                if isinstance(
                    span.span_data,
                    (
                        tracing.GenerationSpanData,
                        tracing.ResponseSpanData,
                        tracing.FunctionSpanData,
                    ),
                ):
                    self._unposted_spans.add(span.span_id)
                else:
                    child_run.post()
                self._runs[span.span_id] = child_run
            except Exception as e:
                logger.exception(f"Error creating span run: {e}")

        def on_span_end(self, span: tracing.Span) -> None:
            run = self._runs.pop(span.span_id, None)
            self._span_data_types.pop(
                span.span_id, None
            )  # Clean up span data type tracking
            if not run:
                return

            try:
                # Extract outputs and metadata
                extracted = agent_utils.extract_span_data(span)
                outputs = extracted.pop("outputs", {})
                inputs = extracted.pop("inputs", {})

                # Update run
                run.outputs = outputs
                if inputs:
                    run.inputs = inputs
                if error := span.error:
                    run.error = str(error)

                # Add OpenAI metadata
                if "metadata" not in run.extra:
                    run.extra["metadata"] = {}
                run.extra["metadata"].update(
                    {
                        "openai_parent_id": span.parent_id,
                        "openai_trace_id": span.trace_id,
                        "openai_span_id": span.span_id,
                    }
                )
                if metadata := extracted.get("metadata"):
                    run.extra["metadata"].update(metadata)
                if invocation_params := extracted.get("invocation_params"):
                    run.extra["invocation_params"] = invocation_params

                if isinstance(span.span_data, tracing.ResponseSpanData):
                    self._first_response_inputs[span.trace_id] = (
                        self._first_response_inputs.get(span.trace_id) or inputs
                    )
                    self._last_response_outputs[span.trace_id] = outputs
                    self._maybe_post_trace(span.trace_id, inputs)
                elif isinstance(span.span_data, tracing.GenerationSpanData):
                    self._first_response_inputs[span.trace_id] = (
                        self._first_response_inputs.get(span.trace_id) or inputs
                    )
                    self._last_response_outputs[span.trace_id] = outputs
                    self._maybe_post_trace(span.trace_id, inputs)

                if span.ended_at:
                    run.end_time = datetime.fromisoformat(span.ended_at)
                else:
                    run.end()

                if span.span_id in self._unposted_spans:
                    self._unposted_spans.discard(span.span_id)
                    run.post()
                else:
                    run.patch(exclude_inputs=True)
            except Exception as e:
                logger.exception(f"Error updating span run: {e}")

        def _maybe_post_trace(self, trace_id: str, inputs: dict) -> None:
            """Post the trace if it hasn't been posted yet."""
            if trace_id in self._unposted_traces:
                trace_run = self._runs.get(trace_id)
                if trace_run:
                    trace_run.inputs = inputs
                    trace_run.post()
                    self._unposted_traces.discard(trace_id)

        def shutdown(self) -> None:
            self.client.flush()

        def force_flush(self) -> None:
            self.client.flush()
