"""LiteLLM 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_litellm")

_originals: dict[str, Any] = {}


def _patch_litellm(tracer_provider: Any, capture_content: bool, capture_usage: bool) -> None:
    try:
        import litellm
    except ImportError as exc:
        raise ImportError("litellm >= 1.0.0 required for plyra-instrument-litellm") from exc

    _originals["completion"] = litellm.completion
    _originals["acompletion"] = litellm.acompletion

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

    @wrapt.decorator
    def _wrap_completion(wrapped: Any, instance: Any, args: tuple, kwargs: dict) -> Any:
        model = kwargs.get("model", args[0] if args else "unknown")
        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_acompletion(wrapped: Any, instance: Any, args: tuple, kwargs: dict) -> Any:
        model = kwargs.get("model", args[0] if args else "unknown")
        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

    litellm.completion = _wrap_completion(litellm.completion)
    litellm.acompletion = _wrap_acompletion(litellm.acompletion)
    logger.debug("plyra-instrument-litellm: patches applied")


def _unpatch_litellm() -> None:
    try:
        import litellm

        if "completion" in _originals:
            litellm.completion = _originals.pop("completion")
        if "acompletion" in _originals:
            litellm.acompletion = _originals.pop("acompletion")
    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)
    # Provider from model prefix (e.g., "anthropic/claude-3-sonnet")
    if "/" in model:
        provider = model.split("/")[0]
        span.set_attribute(SpanAttributes.LLM_PROVIDER, provider)
    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")


def _record_response(span: Any, response: Any, model: str, capture_content: bool, capture_usage: bool) -> None:
    if capture_content and hasattr(response, "choices") and response.choices:
        content = response.choices[0].message.content or ""
        if content:
            span.set_attribute(SpanAttributes.OUTPUT_VALUE, content[:8000])

    if capture_usage and hasattr(response, "usage") and response.usage:
        usage = response.usage
        prompt_toks = getattr(usage, "prompt_tokens", 0) or 0
        completion_toks = getattr(usage, "completion_tokens", 0) or 0
        span.set_attribute(SpanAttributes.LLM_TOKEN_COUNT_PROMPT, prompt_toks)
        span.set_attribute(SpanAttributes.LLM_TOKEN_COUNT_COMPLETION, completion_toks)
        span.set_attribute(SpanAttributes.LLM_TOKEN_COUNT_TOTAL, prompt_toks + completion_toks)

        # LiteLLM provides cost directly
        cost = getattr(response, "_hidden_params", {}).get("response_cost")
        if cost is not None:
            span.set_attribute("llm.cost_usd", float(cost))
        else:
            try:
                from plyra_trace.collector.cost_table import estimate_cost

                estimated = estimate_cost(model, prompt_toks, completion_toks)
                if estimated is not None:
                    span.set_attribute("llm.cost_usd", estimated)
            except Exception:  # noqa: BLE001
                pass
