"""
Cascade SDK - Claude Agent SDK Integration

Provides two approaches for tracing Claude Agent SDK executions:

1. **wrap_claude_query()** (recommended) — wraps the ``query()`` function
   and returns a new async generator with full tracing.  No monkey-patching.

2. **instrument_claude_agents()** (legacy) — monkey-patches the module
   globally.  Kept for backward compatibility.

``wrap_claude_query`` captures:
  - Root agent span with prompt, model, session_id
  - LLM spans per AssistantMessage (tokens, model, completion text)
  - Tool spans per ToolUseBlock ↔ ToolResultBlock (name, input, output)
  - Result metadata (total cost, turns, errors)

Usage::

    from cascade import init_tracing
    from cascade.integrations.claude_agents import wrap_claude_query
    from claude_agent_sdk import query

    init_tracing(project="my_project")
    traced_query = wrap_claude_query(query)

    async for msg in traced_query(prompt="Plan a trip to Tokyo"):
        print(msg)
"""

import logging
import json
import time
import functools
import warnings
from typing import Any, Optional

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

from cascade.tracing import get_tracer

logger = logging.getLogger(__name__)

_MAX_ATTR = 50_000


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


# ============================================================================
# wrap_claude_query  (recommended API)
# ============================================================================

def wrap_claude_query(
    query_fn: Any,
    *,
    name: str = "ClaudeAgentSession",
    session_id: Optional[str] = None,
    record_io: bool = True,
):
    """
    Wrap the Claude Agent SDK ``query()`` function with Cascade tracing.

    Returns a new async-generator function with the same signature as
    ``query()``.  Every call produces a root agent span containing child
    spans for each LLM turn and tool execution.

    Args:
        query_fn: The original ``claude_agent_sdk.query`` function.
        name: Display name for the root span (default ``"ClaudeAgentSession"``).
        session_id: Optional Cascade session ID for multi-turn grouping.
        record_io: Whether to record prompts, completions, and tool I/O
                   as span attributes (default ``True``).

    Usage::

        from claude_agent_sdk import query
        from cascade.integrations.claude_agents import wrap_claude_query

        traced_query = wrap_claude_query(query)

        async for msg in traced_query(prompt="Fix the bug", options=options):
            print(msg)
    """

    async def traced_query(*, prompt: Any, options: Any = None, **kwargs: Any):
        tracer = get_tracer()
        if not tracer:
            async for msg in query_fn(prompt=prompt, options=options, **kwargs):
                yield msg
            return

        # ── Root span ─────────────────────────────────────────────────
        root = tracer.start_span(name, kind=SpanKind.SERVER)
        root.set_attribute("cascade.span_type", "agent")
        if record_io and isinstance(prompt, str):
            root.set_attribute("llm.prompt", _trunc(prompt))
        if session_id:
            root.set_attribute("cascade.session_id", session_id)
        if options and hasattr(options, "model") and options.model:
            root.set_attribute("llm.model", options.model)

        ctx_token = trace.context_api.attach(trace.set_span_in_context(root))

        pending_tools: dict[str, Any] = {}
        total_in = 0
        total_out = 0
        cost_usd = 0.0
        num_turns = 0
        has_error = False
        err_msg = ""
        next_prompt: Optional[str] = _trunc(prompt) if (record_io and isinstance(prompt, str)) else None

        # Track when each LLM turn started so we can record accurate durations.
        # Time is reset after tool results arrive (UserMessage), which is when
        # Claude begins processing the next turn.
        turn_start_ns = time.time_ns()

        try:
            async for msg in query_fn(prompt=prompt, options=options, **kwargs):
                msg_type = type(msg).__name__

                # ── AssistantMessage → LLM span + detect tool calls ───
                if msg_type == "AssistantMessage":
                    prev_turns = num_turns
                    total_in, total_out, num_turns = _handle_assistant(
                        tracer, msg, root, pending_tools,
                        total_in, total_out, num_turns, record_io,
                        start_time_ns=turn_start_ns,
                        first_prompt=next_prompt,
                    )
                    if num_turns > prev_turns:
                        next_prompt = None

                # ── UserMessage → match tool results to pending spans ─
                elif msg_type == "UserMessage":
                    _handle_user(msg, pending_tools, record_io)
                    if record_io:
                        next_prompt = _extract_user_message_summary(msg)
                    # Tool results received → Claude will start its next LLM call now
                    turn_start_ns = time.time_ns()

                # ── ResultMessage → cost, tokens, status ──────────────
                elif msg_type == "ResultMessage":
                    total_in, total_out, cost_usd, num_turns, has_error, err_msg = (
                        _handle_result(msg, root, total_in, total_out, cost_usd, num_turns, record_io)
                    )

                yield msg

        except Exception as exc:
            has_error = True
            err_msg = str(exc)
            root.record_exception(exc)
            raise
        finally:
            for span in pending_tools.values():
                try:
                    span.set_status(Status(StatusCode.ERROR, "span not closed"))
                    span.end()
                except Exception:
                    pass
            pending_tools.clear()

            root.set_attribute("llm.input_tokens", total_in)
            root.set_attribute("llm.output_tokens", total_out)
            root.set_attribute("llm.total_tokens", total_in + total_out)
            root.set_attribute("agent.total_cost_usd", cost_usd)
            root.set_attribute("agent.num_turns", num_turns)
            root.set_status(Status(
                StatusCode.ERROR if has_error else StatusCode.OK,
                err_msg if has_error else None,
            ))
            root.end()
            trace.context_api.detach(ctx_token)

    return traced_query


# ── Message handlers ──────────────────────────────────────────────────────

def _handle_assistant(
    tracer, msg, root, pending_tools, total_in, total_out, num_turns, record_io,
    start_time_ns: Optional[int] = None,
    first_prompt: Optional[str] = None,
):
    """Process an AssistantMessage: create an LLM span and detect tool uses.

    Skips messages that contain no text and no tool calls (e.g.,
    intermediate thinking-only messages the SDK emits internally).
    """
    try:
        content = getattr(msg, "content", None) or []
        text_parts = []
        tool_use_blocks = []
        for block in content:
            btype = type(block).__name__
            if btype == "TextBlock" and hasattr(block, "text") and block.text.strip():
                text_parts.append(block.text)
            elif btype == "ToolUseBlock":
                tool_use_blocks.append(block)

        # Skip AssistantMessages with no meaningful content (thinking-only,
        # empty responses, or SDK internal chatter).
        if not text_parts and not tool_use_blocks:
            return total_in, total_out, num_turns

        end_time_ns = time.time_ns()
        llm = tracer.start_span(
            "anthropic.messages.create",
            kind=SpanKind.CLIENT,
            start_time=start_time_ns,
        )
        llm.set_attribute("cascade.span_type", "llm")
        llm.set_attribute("llm.provider", "anthropic")

        beta_msg = getattr(msg, "message", None)
        if beta_msg:
            model = getattr(beta_msg, "model", None)
            if model:
                llm.set_attribute("llm.model", model)
            usage = getattr(beta_msg, "usage", None)
            if usage:
                in_tok = getattr(usage, "input_tokens", 0) or 0
                out_tok = getattr(usage, "output_tokens", 0) or 0
                llm.set_attribute("llm.input_tokens", in_tok)
                llm.set_attribute("llm.output_tokens", out_tok)
                llm.set_attribute("llm.total_tokens", in_tok + out_tok)
                total_in += in_tok
                total_out += out_tok

        if first_prompt:
            llm.set_attribute("llm.prompt", first_prompt)

        if record_io:
            if text_parts:
                llm.set_attribute("llm.completion", _trunc("\n".join(text_parts)))

        tool_names = []
        if tool_use_blocks:
            llm_ctx = trace.set_span_in_context(llm)
            llm_token = trace.context_api.attach(llm_ctx)
            try:
                for block in tool_use_blocks:
                    tool_names.append(getattr(block, "name", "tool"))
                    _start_tool_span(tracer, block, pending_tools, record_io)
            finally:
                trace.context_api.detach(llm_token)

        if record_io and tool_names:
            llm.set_attribute("llm.tool_calls", ", ".join(tool_names))
            if not text_parts:
                llm.set_attribute("llm.completion", f"[Tool calls: {', '.join(tool_names)}]")

        llm.set_status(Status(StatusCode.OK))
        llm.end(end_time=end_time_ns)
        num_turns += 1

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

    return total_in, total_out, num_turns


def _start_tool_span(tracer, block, pending_tools, record_io):
    """Create a child tool span from a ToolUseBlock."""
    tool_name = getattr(block, "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)
    tool_id = getattr(block, "id", None)
    if tool_id:
        span.set_attribute("tool.use_id", tool_id)
        pending_tools[tool_id] = span
    if record_io and hasattr(block, "input"):
        span.set_attribute("tool.input", _trunc(block.input))


def _extract_user_message_summary(msg) -> Optional[str]:
    """Build an input summary from a UserMessage's tool results."""
    content = getattr(msg, "content", None) or []
    parts = []
    for block in content:
        btype = type(block).__name__
        if btype == "ToolResultBlock":
            tool_id = getattr(block, "tool_use_id", "")
            result = getattr(block, "content", "")
            result_str = _trunc(str(result), 300) if result else ""
            parts.append(f"[tool result] {result_str}")
        elif btype == "TextBlock" and hasattr(block, "text"):
            parts.append(block.text)
    return "\n".join(parts) if parts else None


def _handle_user(msg, pending_tools, record_io):
    """Match tool results in a UserMessage to pending tool spans."""
    content = getattr(msg, "content", None) or []
    for block in content:
        if type(block).__name__ != "ToolResultBlock":
            continue
        tool_id = getattr(block, "tool_use_id", None)
        if not tool_id or tool_id not in pending_tools:
            continue
        span = pending_tools.pop(tool_id)
        try:
            result_content = getattr(block, "content", "")
            if record_io:
                span.set_attribute("tool.output", _trunc(result_content))
            is_err = getattr(block, "is_error", False)
            span.set_status(Status(
                StatusCode.ERROR if is_err else StatusCode.OK,
                _trunc(result_content, 500) if is_err else None,
            ))
        except Exception:
            span.set_status(Status(StatusCode.OK))
        finally:
            span.end()


def _handle_result(msg, root, total_in, total_out, cost_usd, num_turns, record_io):
    """Extract cost, tokens, and status from a ResultMessage."""
    has_error = False
    err_msg = ""
    try:
        if hasattr(msg, "total_cost_usd") and msg.total_cost_usd is not None:
            cost_usd = msg.total_cost_usd
        if hasattr(msg, "num_turns") and msg.num_turns is not None:
            num_turns = msg.num_turns
        usage = getattr(msg, "usage", None)
        if usage:
            total_in = getattr(usage, "input_tokens", total_in) or total_in
            total_out = getattr(usage, "output_tokens", total_out) or total_out
        subtype = getattr(msg, "subtype", "success")
        if subtype != "success":
            has_error = True
            errors = getattr(msg, "errors", None) or []
            err_msg = ", ".join(str(e) for e in errors) if errors else "agent error"
        if record_io:
            result_text = getattr(msg, "result", None)
            if result_text:
                root.set_attribute("llm.completion", _trunc(result_text))
    except Exception as exc:
        logger.debug(f"Cascade: _handle_result error: {exc}")

    return total_in, total_out, cost_usd, num_turns, has_error, err_msg


# ============================================================================
# wrap_claude_client  (for ClaudeSDKClient multi-turn pattern)
# ============================================================================

class TracedClaudeClient:
    """
    Wraps ``ClaudeSDKClient`` with Cascade tracing.

    Each ``query()`` → ``receive_response()`` cycle creates a traced span.

    Usage::

        from claude_agent_sdk import ClaudeSDKClient
        from cascade.integrations.claude_agents import TracedClaudeClient

        async with TracedClaudeClient(ClaudeSDKClient(options), session_id="chat-123") as client:
            await client.query("Hello")
            async for msg in client.receive_response():
                print(msg)
    """

    def __init__(
        self,
        client: Any,
        *,
        name: str = "ClaudeAgentSession",
        session_id: Optional[str] = None,
        record_io: bool = True,
    ):
        self._client = client
        self._name = name
        self._session_id = session_id
        self._record_io = record_io
        self._turn = 0
        self._root: Any = None
        self._ctx_token: Any = None
        self._pending_tools: dict[str, Any] = {}
        self._total_in = 0
        self._total_out = 0
        self._cost_usd = 0.0
        self._num_turns = 0
        self._has_error = False
        self._err_msg = ""
        self._turn_start_ns: int = 0
        self._current_prompt: Optional[str] = None

    async def __aenter__(self):
        await self._client.__aenter__()
        return self

    async def __aexit__(self, *args):
        self._end_root()
        await self._client.__aexit__(*args)

    async def query(self, prompt: Any) -> None:
        self._turn += 1
        tracer = get_tracer()

        # Each turn gets its own root trace (linked by session_id).
        # End the previous turn's root span first.
        if self._root:
            self._end_root()

        # Reset per-turn counters
        self._total_in = 0
        self._total_out = 0
        self._cost_usd = 0.0
        self._num_turns = 0
        self._has_error = False
        self._err_msg = ""

        if tracer:
            span_name = f"{self._name} (turn {self._turn})" if self._turn > 1 else self._name
            self._root = tracer.start_span(span_name, kind=SpanKind.SERVER)
            self._root.set_attribute("cascade.span_type", "agent")
            self._root.set_attribute("agent.turn", self._turn)
            if self._record_io and isinstance(prompt, str):
                self._root.set_attribute("llm.prompt", _trunc(prompt))
            if self._session_id:
                self._root.set_attribute("cascade.session_id", self._session_id)
            self._ctx_token = trace.context_api.attach(trace.set_span_in_context(self._root))

        self._current_prompt = _trunc(prompt) if (self._record_io and isinstance(prompt, str)) else None
        self._turn_start_ns = time.time_ns()
        await self._client.query(prompt)

    async def receive_response(self):
        tracer = get_tracer()
        async for msg in self._client.receive_response():
            if tracer and self._root:
                msg_type = type(msg).__name__
                if msg_type == "AssistantMessage":
                    self._total_in, self._total_out, self._num_turns = _handle_assistant(
                        tracer, msg, self._root, self._pending_tools,
                        self._total_in, self._total_out, self._num_turns, self._record_io,
                        start_time_ns=self._turn_start_ns,
                        first_prompt=self._current_prompt,
                    )
                    self._current_prompt = None
                elif msg_type == "UserMessage":
                    _handle_user(msg, self._pending_tools, self._record_io)
                    if self._record_io:
                        self._current_prompt = _extract_user_message_summary(msg)
                    self._turn_start_ns = time.time_ns()
                elif msg_type == "ResultMessage":
                    self._total_in, self._total_out, self._cost_usd, self._num_turns, self._has_error, self._err_msg = (
                        _handle_result(msg, self._root, self._total_in, self._total_out, self._cost_usd, self._num_turns, self._record_io)
                    )
            yield msg

    async def interrupt(self) -> None:
        await self._client.interrupt()

    def _end_root(self):
        if not self._root:
            return
        for span in self._pending_tools.values():
            try:
                span.set_status(Status(StatusCode.ERROR, "span not closed"))
                span.end()
            except Exception:
                pass
        self._pending_tools.clear()
        self._root.set_attribute("llm.input_tokens", self._total_in)
        self._root.set_attribute("llm.output_tokens", self._total_out)
        self._root.set_attribute("llm.total_tokens", self._total_in + self._total_out)
        self._root.set_attribute("agent.total_cost_usd", self._cost_usd)
        self._root.set_attribute("agent.num_turns", self._num_turns)
        self._root.set_status(Status(
            StatusCode.ERROR if self._has_error else StatusCode.OK,
            self._err_msg if self._has_error else None,
        ))
        self._root.end()
        if self._ctx_token:
            trace.context_api.detach(self._ctx_token)
        self._root = None


# ============================================================================
# instrument_claude_agents  (legacy monkey-patch — backward compat)
# ============================================================================

_instrumented = False
_original_query = None
_original_client_query = None


def instrument_claude_agents() -> None:
    """
    Auto-instrument the Claude Agent SDK via monkey-patching.

    .. deprecated::
        Use :func:`wrap_claude_query` or :class:`TracedClaudeClient` instead.
        They are more robust and capture richer metadata (tokens, cost,
        tool I/O).  This function is kept for backward compatibility only.
    """
    global _instrumented, _original_query, _original_client_query

    warnings.warn(
        "instrument_claude_agents() is deprecated. "
        "Use wrap_claude_query(query) or TracedClaudeClient(client) instead.",
        DeprecationWarning,
        stacklevel=2,
    )

    if _instrumented:
        logger.debug("Claude Agent SDK already instrumented, skipping")
        return

    try:
        import claude_agent_sdk
    except ImportError:
        raise ImportError(
            "Claude Agent SDK not found. Install with: pip install claude-agent-sdk"
        )

    try:
        if hasattr(claude_agent_sdk, "query"):
            _original_query = claude_agent_sdk.query
            claude_agent_sdk.query = _wrap_query_legacy(_original_query, "claude_query")
            logger.debug("Patched claude_agent_sdk.query")

        if hasattr(claude_agent_sdk, "ClaudeSDKClient"):
            client_cls = claude_agent_sdk.ClaudeSDKClient
            if hasattr(client_cls, "query"):
                _original_client_query = client_cls.query
                client_cls.query = _wrap_async_method_legacy(client_cls.query, "ClaudeSDKClient.query")
                logger.debug("Patched ClaudeSDKClient.query")
    except Exception as e:
        logger.warning(f"Cascade: Failed to patch Claude Agent SDK: {e}")
        return

    _instrumented = True
    logger.info("Cascade: Claude Agent SDK auto-instrumentation enabled (legacy)")


def _wrap_query_legacy(original_fn, span_name):
    @functools.wraps(original_fn)
    def wrapper(*args, **kwargs):
        tracer = get_tracer()
        if not tracer:
            return original_fn(*args, **kwargs)
        try:
            prompt = kwargs.get("prompt", args[0] if args else "")
            span = tracer.start_span(span_name, kind=SpanKind.SERVER)
            span.set_attribute("cascade.span_type", "function")
            span.set_attribute("function.name", span_name)
            span.set_attribute("function.input", _trunc(prompt))
            ctx_token = trace.context_api.attach(trace.set_span_in_context(span))
        except Exception:
            return original_fn(*args, **kwargs)

        async_iter = original_fn(*args, **kwargs)

        async def traced_iter():
            outputs = []
            try:
                async for msg in async_iter:
                    if hasattr(msg, "content"):
                        c = msg.content
                        if isinstance(c, str):
                            outputs.append(c)
                        elif isinstance(c, list):
                            for block in c:
                                if hasattr(block, "text"):
                                    outputs.append(block.text)
                    yield msg
            except Exception as e:
                span.set_status(Status(StatusCode.ERROR, str(e)))
                span.record_exception(e)
                raise
            finally:
                if outputs:
                    span.set_attribute("function.output", _trunc("\n".join(outputs)))
                span.set_status(Status(StatusCode.OK))
                span.end()
                trace.context_api.detach(ctx_token)

        return traced_iter()

    return wrapper


def _wrap_async_method_legacy(original_method, span_name):
    @functools.wraps(original_method)
    async def wrapper(self_or_cls, *args, **kwargs):
        tracer = get_tracer()
        if not tracer:
            return await original_method(self_or_cls, *args, **kwargs)
        try:
            prompt = args[0] if args else kwargs.get("prompt", kwargs.get("message", ""))
            with tracer.start_as_current_span(span_name, kind=SpanKind.SERVER):
                span = trace.get_current_span()
                span.set_attribute("cascade.span_type", "function")
                span.set_attribute("function.name", span_name)
                span.set_attribute("function.input", _trunc(prompt))
                try:
                    result = await original_method(self_or_cls, *args, **kwargs)
                    span.set_attribute("function.output", _trunc(result))
                    span.set_status(Status(StatusCode.OK))
                    return result
                except Exception as e:
                    span.set_status(Status(StatusCode.ERROR, str(e)))
                    span.record_exception(e)
                    raise
        except Exception as e:
            if any(p in (type(e).__module__ or "") for p in ("anthropic", "claude")):
                raise
            return await original_method(self_or_cls, *args, **kwargs)

    return wrapper
