"""
OpenAI Agents SDK integration for Amplitude AI SDK.

Provides :class:`AmplitudeTracingProcessor` — a tracing processor that
plugs into the OpenAI Agents SDK's tracing system and automatically
tracks LLM generations, tool calls, agent spans, handoffs, and
guardrails in Amplitude.

Usage::

    from agents import Agent, Runner, set_tracing_export_api_key
    from amplitude import Amplitude
    from amplitude_ai.integrations.openai_agents import AmplitudeTracingProcessor

    amplitude = Amplitude("YOUR_API_KEY")
    processor = AmplitudeTracingProcessor(
        amplitude=amplitude,
        user_id="user-1",
    )

    agent = Agent(name="my-agent", instructions="...")
    result = Runner.run_sync(agent, "Hello!", run_config=RunConfig(
        tracing_processors=[processor],
    ))
"""

from __future__ import annotations

import time
import uuid
from typing import Any, Dict, Optional

try:
    from agents.tracing import TracingProcessor, Trace, Span
    from agents.tracing.spans import (
        GenerationSpanData,
        FunctionSpanData,
        AgentSpanData,
        HandoffSpanData,
        GuardrailSpanData,
    )

    OPENAI_AGENTS_AVAILABLE = True
except ImportError:
    OPENAI_AGENTS_AVAILABLE = False
    TracingProcessor = object  # type: ignore[assignment,misc,unused-ignore]

from amplitude import Amplitude

from ..context import get_active_context
from ..core.privacy import PrivacyConfig
from ..core.tracking import (
    track_ai_message,
    track_span,
    track_tool_call,
    track_user_message,
)
from ..utils.costs import calculate_cost, infer_provider
from ..utils.logger import get_logger
from ..utils.providers import infer_provider_from_model


class AmplitudeTracingProcessor(TracingProcessor):  # type: ignore[misc]
    """OpenAI Agents SDK tracing processor that tracks events in Amplitude.

    Maps OpenAI Agents SDK spans to Amplitude Agent events:

    - ``GenerationSpanData`` -> ``[Agent] AI Response`` (+ ``[Agent] User Message``)
    - ``FunctionSpanData`` -> ``[Agent] Tool Call``
    - ``AgentSpanData`` -> ``[Agent] Span`` (agent lifecycle)
    - ``HandoffSpanData`` -> ``[Agent] Span`` (handoff)
    - ``GuardrailSpanData`` -> ``[Agent] Span`` (guardrail)
    """

    def __init__(
        self,
        amplitude: Amplitude,
        user_id: str,
        privacy_config: Optional[PrivacyConfig] = None,
        session_id: Optional[str] = None,
        trace_id: Optional[str] = None,
        agent_id: Optional[str] = None,
        parent_agent_id: Optional[str] = None,
        customer_org_id: Optional[str] = None,
        agent_version: Optional[str] = None,
        context: Optional[Dict[str, Any]] = None,
        env: Optional[str] = None,
        groups: Optional[Dict[str, Any]] = None,
    ) -> None:
        ctx = get_active_context()

        self.amplitude = amplitude
        self.user_id = user_id or (ctx.user_id if ctx else None) or ""
        self.privacy_config = privacy_config or PrivacyConfig()
        self.session_id = session_id or (ctx.session_id if ctx else None) or str(uuid.uuid4())
        self.trace_id = trace_id or (ctx.trace_id if ctx else None)
        self.agent_id = agent_id or (ctx.agent_id if ctx else None)
        self.parent_agent_id = parent_agent_id or (ctx.parent_agent_id if ctx else None)
        self.customer_org_id = customer_org_id or (ctx.customer_org_id if ctx else None)
        self.agent_version = agent_version or (ctx.agent_version if ctx else None)
        self.context = context or (ctx.context if ctx else None)
        self.env = env or (ctx.env if ctx else None)
        self.groups = groups or (ctx.groups if ctx else None)

        self._turn_counter = 1
        self._trace_ids: Dict[str, str] = {}
        self._span_start_times: Dict[str, float] = {}

    # ------------------------------------------------------------------
    # TracingProcessor interface
    # ------------------------------------------------------------------

    def on_trace_start(self, trace: Any) -> None:
        """Record the trace start and map to an Amplitude trace_id."""
        trace_id_str = getattr(trace, "trace_id", None) or str(uuid.uuid4())
        self._trace_ids[trace_id_str] = self.trace_id or str(uuid.uuid4())

    def on_trace_end(self, trace: Any) -> None:
        """Clean up trace-level state."""
        trace_id_str = getattr(trace, "trace_id", None)
        if trace_id_str:
            self._trace_ids.pop(trace_id_str, None)

    def on_span_start(self, span: Any) -> None:
        """Record span start time."""
        span_id = getattr(span, "span_id", None) or str(uuid.uuid4())
        self._span_start_times[span_id] = time.time()

    def on_span_end(self, span: Any) -> None:
        """Process the completed span and emit Amplitude events."""
        if not OPENAI_AGENTS_AVAILABLE:
            return

        _logger = get_logger(self.amplitude)

        try:
            span_id = getattr(span, "span_id", None) or ""
            start_time = self._span_start_times.pop(span_id, time.time())
            latency_ms = (time.time() - start_time) * 1000

            trace_id_str = getattr(span, "trace_id", None) or ""
            amp_trace_id = self._trace_ids.get(trace_id_str, self.trace_id or str(uuid.uuid4()))

            span_data = getattr(span, "span_data", None)
            if span_data is None:
                return

            if isinstance(span_data, GenerationSpanData):
                self._handle_generation(span_data, latency_ms, amp_trace_id)
            elif isinstance(span_data, FunctionSpanData):
                self._handle_function(span_data, latency_ms, amp_trace_id)
            elif isinstance(span_data, AgentSpanData):
                self._handle_agent_span(span_data, latency_ms, amp_trace_id)
            elif isinstance(span_data, HandoffSpanData):
                self._handle_handoff(span_data, latency_ms, amp_trace_id)
            elif isinstance(span_data, GuardrailSpanData):
                self._handle_guardrail(span_data, latency_ms, amp_trace_id)

        except Exception as exc:
            _logger.warning("OpenAI Agents tracking failed: %s", exc)

    # ------------------------------------------------------------------
    # Internal span handlers
    # ------------------------------------------------------------------

    def _handle_generation(
        self,
        data: Any,
        latency_ms: float,
        trace_id: str,
    ) -> None:
        """Map GenerationSpanData to [Agent] User Message + [Agent] AI Response."""
        model_name = getattr(data, "model", None) or "unknown"
        provider = infer_provider_from_model(model_name)

        input_messages = getattr(data, "input", None) or []
        for msg in input_messages:
            role = msg.get("role", "") if isinstance(msg, dict) else getattr(msg, "role", "")
            content = msg.get("content", "") if isinstance(msg, dict) else getattr(msg, "content", "")
            if role == "user" and content:
                track_user_message(
                    amplitude=self.amplitude,
                    user_id=self.user_id,
                    message_content=str(content),
                    session_id=self.session_id,
                    trace_id=trace_id,
                    turn_id=self._turn_counter,
                    agent_id=self.agent_id,
                    parent_agent_id=self.parent_agent_id,
                    customer_org_id=self.customer_org_id,
                    agent_version=self.agent_version,
                    context=self.context,
                    env=self.env,
                    groups=self.groups,
                    privacy_config=self.privacy_config,
                )
                self._turn_counter += 1

        output = getattr(data, "output", None) or []
        response_text = ""
        tool_calls_list: list[Dict[str, Any]] = []
        for block in output:
            if isinstance(block, dict):
                role = block.get("role", "")
                if role == "assistant":
                    content = block.get("content", "")
                    if isinstance(content, str):
                        response_text += content
                    elif isinstance(content, list):
                        for part in content:
                            if isinstance(part, dict) and part.get("type") == "text":
                                response_text += part.get("text", "")
                    for tc in block.get("tool_calls", []):
                        tool_calls_list.append({
                            "type": "function",
                            "id": tc.get("id", ""),
                            "function": {
                                "name": tc.get("function", {}).get("name", ""),
                                "arguments": tc.get("function", {}).get("arguments", ""),
                            },
                        })
            else:
                role = getattr(block, "role", "")
                if role == "assistant":
                    content = getattr(block, "content", "")
                    if isinstance(content, str):
                        response_text += content

        usage = getattr(data, "usage", None) or {}
        input_tokens = _get_usage_field(usage, "input_tokens", "prompt_tokens")
        output_tokens = _get_usage_field(usage, "output_tokens", "completion_tokens")
        total_tokens = _get_usage_field(usage, "total_tokens")
        prompt_details = getattr(usage, "prompt_tokens_details", None)
        if prompt_details is None and isinstance(usage, dict):
            prompt_details = usage.get("prompt_tokens_details")
        cached_tokens = _get_usage_field(prompt_details or {}, "cached_tokens") or 0

        total_cost_usd = None
        if input_tokens is not None and output_tokens is not None:
            try:
                total_cost_usd = calculate_cost(
                    model_name=model_name,
                    input_tokens=input_tokens,
                    output_tokens=output_tokens,
                    cache_read_input_tokens=cached_tokens,
                    default_provider=infer_provider(model_name),
                )
            except Exception:
                pass

        track_ai_message(
            amplitude=self.amplitude,
            user_id=self.user_id,
            model_name=model_name,
            provider=provider,
            response_content=response_text,
            latency_ms=latency_ms,
            session_id=self.session_id,
            trace_id=trace_id,
            turn_id=self._turn_counter,
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            total_tokens=total_tokens,
            total_cost_usd=total_cost_usd,
            tool_calls=tool_calls_list if tool_calls_list else None,
            agent_id=self.agent_id,
            parent_agent_id=self.parent_agent_id,
            customer_org_id=self.customer_org_id,
            agent_version=self.agent_version,
            context=self.context,
            env=self.env,
            groups=self.groups,
            privacy_config=self.privacy_config,
        )
        self._turn_counter += 1

    def _handle_function(
        self,
        data: Any,
        latency_ms: float,
        trace_id: str,
    ) -> None:
        """Map FunctionSpanData to [Agent] Tool Call."""
        tool_name = getattr(data, "name", None) or "unknown"
        tool_input_raw = getattr(data, "input", None)
        tool_output_raw = getattr(data, "output", None)

        tool_input = None
        if tool_input_raw is not None:
            if isinstance(tool_input_raw, dict):
                tool_input = tool_input_raw
            else:
                tool_input = {"raw": str(tool_input_raw)}

        tool_output = None
        if tool_output_raw is not None:
            tool_output = str(tool_output_raw)

        error = getattr(data, "error", None)
        success = error is None

        track_tool_call(
            amplitude=self.amplitude,
            user_id=self.user_id,
            tool_name=tool_name,
            success=success,
            latency_ms=latency_ms,
            session_id=self.session_id,
            trace_id=trace_id,
            turn_id=self._turn_counter,
            tool_input=tool_input,
            tool_output=tool_output,
            error_message=str(error) if error else None,
            agent_id=self.agent_id,
            parent_agent_id=self.parent_agent_id,
            customer_org_id=self.customer_org_id,
            agent_version=self.agent_version,
            context=self.context,
            env=self.env,
            groups=self.groups,
            privacy_config=self.privacy_config,
        )

    def _handle_agent_span(
        self,
        data: Any,
        latency_ms: float,
        trace_id: str,
    ) -> None:
        """Map AgentSpanData to [Agent] Span (agent lifecycle)."""
        agent_name = getattr(data, "name", None) or "agent"
        output = getattr(data, "output", None)

        output_state = None
        if output is not None:
            output_state = {"output": str(output)}

        track_span(
            amplitude=self.amplitude,
            user_id=self.user_id,
            span_name=f"agent:{agent_name}",
            trace_id=trace_id,
            latency_ms=latency_ms,
            output_state=output_state,
            session_id=self.session_id,
            agent_id=self.agent_id or agent_name,
            parent_agent_id=self.parent_agent_id,
            customer_org_id=self.customer_org_id,
            agent_version=self.agent_version,
            context=self.context,
            env=self.env,
            groups=self.groups,
            privacy_config=self.privacy_config,
        )

    def _handle_handoff(
        self,
        data: Any,
        latency_ms: float,
        trace_id: str,
    ) -> None:
        """Map HandoffSpanData to [Agent] Span (handoff)."""
        from_agent = getattr(data, "from_agent", None) or "unknown"
        to_agent = getattr(data, "to_agent", None) or "unknown"

        track_span(
            amplitude=self.amplitude,
            user_id=self.user_id,
            span_name=f"handoff:{from_agent}->{to_agent}",
            trace_id=trace_id,
            latency_ms=latency_ms,
            input_state={"from_agent": str(from_agent)},
            output_state={"to_agent": str(to_agent)},
            session_id=self.session_id,
            agent_id=self.agent_id,
            parent_agent_id=self.parent_agent_id,
            customer_org_id=self.customer_org_id,
            agent_version=self.agent_version,
            context=self.context,
            env=self.env,
            groups=self.groups,
            privacy_config=self.privacy_config,
        )

    def _handle_guardrail(
        self,
        data: Any,
        latency_ms: float,
        trace_id: str,
    ) -> None:
        """Map GuardrailSpanData to [Agent] Span (guardrail)."""
        guardrail_name = getattr(data, "name", None) or "guardrail"
        triggered = getattr(data, "triggered", False)

        track_span(
            amplitude=self.amplitude,
            user_id=self.user_id,
            span_name=f"guardrail:{guardrail_name}",
            trace_id=trace_id,
            latency_ms=latency_ms,
            output_state={"triggered": triggered},
            is_error=bool(triggered),
            session_id=self.session_id,
            agent_id=self.agent_id,
            parent_agent_id=self.parent_agent_id,
            customer_org_id=self.customer_org_id,
            agent_version=self.agent_version,
            context=self.context,
            env=self.env,
            groups=self.groups,
            privacy_config=self.privacy_config,
        )


# ------------------------------------------------------------------
# Helpers
# ------------------------------------------------------------------



def _get_usage_field(usage: Any, *field_names: str) -> Optional[int]:
    """Extract a token-usage field by trying multiple attribute/key names."""
    for name in field_names:
        if isinstance(usage, dict):
            val = usage.get(name)
        else:
            val = getattr(usage, name, None)
        if val is not None:
            return int(val)
    return None
