"""
Cascade SDK - Pydantic AI Integration

Automatically trace Pydantic AI agents: ``run()``, ``run_sync()``,
``run_stream()``, and ``run_stream_sync()`` are all instrumented.

Each agent run produces a trace tree::

    AgentName                       (root agent span)
    ├─ model_request                (LLM span: model, tokens)
    │  └─ get_weather               (tool span: input, output)
    ├─ model_request                (LLM span: model, tokens)
    └─ result                       (final output)

Two APIs are provided:

1. **setup_pydantic_ai()** — instruments *all* agents globally (one call).
2. **instrument_agent(agent)** — instruments a single Agent instance.

Usage::

    from cascade import init_tracing
    from cascade.integrations.pydantic_ai import setup_pydantic_ai
    from pydantic_ai import Agent

    init_tracing(project="my_project")
    setup_pydantic_ai()

    agent = Agent("openai:gpt-4o", system_prompt="You are helpful.")
    result = agent.run_sync("Hello!")
"""

import logging
import json
import time
import functools
from typing import Any, 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__)

_MAX_ATTR = 50_000
_patched = False


def _trunc(value: Any, limit: int = _MAX_ATTR) -> str:
    try:
        s = value if isinstance(value, str) else json.dumps(value, default=str, ensure_ascii=False)
        return s[:limit] if len(s) > limit else s
    except Exception:
        return "<unserializable>"


# ============================================================================
# setup_pydantic_ai  (global, one-call setup)
# ============================================================================

def setup_pydantic_ai(*, record_io: bool = True) -> None:
    """
    Instrument all Pydantic AI agents with Cascade tracing.

    Call once after ``init_tracing()``.  Every subsequent ``agent.run()``,
    ``agent.run_sync()``, ``agent.run_stream()``, and
    ``agent.run_stream_sync()`` will be traced automatically.

    Args:
        record_io: Whether to record prompts, completions, and tool I/O
                   as span attributes (default ``True``).

    Requires: ``pip install pydantic-ai``
    """
    global _patched

    if _patched:
        logger.debug("Pydantic AI already instrumented, skipping")
        return

    try:
        from pydantic_ai import Agent
    except ImportError:
        raise ImportError(
            "Pydantic AI not found. Install with: pip install pydantic-ai"
        )

    _patch_agent_class(Agent, record_io)
    _patched = True
    logger.info("Cascade: Pydantic AI auto-instrumentation enabled")


# ============================================================================
# instrument_agent  (per-instance, explicit)
# ============================================================================

def instrument_agent(agent: Any, *, record_io: bool = True) -> Any:
    """
    Instrument a single Pydantic AI ``Agent`` instance.

    Returns the same agent (mutated) for convenience.

    Args:
        agent: A ``pydantic_ai.Agent`` instance.
        record_io: Record prompts / completions / tool I/O.

    Usage::

        from pydantic_ai import Agent
        from cascade.integrations.pydantic_ai import instrument_agent

        agent = instrument_agent(Agent("openai:gpt-4o"))
        result = agent.run_sync("Hello!")
    """
    if getattr(agent, "_cascade_instrumented", False):
        return agent

    agent.run = _wrap_run_async(agent.run, agent, record_io)
    agent.run_sync = _wrap_run_sync(agent.run_sync, agent, record_io)
    agent.run_stream = _wrap_run_stream(agent.run_stream, agent, record_io)

    if hasattr(agent, "run_stream_sync"):
        agent.run_stream_sync = _wrap_run_stream_sync(agent.run_stream_sync, agent, record_io)

    agent._cascade_instrumented = True
    return agent


# ============================================================================
# Class-level patching  (used by setup_pydantic_ai)
# ============================================================================

def _patch_agent_class(agent_cls: type, record_io: bool) -> None:
    original_run = agent_cls.run
    original_run_sync = agent_cls.run_sync
    original_run_stream = agent_cls.run_stream

    @functools.wraps(original_run)
    async def patched_run(self, user_prompt, **kwargs):
        tracer = get_tracer()
        if not tracer:
            return await original_run(self, user_prompt, **kwargs)
        return await _traced_run(original_run, self, user_prompt, tracer, record_io, **kwargs)

    @functools.wraps(original_run_sync)
    def patched_run_sync(self, user_prompt, **kwargs):
        tracer = get_tracer()
        if not tracer:
            return original_run_sync(self, user_prompt, **kwargs)
        return _traced_run_sync(original_run_sync, self, user_prompt, tracer, record_io, **kwargs)

    @functools.wraps(original_run_stream)
    def patched_run_stream(self, user_prompt, **kwargs):
        tracer = get_tracer()
        if not tracer:
            return original_run_stream(self, user_prompt, **kwargs)
        return _traced_run_stream_ctx(original_run_stream, self, user_prompt, tracer, record_io, **kwargs)

    agent_cls.run = patched_run
    agent_cls.run_sync = patched_run_sync
    agent_cls.run_stream = patched_run_stream

    if hasattr(agent_cls, "run_stream_sync"):
        original_rss = agent_cls.run_stream_sync

        @functools.wraps(original_rss)
        def patched_rss(self, user_prompt, **kwargs):
            tracer = get_tracer()
            if not tracer:
                return original_rss(self, user_prompt, **kwargs)
            return _traced_run_stream_sync_ctx(original_rss, self, user_prompt, tracer, record_io, **kwargs)

        agent_cls.run_stream_sync = patched_rss


# ============================================================================
# Instance-level wrappers  (used by instrument_agent)
# ============================================================================

def _wrap_run_async(original_fn, agent, record_io):
    @functools.wraps(original_fn)
    async def wrapper(user_prompt, **kwargs):
        tracer = get_tracer()
        if not tracer:
            return await original_fn(user_prompt, **kwargs)
        return await _traced_run(original_fn, agent, user_prompt, tracer, record_io, _skip_self=True, **kwargs)
    return wrapper


def _wrap_run_sync(original_fn, agent, record_io):
    @functools.wraps(original_fn)
    def wrapper(user_prompt, **kwargs):
        tracer = get_tracer()
        if not tracer:
            return original_fn(user_prompt, **kwargs)
        return _traced_run_sync(original_fn, agent, user_prompt, tracer, record_io, _skip_self=True, **kwargs)
    return wrapper


def _wrap_run_stream(original_fn, agent, record_io):
    @functools.wraps(original_fn)
    def wrapper(user_prompt, **kwargs):
        tracer = get_tracer()
        if not tracer:
            return original_fn(user_prompt, **kwargs)
        return _traced_run_stream_ctx(original_fn, agent, user_prompt, tracer, record_io, _skip_self=True, **kwargs)
    return wrapper


def _wrap_run_stream_sync(original_fn, agent, record_io):
    @functools.wraps(original_fn)
    def wrapper(user_prompt, **kwargs):
        tracer = get_tracer()
        if not tracer:
            return original_fn(user_prompt, **kwargs)
        return _traced_run_stream_sync_ctx(original_fn, agent, user_prompt, tracer, record_io, _skip_self=True, **kwargs)
    return wrapper


# ============================================================================
# Core traced execution
# ============================================================================

def _agent_name(agent: Any) -> str:
    return getattr(agent, "name", None) or type(agent).__name__


def _model_name(agent: Any) -> str:
    model = getattr(agent, "model", None)
    if model is None:
        return "unknown"
    if isinstance(model, str):
        return model
    return getattr(model, "model_name", None) or str(model)


async def _traced_run(original_fn, agent, user_prompt, tracer, record_io, *, _skip_self=False, **kwargs):
    span_name = _agent_name(agent)
    root = tracer.start_span(span_name, kind=SpanKind.SERVER)
    root.set_attribute("cascade.span_type", "agent")
    root.set_attribute("llm.model", _model_name(agent))
    session_id = get_session_id()
    if session_id:
        root.set_attribute("cascade.session_id", session_id)
    if record_io:
        root.set_attribute("llm.prompt", _trunc(user_prompt))

    ctx_token = trace.context_api.attach(trace.set_span_in_context(root))
    try:
        if _skip_self:
            result = await original_fn(user_prompt, **kwargs)
        else:
            result = await original_fn(agent, user_prompt, **kwargs)

        _record_child_spans(tracer, result, record_io)
        _set_result_attrs(root, result, record_io)
        root.set_status(Status(StatusCode.OK))
        return result
    except Exception as exc:
        root.set_status(Status(StatusCode.ERROR, str(exc)))
        root.record_exception(exc)
        raise
    finally:
        root.end()
        trace.context_api.detach(ctx_token)


def _traced_run_sync(original_fn, agent, user_prompt, tracer, record_io, *, _skip_self=False, **kwargs):
    span_name = _agent_name(agent)
    root = tracer.start_span(span_name, kind=SpanKind.SERVER)
    root.set_attribute("cascade.span_type", "agent")
    root.set_attribute("llm.model", _model_name(agent))
    session_id = get_session_id()
    if session_id:
        root.set_attribute("cascade.session_id", session_id)
    if record_io:
        root.set_attribute("llm.prompt", _trunc(user_prompt))

    ctx_token = trace.context_api.attach(trace.set_span_in_context(root))
    try:
        if _skip_self:
            result = original_fn(user_prompt, **kwargs)
        else:
            result = original_fn(agent, user_prompt, **kwargs)

        _record_child_spans(tracer, result, record_io)
        _set_result_attrs(root, result, record_io)
        root.set_status(Status(StatusCode.OK))
        return result
    except Exception as exc:
        root.set_status(Status(StatusCode.ERROR, str(exc)))
        root.record_exception(exc)
        raise
    finally:
        root.end()
        trace.context_api.detach(ctx_token)


def _traced_run_stream_ctx(original_fn, agent, user_prompt, tracer, record_io, *, _skip_self=False, **kwargs):
    """Return an async context manager that wraps run_stream with tracing."""
    import contextlib

    @contextlib.asynccontextmanager
    async def wrapper():
        span_name = _agent_name(agent)
        root = tracer.start_span(span_name, kind=SpanKind.SERVER)
        root.set_attribute("cascade.span_type", "agent")
        root.set_attribute("llm.model", _model_name(agent))
        session_id = get_session_id()
        if session_id:
            root.set_attribute("cascade.session_id", session_id)
        if record_io:
            root.set_attribute("llm.prompt", _trunc(user_prompt))

        ctx_token = trace.context_api.attach(trace.set_span_in_context(root))
        try:
            if _skip_self:
                cm = original_fn(user_prompt, **kwargs)
            else:
                cm = original_fn(agent, user_prompt, **kwargs)

            async with cm as streamed_result:
                yield streamed_result

            _set_result_attrs(root, streamed_result, record_io)
            root.set_status(Status(StatusCode.OK))
        except Exception as exc:
            root.set_status(Status(StatusCode.ERROR, str(exc)))
            root.record_exception(exc)
            raise
        finally:
            root.end()
            trace.context_api.detach(ctx_token)

    return wrapper()


def _traced_run_stream_sync_ctx(original_fn, agent, user_prompt, tracer, record_io, *, _skip_self=False, **kwargs):
    """Return a sync context manager that wraps run_stream_sync with tracing."""
    import contextlib

    @contextlib.contextmanager
    def wrapper():
        span_name = _agent_name(agent)
        root = tracer.start_span(span_name, kind=SpanKind.SERVER)
        root.set_attribute("cascade.span_type", "agent")
        root.set_attribute("llm.model", _model_name(agent))
        session_id = get_session_id()
        if session_id:
            root.set_attribute("cascade.session_id", session_id)
        if record_io:
            root.set_attribute("llm.prompt", _trunc(user_prompt))

        ctx_token = trace.context_api.attach(trace.set_span_in_context(root))
        try:
            if _skip_self:
                cm = original_fn(user_prompt, **kwargs)
            else:
                cm = original_fn(agent, user_prompt, **kwargs)

            with cm as streamed_result:
                yield streamed_result

            _set_result_attrs(root, streamed_result, record_io)
            root.set_status(Status(StatusCode.OK))
        except Exception as exc:
            root.set_status(Status(StatusCode.ERROR, str(exc)))
            root.record_exception(exc)
            raise
        finally:
            root.end()
            trace.context_api.detach(ctx_token)

    return wrapper()


# ============================================================================
# Post-run: reconstruct child spans from result messages
# ============================================================================

def _record_child_spans(tracer: Any, result: Any, record_io: bool) -> None:
    """
    Walk result.all_messages() to create LLM and tool child spans.

    Pydantic AI message flow::

        ModelRequest  (UserPromptPart)
        ModelResponse (ToolCallPart... + usage)    ← LLM turn 1
        ModelRequest  (ToolReturnPart...)           ← tool results
        ModelResponse (TextPart + usage)            ← LLM turn 2

    We create LLM spans for each ModelResponse.  When a ModelResponse
    contains ToolCallParts, we create tool spans as **children** of that
    LLM span (so they nest correctly in the trace tree).  Tool outputs
    come from the ToolReturnParts in the *next* ModelRequest.
    """
    try:
        messages = result.all_messages() if callable(getattr(result, "all_messages", None)) else []
    except Exception:
        return

    # Pre-index tool return outputs by tool_call_id so we can attach
    # them to the tool spans created under the preceding LLM span.
    tool_returns: dict[str, Any] = {}
    for msg in messages:
        for part in getattr(msg, "parts", []) or []:
            if type(part).__name__ == "ToolReturnPart":
                tid = getattr(part, "tool_call_id", None)
                if tid:
                    tool_returns[tid] = getattr(part, "content", "")

    # Walk messages in order and create properly nested spans.
    prev_request = None

    for msg in messages:
        msg_cls = type(msg).__name__

        if msg_cls == "ModelRequest":
            prev_request = msg

        elif msg_cls == "ModelResponse":
            _create_llm_and_tool_spans(tracer, msg, prev_request, tool_returns, record_io)
            prev_request = None


def _extract_request_input(request_msg: Any) -> str:
    """Build a human-readable input string from a ModelRequest's parts."""
    parts = getattr(request_msg, "parts", []) or []
    pieces = []
    for part in parts:
        pcls = type(part).__name__
        if pcls == "UserPromptPart":
            pieces.append(getattr(part, "content", ""))
        elif pcls == "SystemPromptPart":
            pieces.append(f"[system] {getattr(part, 'content', '')}")
        elif pcls == "ToolReturnPart":
            name = getattr(part, "tool_name", "tool")
            content = getattr(part, "content", "")
            pieces.append(f"[tool result: {name}] {content}")
    return "\n".join(pieces)


def _create_llm_and_tool_spans(tracer, response_msg, prev_request, tool_returns, record_io):
    """
    Create an LLM span for a ModelResponse.  If the response contains
    ToolCallParts, create tool spans as **children** of the LLM span
    so they nest correctly in the trace tree.
    """
    try:
        llm_span = tracer.start_span("model_request", kind=SpanKind.CLIENT)
        llm_span.set_attribute("cascade.span_type", "llm")

        model_name = getattr(response_msg, "model_name", None)
        if model_name:
            llm_span.set_attribute("llm.model", model_name)

        usage = getattr(response_msg, "usage", None)
        if usage:
            in_tok = getattr(usage, "input_tokens", 0) or 0
            out_tok = getattr(usage, "output_tokens", 0) or 0
            llm_span.set_attribute("llm.input_tokens", in_tok)
            llm_span.set_attribute("llm.output_tokens", out_tok)
            llm_span.set_attribute("llm.total_tokens", in_tok + out_tok)

        if record_io:
            if prev_request:
                input_text = _extract_request_input(prev_request)
                if input_text:
                    llm_span.set_attribute("llm.prompt", _trunc(input_text))

            parts = getattr(response_msg, "parts", []) or []
            text_parts = []
            tool_call_parts_io = []
            for part in parts:
                pcls = type(part).__name__
                if pcls == "TextPart":
                    text_parts.append(getattr(part, "content", ""))
                elif pcls == "ToolCallPart":
                    tc_name = getattr(part, "tool_name", "tool")
                    tc_args = getattr(part, "args", None)
                    if tc_args:
                        tool_call_parts_io.append(f"[tool_call] {tc_name}({_trunc(tc_args, 200)})")
                    else:
                        tool_call_parts_io.append(f"[tool_call] {tc_name}()")
            if text_parts:
                llm_span.set_attribute("llm.completion", _trunc("\n".join(text_parts)))
            elif tool_call_parts_io:
                llm_span.set_attribute("llm.completion", _trunc("\n".join(tool_call_parts_io)))

        # Create tool spans as children of this LLM span by setting it
        # as the active span in context before creating them.
        parts = getattr(response_msg, "parts", []) or []
        tool_call_parts = [p for p in parts if type(p).__name__ == "ToolCallPart"]

        if tool_call_parts:
            ctx = trace.set_span_in_context(llm_span)
            token = trace.context_api.attach(ctx)
            try:
                for tc in tool_call_parts:
                    _create_tool_span(tracer, tc, tool_returns, record_io)
            finally:
                trace.context_api.detach(token)

        llm_span.set_status(Status(StatusCode.OK))
        llm_span.end()

    except Exception as exc:
        logger.debug(f"Cascade: _create_llm_and_tool_spans error: {exc}")


def _create_tool_span(tracer, tool_call_part, tool_returns, record_io):
    """Create a tool span from a ToolCallPart, pairing with its ToolReturnPart output."""
    try:
        tool_name = getattr(tool_call_part, "tool_name", "tool")
        span = tracer.start_span(tool_name, kind=SpanKind.INTERNAL)
        span.set_attribute("cascade.span_type", "tool")
        span.set_attribute("tool.name", tool_name)

        tid = getattr(tool_call_part, "tool_call_id", None)
        if tid:
            span.set_attribute("tool.call_id", tid)

        if record_io:
            args = getattr(tool_call_part, "args", None)
            if args is not None:
                span.set_attribute("tool.input", _trunc(args))
            if tid and tid in tool_returns:
                span.set_attribute("tool.output", _trunc(tool_returns[tid]))

        span.set_status(Status(StatusCode.OK))
        span.end()
    except Exception as exc:
        logger.debug(f"Cascade: _create_tool_span error: {exc}")


def _set_result_attrs(root_span: Any, result: Any, record_io: bool) -> None:
    """Set summary attributes on the root span from the run result."""
    try:
        usage_fn = getattr(result, "usage", None)
        if callable(usage_fn):
            usage = usage_fn()
            total = getattr(usage, "total_tokens", 0) or 0
            req = getattr(usage, "input_tokens", 0) or 0
            resp = getattr(usage, "output_tokens", 0) or 0
            root_span.set_attribute("llm.input_tokens", req)
            root_span.set_attribute("llm.output_tokens", resp)
            root_span.set_attribute("llm.total_tokens", total or (req + resp))

        if record_io:
            output = getattr(result, "output", None)
            if output is not None:
                root_span.set_attribute("llm.completion", _trunc(output))

    except Exception as exc:
        logger.debug(f"Cascade: _set_result_attrs error: {exc}")
