"""
OpenAI Agents SDK TracingProcessor for Raindrop observability.

Implements the TracingProcessor interface to capture agent runs, LLM
generations, tool calls, and handoffs, shipping them to Raindrop.

Usage:
    from raindrop_openai_agents import create_raindrop_openai_agents
    from agents import Agent, Runner

    raindrop = create_raindrop_openai_agents(api_key="rk_...")
    # Processor is auto-registered globally

    agent = Agent(name="Assistant", model="gpt-4o")
    result = Runner.run_sync(agent, "Hello!")
    raindrop["flush"]()
"""

from __future__ import annotations

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

import raindrop.analytics as raindrop


def create_raindrop_openai_agents(
    api_key: str,
    user_id: Optional[str] = None,
    convo_id: Optional[str] = None,
) -> dict:
    """
    Create a Raindrop-instrumented OpenAI Agents SDK processor.

    The processor is automatically registered with the global trace provider.

    Args:
        api_key: Raindrop API key (rk_...).
        user_id: Associate all events with this user.
        convo_id: Conversation ID to group related events.

    Returns:
        Dict with 'processor', 'flush', and 'shutdown'.
    """
    raindrop.init(api_key=api_key)

    processor = RaindropTracingProcessor(user_id=user_id, convo_id=convo_id)

    try:
        from agents.tracing import add_trace_processor
        add_trace_processor(processor)
    except Exception:
        pass

    return {
        "processor": processor,
        "flush": raindrop.flush,
        "shutdown": raindrop.shutdown,
    }


class RaindropTracingProcessor:
    """TracingProcessor implementation that ships events to Raindrop."""

    def __init__(
        self,
        user_id: Optional[str] = None,
        convo_id: Optional[str] = None,
    ):
        self._user_id = user_id
        self._convo_id = convo_id
        self._trace_event_ids: Dict[str, str] = {}
        self._pending_data: Dict[str, Dict[str, Any]] = {}

    def on_trace_start(self, trace: Any) -> None:
        try:
            trace_id = getattr(trace, "trace_id", None)
            if not trace_id:
                return
            event_id = str(uuid.uuid4())
            self._trace_event_ids[trace_id] = event_id
        except Exception:
            pass

    def on_trace_end(self, trace: Any) -> None:
        try:
            trace_id = getattr(trace, "trace_id", None)
            if not trace_id:
                return
            event_id = self._trace_event_ids.pop(trace_id, None)
            pending = self._pending_data.pop(trace_id, {})
            if event_id:
                raindrop.track_ai(
                    user_id=self._user_id or "unknown",
                    event=getattr(trace, "name", "agent_run") or "agent_run",
                    event_id=event_id,
                    input=pending.get("input"),
                    output=pending.get("output"),
                    model=pending.get("model"),
                    properties=pending.get("properties") or None,
                    convo_id=self._convo_id,
                )
        except Exception:
            pass

    def on_span_start(self, span: Any) -> None:
        try:
            trace_id = getattr(span, "trace_id", None)
            if not trace_id or trace_id not in self._trace_event_ids:
                return
        except Exception:
            pass

    def on_span_end(self, span: Any) -> None:
        try:
            trace_id = getattr(span, "trace_id", None)
            if not trace_id or trace_id not in self._trace_event_ids:
                return

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

            span_type = getattr(data, "type", None)

            if span_type == "response":
                self._extract_response_span(trace_id, data)
            elif span_type == "generation":
                self._extract_generation_span(trace_id, data)
        except Exception:
            pass

    def _extract_response_span(self, trace_id: str, data: Any) -> None:
        """Extract data from a 'response' span (OpenAI Agents SDK format)."""
        # NOTE: In multi-response traces (e.g., multi-agent handoffs),
        # later responses overwrite earlier ones. Only the last response's
        # data survives. This is a limitation of the Python SDK's one-shot
        # track_ai call. The TS version preserves all via individual spans.
        pending = self._pending_data.setdefault(trace_id, {})

        input_val = getattr(data, "input", None) or getattr(data, "_input", None)
        if input_val is not None:
            pending["input"] = _format_messages(input_val)

        response = getattr(data, "response", None) or getattr(data, "_response", None)
        if response is None:
            return

        model = getattr(response, "model", None)
        if model:
            pending["model"] = model

        output_parts = getattr(response, "output", None)
        if output_parts:
            texts = []
            for part in output_parts:
                content = getattr(part, "content", None)
                if content:
                    for block in content:
                        text = getattr(block, "text", None)
                        if text:
                            texts.append(text)
            if texts:
                pending["output"] = "\n".join(texts)

        usage = getattr(response, "usage", None)
        if usage:
            props = pending.get("properties", {})
            input_tokens = getattr(usage, "input_tokens", None)
            output_tokens = getattr(usage, "output_tokens", None)
            if input_tokens is not None:
                props["ai.usage.prompt_tokens"] = input_tokens
            if output_tokens is not None:
                props["ai.usage.completion_tokens"] = output_tokens
            if props:
                pending["properties"] = props

    def _extract_generation_span(self, trace_id: str, data: Any) -> None:
        """Extract data from a 'generation' span (generic/older format)."""
        pending = self._pending_data.setdefault(trace_id, {})

        input_val = getattr(data, "input", None)
        if input_val is not None:
            pending["input"] = _format_messages(input_val)

        output_val = getattr(data, "output", None)
        if output_val is not None:
            pending["output"] = _safe_str(output_val)

        model = getattr(data, "model", None)
        if model:
            pending["model"] = model

        usage = getattr(data, "usage", None)
        if usage:
            props = pending.get("properties", {})
            input_tokens = getattr(usage, "input_tokens", None)
            if input_tokens is None and isinstance(usage, dict):
                input_tokens = usage.get("input_tokens")
            output_tokens = getattr(usage, "output_tokens", None)
            if output_tokens is None and isinstance(usage, dict):
                output_tokens = usage.get("output_tokens")
            if input_tokens is not None:
                props["ai.usage.prompt_tokens"] = input_tokens
            if output_tokens is not None:
                props["ai.usage.completion_tokens"] = output_tokens
            if props:
                pending["properties"] = props

    def shutdown(self) -> None:
        self._trace_event_ids.clear()
        self._pending_data.clear()

    def force_flush(self) -> None:
        pass


def _format_messages(input_val: Any) -> Optional[str]:
    """Extract text content from user messages only (no role prefix)."""
    if isinstance(input_val, str):
        return input_val
    if not isinstance(input_val, list):
        return _safe_str(input_val)
    try:
        user_texts = []
        all_texts = []
        for msg in input_val:
            if isinstance(msg, str):
                all_texts.append(msg)
            elif isinstance(msg, dict):
                content = msg.get("content", "")
                text = content if isinstance(content, str) else (_safe_str(content) or "")
                if text:
                    all_texts.append(text)
                    if msg.get("role") == "user":
                        user_texts.append(text)
            else:
                all_texts.append(str(msg))
        texts = user_texts if user_texts else all_texts[-1:]
        return "\n".join(t for t in texts if t) or _safe_str(input_val)
    except Exception:
        return _safe_str(input_val)


def _safe_str(value: Any) -> Optional[str]:
    """Convert a value to string safely."""
    if value is None:
        return None
    if isinstance(value, str):
        return value
    try:
        if hasattr(value, "model_dump_json"):
            return value.model_dump_json()
        return json.dumps(value, default=str)
    except Exception:
        return str(value)
