"""
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
from typing import Any, Dict, Optional

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

from cascade.tracing import get_tracer, get_session_id

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:
    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>"


def _extract_response_text(response: Any) -> Optional[str]:
    """Extract readable text from a Response object's output list."""
    output_items = getattr(response, "output", None)
    if not output_items or not isinstance(output_items, list):
        return None
    texts = []
    for item in output_items:
        item_type = getattr(item, "type", None)
        if item_type == "message":
            for content in getattr(item, "content", []):
                if getattr(content, "type", None) == "output_text":
                    texts.append(getattr(content, "text", ""))
        elif isinstance(item, dict):
            if item.get("type") == "message":
                for content in item.get("content", []):
                    if isinstance(content, dict) and content.get("type") == "output_text":
                        texts.append(content.get("text", ""))
    return "\n".join(texts) if texts else _safe_str(output_items)


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

    def __init__(self):
        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:
        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)

            session_id = get_session_id()
            if session_id:
                span.set_attribute("cascade.session_id", session_id)

            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:
        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
    # ------------------------------------------------------------------

    def on_span_start(self, span_obj: Any) -> None:
        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"

            cascade_type, display_name, attrs = self._map_span(span_type_name, span_data, span_obj)

            otel_span = tracer.start_span(display_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:
        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:
        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:
        pass

    # ------------------------------------------------------------------
    # Internal: map span type + derive display name
    # ------------------------------------------------------------------

    def _map_span(
        self, type_name: str, span_data: Any, span_obj: Any,
    ) -> tuple:
        """Returns (cascade_type, display_name, attrs)."""
        attrs: Dict[str, str] = {}

        # ── AgentSpanData ────────────────────────────────────────────
        if type_name == "AgentSpanData" or (span_data and getattr(span_data, "type", None) == "agent"):
            agent_name = (getattr(span_data, "name", None) if span_data else None) or "Agent"
            attrs["cascade.agent_name"] = agent_name
            tools = getattr(span_data, "tools", None) if span_data else None
            if tools:
                attrs["agent.tools"] = ", ".join(tools)
            handoffs = getattr(span_data, "handoffs", None) if span_data else None
            if handoffs:
                attrs["agent.handoffs"] = ", ".join(handoffs)
            return "agent", agent_name, attrs

        # ── ResponseSpanData (Responses API) ─────────────────────────
        if type_name == "ResponseSpanData" or (span_data and getattr(span_data, "type", None) == "response"):
            response = getattr(span_data, "response", None) if span_data else None
            model = getattr(response, "model", None) if response else None
            display = f"llm.{model}" if model else "llm.response"
            attrs["llm.provider"] = "openai"
            if model:
                attrs["llm.model"] = str(model)
            input_data = getattr(span_data, "input", None) if span_data else None
            if input_data:
                attrs["llm.prompt"] = _safe_str(input_data)
            return "llm", display, attrs

        # ── GenerationSpanData (Chat Completions API) ────────────────
        if type_name == "GenerationSpanData" or (span_data and getattr(span_data, "type", None) == "generation"):
            model = getattr(span_data, "model", None) if span_data else None
            display = f"llm.{model}" if model else "llm.generation"
            attrs["llm.provider"] = "openai"
            if model:
                attrs["llm.model"] = str(model)
            input_data = getattr(span_data, "input", None) if span_data else None
            if input_data:
                attrs["llm.prompt"] = _safe_str(input_data)
            return "llm", display, attrs

        # ── FunctionSpanData ─────────────────────────────────────────
        if type_name == "FunctionSpanData" or (span_data and getattr(span_data, "type", None) == "function"):
            fn_name = (getattr(span_data, "name", None) if span_data else None) or "function"
            attrs["tool.name"] = fn_name
            input_data = getattr(span_data, "input", None) if span_data else None
            if input_data:
                attrs["tool.input"] = _safe_str(input_data)
            return "tool", fn_name, attrs

        # ── HandoffSpanData ──────────────────────────────────────────
        if type_name == "HandoffSpanData" or (span_data and getattr(span_data, "type", None) == "handoff"):
            from_agent = getattr(span_data, "from_agent", None) if span_data else None
            to_agent = getattr(span_data, "to_agent", None) if span_data else None
            display = f"handoff → {to_agent}" if to_agent else "handoff"
            attrs["function.name"] = display
            if from_agent and to_agent:
                attrs["function.input"] = f"Handoff from {from_agent} to {to_agent}"
            return "function", display, attrs

        # ── GuardrailSpanData ────────────────────────────────────────
        if type_name == "GuardrailSpanData" or (span_data and getattr(span_data, "type", None) == "guardrail"):
            gname = (getattr(span_data, "name", None) if span_data else None) or "guardrail"
            attrs["function.name"] = gname
            return "function", gname, attrs

        # ── Default ──────────────────────────────────────────────────
        fallback_name = getattr(span_obj, "name", None) or type_name
        attrs["function.name"] = fallback_name
        return "function", fallback_name, attrs

    # ------------------------------------------------------------------
    # Internal: extract data at span end
    # ------------------------------------------------------------------

    def _extract_end_data(
        self, otel_span: Any, data_type: str, span_data: Any
    ) -> None:
        if not span_data:
            return

        try:
            # ── ResponseSpanData ─────────────────────────────────────
            if data_type == "ResponseSpanData" or getattr(span_data, "type", None) == "response":
                response = getattr(span_data, "response", None)
                if response:
                    model = getattr(response, "model", None)
                    if model:
                        otel_span.set_attribute("llm.model", str(model))

                    completion_text = _extract_response_text(response)
                    if completion_text:
                        otel_span.set_attribute("llm.completion", _safe_str(completion_text))

                    usage = getattr(response, "usage", None)
                    if usage:
                        in_tok = getattr(usage, "input_tokens", None)
                        out_tok = getattr(usage, "output_tokens", None)
                        if in_tok is not None:
                            otel_span.set_attribute("llm.input_tokens", int(in_tok))
                        if out_tok is not None:
                            otel_span.set_attribute("llm.output_tokens", int(out_tok))
                        if in_tok is not None and out_tok is not None:
                            otel_span.set_attribute("llm.total_tokens", int(in_tok) + int(out_tok))

                input_data = getattr(span_data, "input", None)
                if input_data:
                    otel_span.set_attribute("llm.prompt", _safe_str(input_data))

            # ── GenerationSpanData ───────────────────────────────────
            elif data_type == "GenerationSpanData" or getattr(span_data, "type", None) == "generation":
                input_data = getattr(span_data, "input", None)
                if input_data:
                    otel_span.set_attribute("llm.prompt", _safe_str(input_data))

                model = getattr(span_data, "model", None)
                if model:
                    otel_span.set_attribute("llm.model", str(model))

                output = getattr(span_data, "output", None)
                if output:
                    otel_span.set_attribute("llm.completion", _safe_str(output))

                usage = getattr(span_data, "usage", None)
                if usage:
                    in_tok = usage.get("input_tokens") if isinstance(usage, dict) else getattr(usage, "input_tokens", None)
                    out_tok = usage.get("output_tokens") if isinstance(usage, dict) else getattr(usage, "output_tokens", None)
                    if in_tok is not None:
                        otel_span.set_attribute("llm.input_tokens", int(in_tok))
                    if out_tok is not None:
                        otel_span.set_attribute("llm.output_tokens", int(out_tok))
                    if in_tok is not None and out_tok is not None:
                        otel_span.set_attribute("llm.total_tokens", int(in_tok) + int(out_tok))

            # ── FunctionSpanData ─────────────────────────────────────
            elif data_type == "FunctionSpanData" or getattr(span_data, "type", None) == "function":
                input_data = getattr(span_data, "input", None)
                if input_data:
                    otel_span.set_attribute("tool.input", _safe_str(input_data))
                output = getattr(span_data, "output", None)
                if output:
                    otel_span.set_attribute("tool.output", _safe_str(output))

            # ── HandoffSpanData ──────────────────────────────────────
            elif data_type == "HandoffSpanData" or getattr(span_data, "type", None) == "handoff":
                from_agent = getattr(span_data, "from_agent", None)
                to_agent = getattr(span_data, "to_agent", None)
                if from_agent and to_agent:
                    otel_span.set_attribute("function.output", f"Handed off from {from_agent} to {to_agent}")

            # ── GuardrailSpanData ────────────────────────────────────
            elif data_type == "GuardrailSpanData" or getattr(span_data, "type", None) == "guardrail":
                triggered = getattr(span_data, "triggered", None)
                if triggered is not None:
                    otel_span.set_attribute("function.output", f"triggered={triggered}")

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