"""
Anthropic SDK instrumentation.

Patches Anthropic SDK methods to capture:
- Messages (sync and async)
- Streaming responses
- Tool use

Based on ARCHITECTURE_V2.md Section 6.3.
"""

from __future__ import annotations

import logging
import time
from typing import Any, AsyncIterator, Callable, Dict, Iterator, Optional, TypeVar

import wrapt

from risicare.integrations._base import scrub_sensitive

logger = logging.getLogger(__name__)

T = TypeVar("T")

# Track if already instrumented
_instrumented = False


def instrument_anthropic(module: Any) -> None:
    """
    Apply instrumentation to Anthropic module.

    This patches:
    - anthropic.resources.messages.Messages.create
    - anthropic.resources.messages.AsyncMessages.create
    - anthropic.resources.messages.Messages.stream (context manager)
    - anthropic.resources.messages.AsyncMessages.stream (async context manager)
    """
    global _instrumented
    if _instrumented:
        return

    try:
        # Patch sync messages.create
        wrapt.wrap_function_wrapper(
            module,
            "resources.messages.Messages.create",
            _wrap_messages_create,
        )

        # Patch async messages.create
        wrapt.wrap_function_wrapper(
            module,
            "resources.messages.AsyncMessages.create",
            _wrap_async_messages_create,
        )

        _instrumented = True
        logger.debug("Instrumented Anthropic SDK")

    except Exception as e:
        logger.warning(f"Failed to instrument Anthropic SDK: {e}")


def _get_tracer() -> Optional[Any]:
    """Get the Risicare tracer if available."""
    try:
        from risicare.tracer import get_tracer

        return get_tracer()
    except ImportError:
        return None


def _get_config() -> Optional[Any]:
    """Get the Risicare client config if available."""
    try:
        from risicare.client import get_client

        client = get_client()
        return client.config if client else None
    except ImportError:
        return None


def _should_trace_content() -> bool:
    """Check if content should be traced. Defaults to False (fail-closed for privacy)."""
    config = _get_config()
    return config.trace_content if config else False


def _extract_model_info(kwargs: Dict[str, Any]) -> Dict[str, Any]:
    """Extract model information from request kwargs.

    Only JSON-primitive values are included — SDK sentinels are skipped.
    """
    from risicare.integrations._base import is_json_primitive

    stream = kwargs.get("stream", False)
    temperature = kwargs.get("temperature")
    max_tokens = kwargs.get("max_tokens")
    tools = kwargs.get("tools")

    info: Dict[str, Any] = {
        "gen_ai.system": "anthropic",
        "gen_ai.request.model": kwargs.get("model", "unknown"),
        "gen_ai.request.has_tools": isinstance(tools, (list, tuple)) and len(tools) > 0,
    }

    if is_json_primitive(max_tokens):
        info["gen_ai.request.max_tokens"] = max_tokens
    if is_json_primitive(temperature):
        info["gen_ai.request.temperature"] = temperature
    if is_json_primitive(stream):
        info["gen_ai.request.stream"] = stream

    return info


def _capture_messages(span: Any, messages: list, trace_content: bool) -> None:
    """Capture input messages to span."""
    if not trace_content:
        span.set_attribute("gen_ai.prompt.count", len(messages))
        return

    for i, msg in enumerate(messages[:10]):  # Limit to 10 messages
        span.set_attribute(f"gen_ai.prompt.{i}.role", msg.get("role", ""))

        content = msg.get("content", "")
        if isinstance(content, str):
            span.set_attribute(
                f"gen_ai.prompt.{i}.content",
                content[:10000] if len(content) > 10000 else content,
            )
        elif isinstance(content, list):
            # Handle content blocks
            text_parts = []
            for block in content:
                if isinstance(block, dict) and block.get("type") == "text":
                    text_parts.append(block.get("text", ""))
            combined = "\n".join(text_parts)
            span.set_attribute(
                f"gen_ai.prompt.{i}.content",
                combined[:10000] if len(combined) > 10000 else combined,
            )


def _capture_response(
    span: Any,
    response: Any,
    latency_ms: float,
    trace_content: bool,
) -> None:
    """Capture response data into span."""
    # Model info
    span.set_attribute("gen_ai.response.model", getattr(response, "model", "unknown"))
    span.set_attribute("gen_ai.response.id", getattr(response, "id", ""))
    span.set_attribute("gen_ai.latency_ms", latency_ms)
    span.set_attribute("gen_ai.response.stop_reason", getattr(response, "stop_reason", ""))

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

    # Content
    content_blocks = getattr(response, "content", [])
    span.set_attribute("gen_ai.response.content_blocks", len(content_blocks))

    if trace_content and content_blocks:
        text_parts = []
        tool_uses = []

        for block in content_blocks:
            block_type = getattr(block, "type", "")
            if block_type == "text":
                text_parts.append(getattr(block, "text", ""))
            elif block_type == "tool_use":
                tool_uses.append({
                    "name": getattr(block, "name", ""),
                    "id": getattr(block, "id", ""),
                })

        if text_parts:
            combined = "\n".join(text_parts)
            span.set_attribute(
                "gen_ai.completion.content",
                combined[:10000] if len(combined) > 10000 else combined,
            )

        if tool_uses:
            span.set_attribute("gen_ai.completion.tool_uses", len(tool_uses))
            for j, tu in enumerate(tool_uses[:5]):  # Limit to 5
                span.set_attribute(f"gen_ai.completion.tool_use.{j}.name", tu["name"])


def _wrap_messages_create(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for anthropic.messages.create (sync)."""
    tracer = _get_tracer()
    if tracer is None or not tracer.is_enabled:
        return wrapped(*args, **kwargs)

    from risicare.integrations._dedup import is_provider_suppressed

    if is_provider_suppressed():
        return wrapped(*args, **kwargs)

    from risicare_core import SpanKind

    model = kwargs.get("model", "unknown")
    trace_content = _should_trace_content()
    stream = kwargs.get("stream", False)

    # NON-STREAMING: use context manager (span auto-ended on block exit)
    if not stream:
        with tracer.start_span(
            name=f"anthropic.messages.create/{model}",
            kind=SpanKind.LLM_CALL,
            attributes=_extract_model_info(kwargs),
        ) as span:
            # Capture system message
            system_msg = kwargs.get("system")
            if system_msg and trace_content:
                if isinstance(system_msg, str):
                    span.set_attribute(
                        "gen_ai.system_prompt",
                        system_msg[:5000] if len(system_msg) > 5000 else system_msg,
                    )

            # Capture input messages
            messages = kwargs.get("messages", [])
            _capture_messages(span, messages, trace_content)

            try:
                start_time = time.perf_counter()
                response = wrapped(*args, **kwargs)
                latency_ms = (time.perf_counter() - start_time) * 1000
                _capture_response(span, response, latency_ms, trace_content)
                return response
            except Exception as e:
                span.record_exception(e)
                span.set_attribute("error", True)
                span.set_attribute("error.type", type(e).__name__)
                span.set_attribute("error.message", scrub_sensitive(str(e)[:2000]))
                raise

    # STREAMING: manual span lifecycle — span.end() deferred to generator finally
    span = tracer.create_span(
        name=f"anthropic.messages.create/{model}",
        kind=SpanKind.LLM_CALL,
        attributes=_extract_model_info(kwargs),
    )

    # Capture system message
    system_msg = kwargs.get("system")
    if system_msg and trace_content:
        if isinstance(system_msg, str):
            span.set_attribute(
                "gen_ai.system_prompt",
                system_msg[:5000] if len(system_msg) > 5000 else system_msg,
            )

    # Capture input messages
    messages = kwargs.get("messages", [])
    _capture_messages(span, messages, trace_content)

    try:
        start_time = time.perf_counter()
        result = wrapped(*args, **kwargs)
        return _wrap_stream_sync(result, span, start_time, trace_content, tracer)
    except Exception as e:
        span.record_exception(e)
        span.set_attribute("error", True)
        span.set_attribute("error.type", type(e).__name__)
        span.set_attribute("error.message", scrub_sensitive(str(e)[:2000]))
        span.end()
        tracer.export_span(span)
        raise


async def _wrap_async_messages_create(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for anthropic.messages.create (async)."""
    tracer = _get_tracer()
    if tracer is None or not tracer.is_enabled:
        return await wrapped(*args, **kwargs)

    from risicare.integrations._dedup import is_provider_suppressed

    if is_provider_suppressed():
        return await wrapped(*args, **kwargs)

    from risicare_core import SpanKind

    model = kwargs.get("model", "unknown")
    trace_content = _should_trace_content()
    stream = kwargs.get("stream", False)

    # NON-STREAMING: use context manager (span auto-ended on block exit)
    if not stream:
        with tracer.start_span(
            name=f"anthropic.messages.create/{model}",
            kind=SpanKind.LLM_CALL,
            attributes=_extract_model_info(kwargs),
        ) as span:
            # Capture system message
            system_msg = kwargs.get("system")
            if system_msg and trace_content:
                if isinstance(system_msg, str):
                    span.set_attribute(
                        "gen_ai.system_prompt",
                        system_msg[:5000] if len(system_msg) > 5000 else system_msg,
                    )

            # Capture input messages
            messages = kwargs.get("messages", [])
            _capture_messages(span, messages, trace_content)

            try:
                start_time = time.perf_counter()
                response = await wrapped(*args, **kwargs)
                latency_ms = (time.perf_counter() - start_time) * 1000
                _capture_response(span, response, latency_ms, trace_content)
                return response
            except Exception as e:
                span.record_exception(e)
                span.set_attribute("error", True)
                span.set_attribute("error.type", type(e).__name__)
                span.set_attribute("error.message", scrub_sensitive(str(e)[:2000]))
                raise

    # STREAMING: manual span lifecycle — span.end() deferred to generator finally
    span = tracer.create_span(
        name=f"anthropic.messages.create/{model}",
        kind=SpanKind.LLM_CALL,
        attributes=_extract_model_info(kwargs),
    )

    # Capture system message
    system_msg = kwargs.get("system")
    if system_msg and trace_content:
        if isinstance(system_msg, str):
            span.set_attribute(
                "gen_ai.system_prompt",
                system_msg[:5000] if len(system_msg) > 5000 else system_msg,
            )

    # Capture input messages
    messages = kwargs.get("messages", [])
    _capture_messages(span, messages, trace_content)

    try:
        start_time = time.perf_counter()
        result = await wrapped(*args, **kwargs)
        return _wrap_stream_async(result, span, start_time, trace_content, tracer)
    except Exception as e:
        span.record_exception(e)
        span.set_attribute("error", True)
        span.set_attribute("error.type", type(e).__name__)
        span.set_attribute("error.message", scrub_sensitive(str(e)[:2000]))
        span.end()
        tracer.export_span(span)
        raise


def _wrap_stream_sync(
    stream: Iterator[Any],
    span: Any,
    start_time: float,
    trace_content: bool,
    tracer: Any,
) -> Iterator[Any]:
    """Wrap a sync streaming response."""
    event_count = 0
    content_buffer = []
    content_length = 0
    usage_data = {}

    try:
        for event in stream:
            event_count += 1
            event_type = getattr(event, "type", "")

            # Accumulate text content (bounded)
            if trace_content and content_length < 10000 and event_type == "content_block_delta":
                delta = getattr(event, "delta", None)
                if delta and getattr(delta, "type", "") == "text_delta":
                    text = getattr(delta, "text", "")
                    if text:
                        content_buffer.append(text)
                        content_length += len(text)

            # Capture usage from message_delta
            if event_type == "message_delta":
                usage = getattr(event, "usage", None)
                if usage:
                    usage_data["output_tokens"] = getattr(usage, "output_tokens", 0)

            # Capture usage from message_start
            if event_type == "message_start":
                message = getattr(event, "message", None)
                if message:
                    usage = getattr(message, "usage", None)
                    if usage:
                        usage_data["input_tokens"] = getattr(usage, "input_tokens", 0)

            yield event

    except Exception as e:
        span.record_exception(e)
        span.set_attribute("error", True)
        span.set_attribute("error.type", type(e).__name__)
        raise
    finally:
        latency_ms = (time.perf_counter() - start_time) * 1000
        span.set_attribute("gen_ai.latency_ms", latency_ms)
        span.set_attribute("gen_ai.response.events", event_count)
        span.set_attribute("gen_ai.response.stream", True)

        if usage_data:
            input_tokens = usage_data.get("input_tokens", 0)
            output_tokens = usage_data.get("output_tokens", 0)
            span.set_attribute("gen_ai.usage.prompt_tokens", input_tokens)
            span.set_attribute("gen_ai.usage.completion_tokens", output_tokens)
            span.set_attribute("gen_ai.usage.total_tokens", input_tokens + output_tokens)

        if trace_content and content_buffer:
            full_content = "".join(content_buffer)
            span.set_attribute(
                "gen_ai.completion.content",
                full_content[:10000] if len(full_content) > 10000 else full_content,
            )

        from risicare_core import SpanStatus
        if span.status == SpanStatus.UNSET:
            span.mark_ok()
        span.end()
        tracer.export_span(span)


async def _wrap_stream_async(
    stream: AsyncIterator[Any],
    span: Any,
    start_time: float,
    trace_content: bool,
    tracer: Any,
) -> AsyncIterator[Any]:
    """Wrap an async streaming response."""
    event_count = 0
    content_buffer = []
    content_length = 0
    usage_data = {}

    try:
        async for event in stream:
            event_count += 1
            event_type = getattr(event, "type", "")

            # Accumulate text content (bounded)
            if trace_content and content_length < 10000 and event_type == "content_block_delta":
                delta = getattr(event, "delta", None)
                if delta and getattr(delta, "type", "") == "text_delta":
                    text = getattr(delta, "text", "")
                    if text:
                        content_buffer.append(text)
                        content_length += len(text)

            # Capture usage from message_delta
            if event_type == "message_delta":
                usage = getattr(event, "usage", None)
                if usage:
                    usage_data["output_tokens"] = getattr(usage, "output_tokens", 0)

            # Capture usage from message_start
            if event_type == "message_start":
                message = getattr(event, "message", None)
                if message:
                    usage = getattr(message, "usage", None)
                    if usage:
                        usage_data["input_tokens"] = getattr(usage, "input_tokens", 0)

            yield event

    except Exception as e:
        span.record_exception(e)
        span.set_attribute("error", True)
        span.set_attribute("error.type", type(e).__name__)
        raise
    finally:
        latency_ms = (time.perf_counter() - start_time) * 1000
        span.set_attribute("gen_ai.latency_ms", latency_ms)
        span.set_attribute("gen_ai.response.events", event_count)
        span.set_attribute("gen_ai.response.stream", True)

        if usage_data:
            input_tokens = usage_data.get("input_tokens", 0)
            output_tokens = usage_data.get("output_tokens", 0)
            span.set_attribute("gen_ai.usage.prompt_tokens", input_tokens)
            span.set_attribute("gen_ai.usage.completion_tokens", output_tokens)
            span.set_attribute("gen_ai.usage.total_tokens", input_tokens + output_tokens)

        if trace_content and content_buffer:
            full_content = "".join(content_buffer)
            span.set_attribute(
                "gen_ai.completion.content",
                full_content[:10000] if len(full_content) > 10000 else full_content,
            )

        from risicare_core import SpanStatus
        if span.status == SpanStatus.UNSET:
            span.mark_ok()
        span.end()
        tracer.export_span(span)
