"""
Cascade SDK - OpenAI Agents SDK Auto-Instrumentation

Automatically traces all OpenAI Agents SDK executions by registering
a custom TracingProcessor that translates agent spans into Cascade spans.

Usage::

    from cascade import init_tracing
    from cascade.integrations import instrument_openai_agents

    init_tracing(project="my_openai_agent")
    instrument_openai_agents()

    # Existing OpenAI Agents code -- no changes needed
    from agents import Agent, Runner
    agent = Agent(name="MyAgent", ...)
    result = Runner.run_sync(agent, "Hello")
"""

import logging
import json
import time
from typing import Any, Dict, Optional

from opentelemetry import trace
from opentelemetry.trace import Status, StatusCode, SpanKind

from cascade.tracing import get_tracer

logger = logging.getLogger(__name__)

_instrumented = False


def instrument_openai_agents() -> None:
    """
    Auto-instrument the OpenAI Agents SDK with Cascade tracing.

    Call this once after ``init_tracing()`` and before running any
    agent code. All subsequent ``Runner.run()`` / ``run_sync()`` /
    ``run_streamed()`` calls will be traced automatically.

    Requires: ``pip install openai-agents``

    Raises:
        ImportError: If the openai-agents package is not installed.
    """
    global _instrumented
    if _instrumented:
        logger.debug("OpenAI Agents SDK already instrumented, skipping")
        return

    try:
        from agents.tracing import add_trace_processor
    except ImportError:
        raise ImportError(
            "OpenAI Agents SDK not found. Install with: "
            "pip install openai-agents"
        )

    try:
        processor = CascadeTracingProcessor()
        add_trace_processor(processor)
    except Exception as e:
        logger.warning(f"Cascade: Failed to register OpenAI Agents tracing processor: {e}")
        return

    _instrumented = True
    logger.info("Cascade: OpenAI Agents SDK auto-instrumentation enabled")


def _safe_str(value: Any, max_len: int = 4000) -> str:
    """Safely convert a value to a truncated string."""
    try:
        if isinstance(value, str):
            s = value
        elif isinstance(value, dict) or isinstance(value, list):
            s = json.dumps(value, default=str, ensure_ascii=False)
        else:
            s = str(value)
        return s[:max_len] if len(s) > max_len else s
    except Exception:
        return "<unserializable>"


class CascadeTracingProcessor:
    """
    OpenAI Agents SDK TracingProcessor that translates agent events
    into Cascade OpenTelemetry spans.

    Registered via ``add_trace_processor()`` and receives callbacks for
    every trace and span lifecycle event.
    """

    def __init__(self):
        # trace_id / span_id -> (otel_span, ctx_token)
        self._active_traces: Dict[str, Any] = {}
        self._active_spans: Dict[str, Any] = {}

    # ------------------------------------------------------------------
    # Trace lifecycle (Runner.run / run_sync / run_streamed)
    # ------------------------------------------------------------------

    def on_trace_start(self, trace_obj: Any) -> None:
        """Called when a Runner.run() trace begins."""
        try:
            tracer = get_tracer()
            if not tracer:
                return

            trace_id = getattr(trace_obj, "trace_id", None) or str(id(trace_obj))
            name = getattr(trace_obj, "name", None) or getattr(trace_obj, "workflow_name", "AgentRun")

            span = tracer.start_span(name, kind=SpanKind.SERVER)
            span.set_attribute("cascade.span_type", "function")
            span.set_attribute("function.name", name)

            group_id = getattr(trace_obj, "group_id", None)
            if group_id:
                span.set_attribute("cascade.group_id", str(group_id))
            metadata = getattr(trace_obj, "metadata", None)
            if metadata and isinstance(metadata, dict):
                for k, v in metadata.items():
                    span.set_attribute(f"cascade.{k}", _safe_str(v, 200))

            ctx_token = trace.context_api.attach(trace.set_span_in_context(span))
            self._active_traces[trace_id] = {"span": span, "ctx_token": ctx_token, "name": name}
        except Exception as e:
            logger.debug(f"Cascade: on_trace_start tracing failed: {e}")

    def on_trace_end(self, trace_obj: Any) -> None:
        """Called when a Runner.run() trace completes."""
        try:
            trace_id = getattr(trace_obj, "trace_id", None) or str(id(trace_obj))
            entry = self._active_traces.pop(trace_id, None)
            if not entry:
                return
            entry["span"].set_status(Status(StatusCode.OK))
            entry["span"].end()
            if entry.get("ctx_token"):
                trace.context_api.detach(entry["ctx_token"])
        except Exception as e:
            logger.debug(f"Cascade: on_trace_end tracing failed: {e}")

    # ------------------------------------------------------------------
    # Span lifecycle (agent, generation, function, guardrail, handoff)
    # ------------------------------------------------------------------

    def on_span_start(self, span_obj: Any) -> None:
        """Called when any span begins within a trace."""
        try:
            tracer = get_tracer()
            if not tracer:
                return

            span_id = getattr(span_obj, "span_id", None) or str(id(span_obj))
            span_data = getattr(span_obj, "span_data", None)

            span_type_name = type(span_data).__name__ if span_data else "unknown"
            name = getattr(span_obj, "name", None) or span_type_name

            cascade_type, attrs = self._map_span_type(span_type_name, span_data, name)

            otel_span = tracer.start_span(name, kind=SpanKind.INTERNAL)
            otel_span.set_attribute("cascade.span_type", cascade_type)

            for k, v in attrs.items():
                otel_span.set_attribute(k, v)

            ctx_token = trace.context_api.attach(trace.set_span_in_context(otel_span))
            self._active_spans[span_id] = {
                "span": otel_span,
                "ctx_token": ctx_token,
                "type": cascade_type,
                "data_type": span_type_name,
            }
        except Exception as e:
            logger.debug(f"Cascade: on_span_start tracing failed: {e}")

    def on_span_end(self, span_obj: Any) -> None:
        """Called when a span completes."""
        try:
            span_id = getattr(span_obj, "span_id", None) or str(id(span_obj))
            entry = self._active_spans.pop(span_id, None)
            if not entry:
                return

            otel_span = entry["span"]
            span_data = getattr(span_obj, "span_data", None)

            self._extract_end_data(otel_span, entry["data_type"], span_data)

            otel_span.set_status(Status(StatusCode.OK))
            otel_span.end()
            if entry.get("ctx_token"):
                trace.context_api.detach(entry["ctx_token"])
        except Exception as e:
            logger.debug(f"Cascade: on_span_end tracing failed: {e}")

    # ------------------------------------------------------------------
    # Shutdown / flush
    # ------------------------------------------------------------------

    def shutdown(self) -> None:
        """Clean up resources."""
        # End any dangling spans
        for entry in list(self._active_spans.values()):
            try:
                entry["span"].end()
                if entry.get("ctx_token"):
                    trace.context_api.detach(entry["ctx_token"])
            except Exception:
                pass
        self._active_spans.clear()
        for entry in list(self._active_traces.values()):
            try:
                entry["span"].end()
                if entry.get("ctx_token"):
                    trace.context_api.detach(entry["ctx_token"])
            except Exception:
                pass
        self._active_traces.clear()

    def force_flush(self) -> None:
        """Force processing of queued items."""
        pass

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _map_span_type(
        self, type_name: str, span_data: Any, name: str
    ) -> tuple:
        """Map OpenAI Agents SDK span type to Cascade span type + attributes."""
        attrs: Dict[str, str] = {}

        if "Agent" in type_name or type_name == "AgentSpanData":
            cascade_type = "agent"
            agent_name = name
            if span_data:
                agent_name = getattr(span_data, "name", name) or name
            attrs["cascade.agent_name"] = agent_name
            return cascade_type, attrs

        if "Generation" in type_name or type_name == "GenerationSpanData":
            cascade_type = "llm"
            if span_data:
                model = getattr(span_data, "model", None) or "unknown"
                attrs["llm.model"] = model
                attrs["llm.provider"] = "openai"

                # Input
                input_data = getattr(span_data, "input", None)
                if input_data:
                    attrs["llm.prompt"] = _safe_str(input_data)

            return cascade_type, attrs

        if "Function" in type_name or type_name == "FunctionSpanData":
            cascade_type = "tool"
            if span_data:
                fn_name = getattr(span_data, "name", name) or name
                attrs["tool.name"] = fn_name
                input_data = getattr(span_data, "input", None)
                if input_data:
                    attrs["tool.input"] = _safe_str(input_data)
            return cascade_type, attrs

        if "Guardrail" in type_name or type_name == "GuardrailSpanData":
            cascade_type = "function"
            attrs["function.name"] = name
            return cascade_type, attrs

        if "Handoff" in type_name or type_name == "HandoffSpanData":
            cascade_type = "function"
            attrs["function.name"] = name
            if span_data:
                from_agent = getattr(span_data, "from_agent", None)
                to_agent = getattr(span_data, "to_agent", None)
                if from_agent:
                    attrs["function.input"] = f"Handoff from {from_agent} to {to_agent}"
            return cascade_type, attrs

        # Default: function span
        cascade_type = "function"
        attrs["function.name"] = name
        return cascade_type, attrs

    def _extract_end_data(
        self, otel_span: Any, data_type: str, span_data: Any
    ) -> None:
        """Extract completion data from a finished span."""
        if not span_data:
            return

        try:
            if "Generation" in data_type:
                # LLM completion
                output = getattr(span_data, "output", None)
                if output:
                    otel_span.set_attribute("llm.completion", _safe_str(output))

                # Token usage
                usage = getattr(span_data, "usage", None) or getattr(span_data, "token_usage", None)
                if usage:
                    if hasattr(usage, "input_tokens"):
                        otel_span.set_attribute("llm.input_tokens", usage.input_tokens)
                        otel_span.set_attribute("llm.output_tokens", usage.output_tokens)
                        otel_span.set_attribute("llm.total_tokens",
                                                usage.input_tokens + usage.output_tokens)
                    elif isinstance(usage, dict):
                        otel_span.set_attribute("llm.input_tokens", usage.get("input_tokens", 0))
                        otel_span.set_attribute("llm.output_tokens", usage.get("output_tokens", 0))
                        otel_span.set_attribute("llm.total_tokens", usage.get("total_tokens", 0))

            elif "Function" in data_type:
                output = getattr(span_data, "output", None)
                if output:
                    otel_span.set_attribute("tool.output", _safe_str(output))

            elif "Guardrail" in data_type or "Handoff" in data_type:
                output = getattr(span_data, "output", None) or getattr(span_data, "result", None)
                if output:
                    otel_span.set_attribute("function.output", _safe_str(output))

        except Exception as e:
            logger.debug(f"Error extracting span end data: {e}")
