"""Zero-code-change instrumentation for the Anthropic Python SDK."""
from __future__ import annotations

import functools
import json
import logging
import time
from typing import Any

from ..client import get_client, get_tracer
from ..context import reset_current_span, set_current_span
from ..types import (
    ANTHROPIC_KEY_CONTENT,
    ANTHROPIC_KEY_MESSAGES,
    ANTHROPIC_KEY_MODEL,
    ANTHROPIC_KEY_SYSTEM,
    ANTHROPIC_KEY_TOOLS,
    META_CACHE_CREATION_TOKENS,
    META_CACHE_READ_TOKENS,
    META_MAX_TOKENS,
    META_PROVIDER,
    META_RESPONSE_ID,
    META_STOP_REASON,
    META_STOP_SEQUENCES,
    META_TEMPERATURE,
    META_THINKING,
    META_TIME_TO_FIRST_TOKEN,
    META_TOOL_CHOICE,
    META_TOP_K,
    META_TOP_P,
    SpanKind,
    SpanStatus,
)
from ._streaming import (
    AsyncMessageStreamManagerProxy,
    AsyncStreamProxy,
    MessageStreamManagerProxy,
    StreamProxy,
)

logger = logging.getLogger("staso")

_patched = False


def _safe_json(obj: Any) -> str:
    try:
        return json.dumps(obj, default=str)
    except (TypeError, ValueError, OverflowError):
        return "{}"


def _dump_content(content: Any) -> list[dict]:
    """Extract content blocks from an Anthropic response as serializable dicts."""
    if not content:
        return []
    result = []
    for block in content:
        if hasattr(block, "model_dump"):
            result.append(block.model_dump())
        elif isinstance(block, dict):
            result.append(block)
        else:
            result.append({"type": "unknown", "value": str(block)})
    return result


def patch_anthropic() -> None:
    """Monkey-patch ``anthropic.resources.Messages.create`` (sync and async)
    to automatically emit LLM spans with comprehensive data capture.

    Safe to call multiple times -- only patches once.
    """
    global _patched
    if _patched:
        return

    try:
        from anthropic.resources import AsyncMessages, Messages
    except ImportError:
        raise ImportError(
            "anthropic package is required for this integration. "
            "Install it with: pip install 'staso[anthropic]'"
        )

    _wrap_sync(Messages)
    _wrap_async(AsyncMessages)
    _wrap_stream_sync(Messages)
    _wrap_stream_async(AsyncMessages)
    _patched = True
    logger.debug("staso: anthropic SDK patched")


def _extract_model(args: tuple, kwargs: dict) -> str:
    return str(kwargs.get("model", args[0] if args else ""))


def _capture_input(model: str, kwargs: dict) -> dict:
    """Build the span input dict from request kwargs."""
    inp: dict[str, Any] = {
        ANTHROPIC_KEY_MODEL: model,
    }

    # System prompt (full, not truncated)
    system = kwargs.get("system")
    if system is not None:
        inp[ANTHROPIC_KEY_SYSTEM] = system

    # Tool definitions
    tools = kwargs.get("tools")
    if tools:
        inp[ANTHROPIC_KEY_TOOLS] = tools

    # Conditionally include messages
    client = get_client()
    if client and client.config.capture_messages:
        messages = kwargs.get("messages")
        if messages:
            inp[ANTHROPIC_KEY_MESSAGES] = messages

    return inp


def _capture_metadata(kwargs: dict) -> dict[str, Any]:
    """Extract request parameters into metadata dict."""
    meta: dict[str, Any] = {META_PROVIDER: "anthropic"}

    _optional_keys = [
        ("temperature", META_TEMPERATURE),
        ("max_tokens", META_MAX_TOKENS),
        ("top_p", META_TOP_P),
        ("top_k", META_TOP_K),
        ("stop_sequences", META_STOP_SEQUENCES),
        ("tool_choice", META_TOOL_CHOICE),
        ("thinking", META_THINKING),
    ]
    for kwarg_key, meta_key in _optional_keys:
        val = kwargs.get(kwarg_key)
        if val is not None:
            meta[meta_key] = val

    return meta


def _capture_response(span: Any, response: Any) -> None:
    """Extract response data onto span fields."""
    # Output: full content blocks
    content = getattr(response, "content", None)
    span.output = {ANTHROPIC_KEY_CONTENT: _dump_content(content)}

    # Response metadata
    stop_reason = getattr(response, "stop_reason", None)
    if stop_reason:
        span.metadata[META_STOP_REASON] = stop_reason

    response_id = getattr(response, "id", None)
    if response_id:
        span.metadata[META_RESPONSE_ID] = response_id

    # Token usage
    usage = getattr(response, "usage", None)
    if usage:
        span.input_tokens = getattr(usage, "input_tokens", 0) or 0
        span.output_tokens = getattr(usage, "output_tokens", 0) or 0
        span.total_tokens = span.input_tokens + span.output_tokens

        # Cache token breakdown
        cache_read = getattr(usage, "cache_read_input_tokens", None)
        if cache_read:
            span.metadata[META_CACHE_READ_TOKENS] = cache_read

        cache_creation = getattr(usage, "cache_creation_input_tokens", None)
        if cache_creation:
            span.metadata[META_CACHE_CREATION_TOKENS] = cache_creation

    # Thinking detection
    if content:
        has_thinking = any(getattr(b, "type", None) == "thinking" for b in content)
        if has_thinking:
            span.metadata["has_thinking"] = True


def _wrap_sync(cls: type) -> None:
    original = cls.create

    @functools.wraps(original)
    def _traced(self: Any, *args: Any, **kwargs: Any) -> Any:
        tracer = get_tracer()
        model = _extract_model(args, kwargs)

        if kwargs.get("stream"):
            # Streaming: manual span lifecycle
            span = tracer.start_span("anthropic.messages.create", kind=SpanKind.LLM)
            span.model = model
            span.input = _capture_input(model, kwargs)
            span.metadata.update(_capture_metadata(kwargs))
            token = set_current_span(span)
            start_time = time.time()
            try:
                stream = original(self, *args, **kwargs)
                return StreamProxy(stream, span, start_time, token)
            except Exception as exc:
                span.record_exception(exc)
                span.end()
                reset_current_span(token)
                raise
        else:
            with tracer.start_as_current_span(
                "anthropic.messages.create", kind=SpanKind.LLM
            ) as span:
                span.model = model
                span.input = _capture_input(model, kwargs)
                span.metadata.update(_capture_metadata(kwargs))
                try:
                    response = original(self, *args, **kwargs)
                    _capture_response(span, response)
                    span.status = SpanStatus.OK
                    return response
                except Exception as exc:
                    span.record_exception(exc)
                    raise

    cls.create = _traced  # type: ignore[assignment]


def _wrap_async(cls: type) -> None:
    original = cls.create

    @functools.wraps(original)
    async def _traced(self: Any, *args: Any, **kwargs: Any) -> Any:
        tracer = get_tracer()
        model = _extract_model(args, kwargs)

        if kwargs.get("stream"):
            span = tracer.start_span("anthropic.messages.create", kind=SpanKind.LLM)
            span.model = model
            span.input = _capture_input(model, kwargs)
            span.metadata.update(_capture_metadata(kwargs))
            token = set_current_span(span)
            start_time = time.time()
            try:
                stream = await original(self, *args, **kwargs)
                return AsyncStreamProxy(stream, span, start_time, token)
            except Exception as exc:
                span.record_exception(exc)
                span.end()
                reset_current_span(token)
                raise
        else:
            with tracer.start_as_current_span(
                "anthropic.messages.create", kind=SpanKind.LLM
            ) as span:
                span.model = model
                span.input = _capture_input(model, kwargs)
                span.metadata.update(_capture_metadata(kwargs))
                try:
                    response = await original(self, *args, **kwargs)
                    _capture_response(span, response)
                    span.status = SpanStatus.OK
                    return response
                except Exception as exc:
                    span.record_exception(exc)
                    raise

    cls.create = _traced  # type: ignore[assignment]


def _wrap_stream_sync(cls: type) -> None:
    """Patch Messages.stream to wrap with MessageStreamManagerProxy."""
    if not hasattr(cls, "stream"):
        return
    original = cls.stream

    @functools.wraps(original)
    def _traced(self: Any, *args: Any, **kwargs: Any) -> Any:
        tracer = get_tracer()
        model = _extract_model(args, kwargs)

        span = tracer.start_span("anthropic.messages.stream", kind=SpanKind.LLM)
        span.model = model
        span.input = _capture_input(model, kwargs)
        span.metadata.update(_capture_metadata(kwargs))
        token = set_current_span(span)
        start_time = time.time()

        try:
            manager = original(self, *args, **kwargs)
            return MessageStreamManagerProxy(manager, span, start_time, token)
        except Exception as exc:
            span.record_exception(exc)
            span.end()
            reset_current_span(token)
            raise

    cls.stream = _traced  # type: ignore[assignment]


def _wrap_stream_async(cls: type) -> None:
    """Patch AsyncMessages.stream."""
    if not hasattr(cls, "stream"):
        return
    original = cls.stream

    @functools.wraps(original)
    def _traced(self: Any, *args: Any, **kwargs: Any) -> Any:
        tracer = get_tracer()
        model = _extract_model(args, kwargs)

        span = tracer.start_span("anthropic.messages.stream", kind=SpanKind.LLM)
        span.model = model
        span.input = _capture_input(model, kwargs)
        span.metadata.update(_capture_metadata(kwargs))
        token = set_current_span(span)
        start_time = time.time()

        try:
            manager = original(self, *args, **kwargs)
            return AsyncMessageStreamManagerProxy(manager, span, start_time, token)
        except Exception as exc:
            span.record_exception(exc)
            span.end()
            reset_current_span(token)
            raise

    cls.stream = _traced  # type: ignore[assignment]
