"""
Cascade SDK - Claude Agent SDK Auto-Instrumentation

Automatically traces Claude Agent SDK executions by monkey-patching
the ``query()`` function and ``ClaudeSDKClient`` methods.

Note: The Claude Agent SDK does not yet have a native tracing/callback
interface (GitHub issue #452). This integration uses monkey-patching
which is inherently more fragile than callback-based approaches.

Usage::

    from cascade import init_tracing
    from cascade.integrations import instrument_claude_agents

    init_tracing(project="my_claude_agent")
    instrument_claude_agents()

    # Existing Claude Agent SDK code -- no changes needed
    from claude_agent_sdk import query
    result = await query("Plan a trip to Tokyo")
"""

import logging
import json
import time
import functools
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
_original_query = None
_original_client_query = None


def instrument_claude_agents() -> None:
    """
    Auto-instrument the Claude Agent SDK with Cascade tracing.

    Call this once after ``init_tracing()`` and before running any
    Claude Agent SDK code. All subsequent ``query()`` and
    ``ClaudeSDKClient`` calls will be traced automatically.

    Requires: ``pip install claude-agent-sdk``

    Raises:
        ImportError: If the claude-agent-sdk package is not installed.
    """
    global _instrumented, _original_query, _original_client_query

    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(_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(client_cls.query, "ClaudeSDKClient.query")
                logger.debug("Patched ClaudeSDKClient.query")

        _patch_tool_decorator(claude_agent_sdk)
    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")


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


def _wrap_query(original_fn: Any, span_name: str) -> Any:
    """
    Wrap the ``query()`` function which returns an AsyncIterator.

    Claude Agent SDK's ``query()`` is a regular function that returns
    an async iterator. We wrap it to start a span when called and
    collect output as messages are yielded.
    """

    @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", _safe_str(prompt))
            ctx_token = trace.context_api.attach(trace.set_span_in_context(span))
        except Exception as e:
            logger.debug(f"Cascade: _wrap_query span creation failed: {e}")
            return original_fn(*args, **kwargs)

        # Get the original async iterator
        async_iter = original_fn(*args, **kwargs)

        # Wrap the async iterator to capture output and close the span
        async def traced_iter():
            outputs = []
            try:
                async for msg in async_iter:
                    # Capture text content from result messages
                    msg_type = type(msg).__name__
                    if hasattr(msg, "content"):
                        content = msg.content
                        if isinstance(content, str):
                            outputs.append(content)
                        elif isinstance(content, list):
                            for block in content:
                                if hasattr(block, "text"):
                                    outputs.append(block.text)
                    yield msg
            except Exception as e:
                span.set_attribute("function.error", str(e))
                span.set_status(Status(StatusCode.ERROR, str(e)))
                span.record_exception(e)
                raise
            finally:
                if outputs:
                    span.set_attribute("function.output", _safe_str("\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(original_method: Any, span_name: str) -> Any:
    """Wrap an async method on a class in a Cascade span."""

    @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) as span_ctx:
                span = trace.get_current_span()
                span.set_attribute("cascade.span_type", "function")
                span.set_attribute("function.name", span_name)
                span.set_attribute("function.input", _safe_str(prompt))

                try:
                    result = await original_method(self_or_cls, *args, **kwargs)
                    span.set_attribute("function.output", _safe_str(result))
                    span.set_status(Status(StatusCode.OK))
                    return result
                except Exception as e:
                    span.set_attribute("function.error", str(e))
                    span.set_status(Status(StatusCode.ERROR, str(e)))
                    span.record_exception(e)
                    raise
        except Exception as e:
            if any(provider in (type(e).__module__ or "") for provider in ("anthropic", "claude")):
                raise
            logger.debug(f"Cascade: _wrap_async_method tracing failed: {e}")
            return await original_method(self_or_cls, *args, **kwargs)

    return wrapper


def _patch_tool_decorator(sdk_module: Any) -> None:
    """
    Patch the Claude Agent SDK's ``@tool`` decorator so that tool
    executions are automatically wrapped in Cascade tool spans.
    """
    if not hasattr(sdk_module, "tool"):
        return

    original_tool = sdk_module.tool

    def patched_tool(*tool_args, **tool_kwargs):
        """Wrapping the @tool decorator to add Cascade tracing."""
        # Call the original @tool decorator
        result = original_tool(*tool_args, **tool_kwargs)

        if not callable(result):
            return result

        # The result is the decorated tool function -- wrap it
        original_handler = result

        if hasattr(original_handler, "__wrapped__"):
            # Already wrapped
            return original_handler

        import inspect

        if inspect.iscoroutinefunction(original_handler):
            @functools.wraps(original_handler)
            async def traced_tool_handler(args_dict):
                tracer = get_tracer()
                if not tracer:
                    return await original_handler(args_dict)

                tool_name = tool_args[0] if tool_args else getattr(original_handler, "__name__", "tool")

                with tracer.start_as_current_span(tool_name, kind=SpanKind.INTERNAL) as span_ctx:
                    span = trace.get_current_span()
                    span.set_attribute("cascade.span_type", "tool")
                    span.set_attribute("tool.name", tool_name)
                    span.set_attribute("tool.input", _safe_str(args_dict))

                    try:
                        result = await original_handler(args_dict)
                        span.set_attribute("tool.output", _safe_str(result))
                        span.set_status(Status(StatusCode.OK))
                        return result
                    except Exception as e:
                        span.set_attribute("tool.error", str(e))
                        span.set_status(Status(StatusCode.ERROR, str(e)))
                        span.record_exception(e)
                        raise

            traced_tool_handler.__wrapped__ = True
            return traced_tool_handler
        else:
            @functools.wraps(original_handler)
            def traced_tool_handler(args_dict):
                tracer = get_tracer()
                if not tracer:
                    return original_handler(args_dict)

                tool_name = tool_args[0] if tool_args else getattr(original_handler, "__name__", "tool")

                with tracer.start_as_current_span(tool_name, kind=SpanKind.INTERNAL) as span_ctx:
                    span = trace.get_current_span()
                    span.set_attribute("cascade.span_type", "tool")
                    span.set_attribute("tool.name", tool_name)
                    span.set_attribute("tool.input", _safe_str(args_dict))

                    try:
                        result = original_handler(args_dict)
                        span.set_attribute("tool.output", _safe_str(result))
                        span.set_status(Status(StatusCode.OK))
                        return result
                    except Exception as e:
                        span.set_attribute("tool.error", str(e))
                        span.set_status(Status(StatusCode.ERROR, str(e)))
                        span.record_exception(e)
                        raise

            traced_tool_handler.__wrapped__ = True
            return traced_tool_handler

    sdk_module.tool = patched_tool
    logger.debug("Patched claude_agent_sdk.tool decorator")
