from __future__ import annotations

import json
import logging
from typing import Any, Optional

import agents.tracing as oai
from agents import add_trace_processor
from agents._run_impl import TraceCtxManager
from agents.tracing.setup import GLOBAL_TRACE_PROVIDER
from pydantic import BaseModel

from mlflow import MlflowClient
from mlflow.entities.span import LiveSpan, SpanType
from mlflow.entities.span_event import SpanEvent
from mlflow.entities.span_status import SpanStatus, SpanStatusCode
from mlflow.openai import FLAVOR_NAME
from mlflow.tracing.constant import SpanAttributeKey
from mlflow.tracing.utils import end_client_span_or_trace, start_client_span_or_trace
from mlflow.types.chat import (
    ChatMessage,
    ChatTool,
    Function,
    FunctionToolDefinition,
    TextContentPart,
    ToolCall,
)
from mlflow.utils.autologging_utils.safety import safe_patch

_logger = logging.getLogger(__name__)


class OpenAISpanType:
    """
    https://github.com/openai/openai-agents-python/blob/main/src/agents/tracing/span_data.py#L11
    """

    AGENT = "agent"
    FUNCTION = "function"
    GENERATION = "generation"
    RESPONSE = "response"
    HANDOFF = "handoff"
    CUSTOM = "custom"
    GUARDRAIL = "guardrail"


_SPAN_TYPE_MAP = {
    OpenAISpanType.AGENT: SpanType.AGENT,
    OpenAISpanType.FUNCTION: SpanType.TOOL,
    OpenAISpanType.GENERATION: SpanType.CHAT_MODEL,
    OpenAISpanType.RESPONSE: SpanType.CHAT_MODEL,
    OpenAISpanType.GUARDRAIL: SpanType.TOOL,
    # Default to chain type
}


def add_mlflow_trace_processor():
    processors = GLOBAL_TRACE_PROVIDER._multi_processor._processors

    if any(isinstance(p, MlflowOpenAgentTracingProcessor) for p in processors):
        return

    add_trace_processor(MlflowOpenAgentTracingProcessor())


def remove_mlflow_trace_processor():
    processors = GLOBAL_TRACE_PROVIDER._multi_processor._processors
    non_mlflow_processors = [
        p for p in processors if not isinstance(p, MlflowOpenAgentTracingProcessor)
    ]
    GLOBAL_TRACE_PROVIDER._multi_processor._processors = non_mlflow_processors


class MlflowOpenAgentTracingProcessor(oai.TracingProcessor):
    def __init__(
        self,
        project_name: Optional[str] = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        self._span_id_to_mlflow_span: dict[str, LiveSpan] = {}
        self._project_name = project_name
        self._mlflow_client = MlflowClient()

        # Patch TraceCtxManager to handle exceptions from the agent properly
        # The original implementation does not propagate exception to the root span,
        # resulting in the trace to have status OK even if there is an exception.
        def _patched_exit(original, instance, exc_type, exc_val, exc_tb):
            try:
                if exc_val and instance.trace:
                    span = self._span_id_to_mlflow_span.get(instance.trace.trace_id)
                    span.add_event(SpanEvent.from_exception(exc_val))
                    span.set_status(SpanStatusCode.ERROR)
            except Exception:
                _logger.debug("Failed to handle exception in MLflow trace", exc_info=True)

            return original(instance, exc_type, exc_val, exc_tb)

        safe_patch(
            FLAVOR_NAME,
            TraceCtxManager,
            "__exit__",
            _patched_exit,
        )

    def on_trace_start(self, trace: oai.Trace) -> None:
        try:
            mlflow_span = start_client_span_or_trace(
                client=self._mlflow_client,
                name=trace.name,
                span_type=SpanType.AGENT,
                # TODO: Trace object doesn't contain input/output. Can we get it somehow?
                inputs="",
                attributes=trace.metadata,
            )
            # NB: Trace ID has different prefix as span ID so will not conflict
            self._span_id_to_mlflow_span[trace.trace_id] = mlflow_span

            if trace.group_id:
                # Group ID is used for grouping multiple agent executions together
                mlflow_span.set_tag("group_id", trace.group_id)

            original_exit = trace.__exit__

            # Patch __exit__ method to handle exception properly
            def _patched_exit(self, exc_type, exc_val, exc_tb):
                if exc_val:
                    mlflow_span.add_event(SpanEvent.from_exception(exc_val))
                    mlflow_span.set_status(SpanStatusCode.ERROR)

                original_exit(exc_type, exc_val, exc_tb)

            safe_patch(
                FLAVOR_NAME,
                trace.__class__,
                "__exit__",
                _patched_exit,
            )

        except Exception:
            _logger.debug("Failed to start MLflow trace", exc_info=True)

    def on_trace_end(self, trace: oai.Trace) -> None:
        try:
            mlflow_span = self._span_id_to_mlflow_span.pop(trace.trace_id, None)
            end_client_span_or_trace(
                client=self._mlflow_client,
                span=mlflow_span,
                status=mlflow_span.status,
                outputs="",
            )
        except Exception:
            _logger.debug("Failed to end MLflow trace", exc_info=True)

    def on_span_start(self, span: oai.Span[Any]) -> None:
        try:
            parent_mlflow_span = self._span_id_to_mlflow_span.get(span.parent_id)

            # Parent might be a trace
            if not parent_mlflow_span:
                parent_mlflow_span = self._span_id_to_mlflow_span.get(span.trace_id)

            inputs, _, attributes = _parse_span_data(span.span_data)

            mlflow_span = start_client_span_or_trace(
                client=self._mlflow_client,
                name=_get_span_name(span.span_data),
                span_type=_SPAN_TYPE_MAP.get(span.span_data.type, SpanType.CHAIN),
                parent_span=parent_mlflow_span,
                inputs=inputs,
                attributes=attributes,
            )
            self._span_id_to_mlflow_span[span.span_id] = mlflow_span
        except Exception:
            _logger.debug("Failed to start MLflow span", exc_info=True)

    def on_span_end(self, span: oai.Span[Any]) -> None:
        try:
            # parsed_span_data = parse_spandata(span.span_data)
            mlflow_span = self._span_id_to_mlflow_span.pop(span.span_id, None)

            inputs, outputs, attributes = _parse_span_data(span.span_data)

            mlflow_span.set_inputs(inputs)
            mlflow_span.set_outputs(outputs)
            mlflow_span.set_attributes(attributes)

            if span.error:
                status = SpanStatus(
                    status_code=SpanStatusCode.ERROR,
                    description=span.error["message"],
                )
                mlflow_span.add_event(
                    SpanEvent(
                        name="exception",
                        attributes={
                            "exception.message": span.error["message"],
                            "exception.type": "",
                            "exception.stacktrace": json.dumps(span.error["data"]),
                        },
                    )
                )
            else:
                status = SpanStatusCode.OK

            end_client_span_or_trace(
                client=self._mlflow_client,
                span=mlflow_span,
                status=status,
            )
        except Exception:
            _logger.debug("Failed to end MLflow span", exc_info=True)

    def force_flush(self) -> None:
        # MLflow doesn't need flush but this method is required by the interface
        pass

    def shutdown(self) -> None:
        self.force_flush()


def _get_span_name(span_data: oai.SpanData) -> str:
    if hasattr(span_data, "name"):
        return span_data.name
    elif isinstance(span_data, oai.GenerationSpanData):
        return "Generation"
    elif isinstance(span_data, oai.ResponseSpanData):
        return "Response"
    elif isinstance(span_data, oai.HandoffSpanData):
        return "Handoff"
    else:
        return "Unknown"


def _parse_span_data(span_data: oai.SpanData) -> tuple[Any, Any, dict[str, Any]]:
    inputs = None
    outputs = None
    attributes = {}

    if span_data.type == OpenAISpanType.AGENT:
        attributes = {
            "handoffs": span_data.handoffs,
            "tools": span_data.tools,
            "output_type": span_data.output_type,
        }
        outputs = {"output_type": span_data.output_type}

    elif span_data.type == OpenAISpanType.FUNCTION:
        try:
            inputs = json.loads(span_data.input)
        except Exception:
            inputs = span_data.input
        outputs = span_data.output

    elif span_data.type == OpenAISpanType.GENERATION:
        inputs = span_data.input
        outputs = span_data.output
        attributes = {
            "model": span_data.model,
            "model_config": span_data.model_config,
            "usage": span_data.usage,
        }

    elif span_data.type == OpenAISpanType.RESPONSE:
        inputs, outputs, attributes = _parse_response_span_data(span_data)

    elif span_data.type == OpenAISpanType.HANDOFF:
        inputs = {"from_agent": span_data.from_agent}
        outputs = {"to_agent": span_data.to_agent}

    elif span_data.type == OpenAISpanType.CUSTOM:
        outputs = span_data.data

    elif span_data.type == OpenAISpanType.GUARDRAIL:
        outputs = {"triggered": span_data.triggered}

    return inputs, outputs, attributes


def _parse_response_span_data(span_data: oai.ResponseSpanData) -> tuple[Any, Any, dict[str, Any]]:
    inputs = span_data.input
    response = span_data.response
    response_dict = response.model_dump() if response else {}
    outputs = response_dict.get("output")
    attributes = {k: v for k, v in response_dict.items() if k != "output"}

    # Extract chat messages
    messages = []
    if response and response.instructions:
        messages.append(ChatMessage(role="system", content=span_data.response.instructions))
    if span_data.input:
        parsed = [_parse_message_like(m) for m in span_data.input]
        messages.extend([m for m in parsed if m is not None])
    if response and response.output:
        parsed = [_parse_message_like(m) for m in span_data.response.output]
        messages.extend(parsed)
    attributes[SpanAttributeKey.CHAT_MESSAGES] = [m.model_dump_compat() for m in messages]

    # Extract chat tools
    chat_tools = []
    for tool in response_dict.get("tools", []):
        try:
            tool = ChatTool(
                type="function",
                function=FunctionToolDefinition(
                    name=tool["name"],
                    description=tool.get("description"),
                    parameters=tool.get("parameters"),
                    strict=tool.get("strict"),
                ),
            )
            chat_tools.append(tool)
        except Exception as e:
            _logger.debug(f"Failed to parse chat tool: {tool}. Error: {e}")

    if chat_tools:
        attributes[SpanAttributeKey.CHAT_TOOLS] = chat_tools

    return inputs, outputs, attributes


def _parse_message_like(message_like: Any) -> Optional[ChatMessage]:
    try:
        return ChatMessage.validate_compat(message_like)
    except Exception:
        pass

    if isinstance(message_like, BaseModel):
        message_like = message_like.model_dump()

    msg_type = message_like["type"]
    if msg_type == "message":
        content = []
        refusal = None
        for content_block in message_like["content"]:
            # Content is a list of either text or refusal https://github.com/openai/openai-python/blob/9dea82fb8cdd06683f9e8033b54cff219789af7f/src/openai/types/responses/response_output_message.py#L13C38-L13C56
            if "text" in content_block:
                content.append(TextContentPart(type="text", text=content_block["text"]))
            elif "refusal" in content_block:
                refusal = content_block["refusal"]
            else:
                _logger.debug(f"Unknown content type in message: {content_block}")
        return ChatMessage(
            role=message_like["role"],
            content=content,
            refusal=refusal,
        )
    elif msg_type == "function_call":
        return ChatMessage(
            role="assistant",
            content="",
            tool_calls=[
                ToolCall(
                    id=message_like["call_id"],
                    function=Function(
                        name=message_like["name"],
                        arguments=message_like["arguments"],
                    ),
                )
            ],
        )
    elif msg_type == "function_call_output":
        return ChatMessage(
            role="tool",
            content=message_like["output"],
            tool_call_id=message_like["call_id"],
        )

    # Ignore unknown message types.
    # Response API supports the following additional message types, which is not
    # supported by our chat standard schema yet:
    # https://github.com/openai/openai-python/blob/9dea82fb8cdd06683f9e8033b54cff219789af7f/src/openai/types/responses/response_output_item.py#L16
    # - File search tool call
    # - Web search tool call
    # - Computer tool call
    # - Reasoning
    _logger.debug(f"Unknown message type: {msg_type}")