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

import functools
import json
import logging
from typing import Any

from openinference.semconv.trace import OpenInferenceSpanKindValues
from openinference.semconv.trace import SpanAttributes as OI
from opentelemetry import trace
from opentelemetry.trace import StatusCode

logger = logging.getLogger("staso")

_patched = False


def _truncate(text: Any, max_len: int = 500) -> str:
    s = str(text) if text else ""
    return s[:max_len] + ("..." if len(s) > max_len else "")


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


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

    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)
    _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 _wrap_sync(cls: type) -> None:
    original = cls.create

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

        with tracer.start_as_current_span(
            "anthropic.messages.create",
            attributes={
                OI.OPENINFERENCE_SPAN_KIND: OpenInferenceSpanKindValues.LLM.value,
                OI.LLM_MODEL_NAME: model,
                OI.LLM_PROVIDER: "anthropic",
                OI.INPUT_VALUE: _safe_json({
                    "model": model,
                    "max_tokens": kwargs.get("max_tokens"),
                    "temperature": kwargs.get("temperature"),
                    "system": _truncate(kwargs.get("system", "")),
                }),
            },
        ) as span:
            try:
                response = original(self, *args, **kwargs)
                usage = getattr(response, "usage", None)
                input_tok = getattr(usage, "input_tokens", 0) if usage else 0
                output_tok = getattr(usage, "output_tokens", 0) if usage else 0
                span.set_attribute(OI.LLM_TOKEN_COUNT_PROMPT, input_tok)
                span.set_attribute(OI.LLM_TOKEN_COUNT_COMPLETION, output_tok)
                span.set_attribute(OI.LLM_TOKEN_COUNT_TOTAL, input_tok + output_tok)
                span.set_attribute(OI.OUTPUT_VALUE, _safe_json({
                    "stop_reason": getattr(response, "stop_reason", None),
                }))
                span.set_status(StatusCode.OK)
                return response
            except Exception as exc:
                span.set_status(StatusCode.ERROR, str(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 = trace.get_tracer("staso")
        model = _extract_model(args, kwargs)

        with tracer.start_as_current_span(
            "anthropic.messages.create",
            attributes={
                OI.OPENINFERENCE_SPAN_KIND: OpenInferenceSpanKindValues.LLM.value,
                OI.LLM_MODEL_NAME: model,
                OI.LLM_PROVIDER: "anthropic",
                OI.INPUT_VALUE: _safe_json({
                    "model": model,
                    "max_tokens": kwargs.get("max_tokens"),
                    "temperature": kwargs.get("temperature"),
                    "system": _truncate(kwargs.get("system", "")),
                }),
            },
        ) as span:
            try:
                response = await original(self, *args, **kwargs)
                usage = getattr(response, "usage", None)
                input_tok = getattr(usage, "input_tokens", 0) if usage else 0
                output_tok = getattr(usage, "output_tokens", 0) if usage else 0
                span.set_attribute(OI.LLM_TOKEN_COUNT_PROMPT, input_tok)
                span.set_attribute(OI.LLM_TOKEN_COUNT_COMPLETION, output_tok)
                span.set_attribute(OI.LLM_TOKEN_COUNT_TOTAL, input_tok + output_tok)
                span.set_attribute(OI.OUTPUT_VALUE, _safe_json({
                    "stop_reason": getattr(response, "stop_reason", None),
                }))
                span.set_status(StatusCode.OK)
                return response
            except Exception as exc:
                span.set_status(StatusCode.ERROR, str(exc))
                span.record_exception(exc)
                raise

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