"""Anthropic SDK monkey-patch for plyra-trace."""

from __future__ import annotations

import json
import logging
from typing import Any

import wrapt
from opentelemetry import trace as trace_api

from plyra_trace.semconv import SpanAttributes, SpanKind

logger = logging.getLogger("plyra_instrument_anthropic")

_originals: dict[str, Any] = {}


def _patch_anthropic(tracer_provider: Any, capture_content: bool, capture_usage: bool) -> None:
    try:
        import anthropic.resources.messages as msg_mod
    except ImportError as exc:
        raise ImportError("anthropic >= 0.20.0 required for plyra-instrument-anthropic") from exc

    _originals["messages_create"] = msg_mod.Messages.create
    _originals["async_messages_create"] = msg_mod.AsyncMessages.create

    if tracer_provider is not None:
        tracer = tracer_provider.get_tracer("plyra-instrument-anthropic")
    else:
        tracer = trace_api.get_tracer("plyra-instrument-anthropic")

    @wrapt.decorator
    def _wrap_messages_create(wrapped: Any, instance: Any, args: tuple, kwargs: dict) -> Any:
        model = kwargs.get("model", "claude")
        with tracer.start_as_current_span(f"llm.{model}") as span:
            _set_attrs(span, model, kwargs, capture_content)
            try:
                response = wrapped(*args, **kwargs)
            except Exception as e:
                span.set_status(trace_api.StatusCode.ERROR, str(e))
                raise
            _record_response(span, response, model, capture_content, capture_usage)
            return response

    @wrapt.decorator
    async def _wrap_async_messages_create(wrapped: Any, instance: Any, args: tuple, kwargs: dict) -> Any:
        model = kwargs.get("model", "claude")
        with tracer.start_as_current_span(f"llm.{model}") as span:
            _set_attrs(span, model, kwargs, capture_content)
            try:
                response = await wrapped(*args, **kwargs)
            except Exception as e:
                span.set_status(trace_api.StatusCode.ERROR, str(e))
                raise
            _record_response(span, response, model, capture_content, capture_usage)
            return response

    msg_mod.Messages.create = _wrap_messages_create(msg_mod.Messages.create)  # type: ignore[method-assign]
    msg_mod.AsyncMessages.create = _wrap_async_messages_create(msg_mod.AsyncMessages.create)  # type: ignore[method-assign]
    logger.debug("plyra-instrument-anthropic: patches applied")


def _unpatch_anthropic() -> None:
    try:
        import anthropic.resources.messages as msg_mod

        if "messages_create" in _originals:
            msg_mod.Messages.create = _originals.pop("messages_create")  # type: ignore[method-assign]
        if "async_messages_create" in _originals:
            msg_mod.AsyncMessages.create = _originals.pop("async_messages_create")  # type: ignore[method-assign]
    except ImportError:
        pass


def _set_attrs(span: Any, model: str, kwargs: dict, capture_content: bool) -> None:
    span.set_attribute(SpanAttributes.SPAN_KIND, SpanKind.LLM)
    span.set_attribute(SpanAttributes.OPENINFERENCE_SPAN_KIND, "LLM")
    span.set_attribute(SpanAttributes.LLM_MODEL_NAME, model)
    span.set_attribute(SpanAttributes.LLM_PROVIDER, "anthropic")
    span.set_attribute(SpanAttributes.LLM_SYSTEM, "anthropic")

    if capture_content:
        msgs = kwargs.get("messages", [])
        if msgs:
            span.set_attribute(SpanAttributes.INPUT_VALUE, json.dumps(msgs)[:8000])
            span.set_attribute(SpanAttributes.INPUT_MIME_TYPE, "application/json")
        if kwargs.get("system"):
            span.set_attribute("llm.system_prompt", str(kwargs["system"])[:2000])

    params = {k: v for k, v in kwargs.items() if k in ("temperature", "max_tokens", "top_p", "top_k") and v is not None}
    if params:
        span.set_attribute(SpanAttributes.LLM_INVOCATION_PARAMETERS, json.dumps(params))


def _record_response(span: Any, response: Any, model: str, capture_content: bool, capture_usage: bool) -> None:
    if capture_content and hasattr(response, "content") and response.content:
        text_parts = [b.text for b in response.content if hasattr(b, "text")]
        content = " ".join(text_parts)
        if content:
            span.set_attribute(SpanAttributes.OUTPUT_VALUE, content[:8000])
            span.set_attribute(SpanAttributes.OUTPUT_MIME_TYPE, "text/plain")

    if capture_usage and hasattr(response, "usage") and response.usage:
        usage = response.usage
        input_toks = getattr(usage, "input_tokens", 0) or 0
        output_toks = getattr(usage, "output_tokens", 0) or 0
        span.set_attribute(SpanAttributes.LLM_TOKEN_COUNT_PROMPT, input_toks)
        span.set_attribute(SpanAttributes.LLM_TOKEN_COUNT_COMPLETION, output_toks)
        span.set_attribute(SpanAttributes.LLM_TOKEN_COUNT_TOTAL, input_toks + output_toks)

        try:
            from plyra_trace.collector.cost_table import estimate_cost

            cost = estimate_cost(model, input_toks, output_toks)
            if cost is not None:
                span.set_attribute("llm.cost_usd", cost)
        except Exception:  # noqa: BLE001
            pass
