"""patch_litellm() — automatic trace capture for litellm.completion / acompletion."""

from __future__ import annotations

import functools
import logging
import time
from typing import Any

from carrot_ai._context import current_trace_id
from carrot_ai._wrap import (
    _AsyncStreamProxy,
    _OpenAIAccumulator,
    _SyncStreamProxy,
    _format_time,
    _normalize_openai,
    _safe_serialize,
)

logger = logging.getLogger("carrot_ai")

_patched = False


def patch_litellm(
    *,
    name: str | None = None,
    tags: list[str] | None = None,
) -> None:
    """Monkey-patch ``litellm.completion`` and ``litellm.acompletion``.

    After calling this, every ``litellm.completion()`` and
    ``litellm.acompletion()`` call is automatically traced — no other
    code changes required.

    Args:
        name: Override the default trace name (``"chat.completions"``).
        tags: Tags attached to every trace.

    Usage::

        import carrot_ai
        import litellm

        carrot_ai.init()
        carrot_ai.patch_litellm(name="my-service")

        response = litellm.completion(
            model="openai/gpt-4o",
            messages=[{"role": "user", "content": "Hello"}],
        )
    """
    global _patched
    if _patched:
        return

    try:
        import litellm  # noqa: F811
    except ImportError as exc:
        raise ImportError(
            "litellm is not installed. Install it with: pip install litellm"
        ) from exc

    extra: dict[str, Any] = {}
    if name:
        extra["name"] = name
    if tags:
        extra["tags"] = tags

    _patch(litellm, "completion", is_async=False, extra=extra)
    _patch(litellm, "acompletion", is_async=True, extra=extra)
    _patched = True


# ---------------------------------------------------------------------------
# Internal patching
# ---------------------------------------------------------------------------


def _patch(
    module: Any, attr: str, *, is_async: bool, extra: dict[str, Any]
) -> None:
    original = getattr(module, attr)

    if is_async:

        @functools.wraps(original)
        async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
            return await _trace_async(original, extra, args, kwargs)

        setattr(module, attr, async_wrapper)
    else:

        @functools.wraps(original)
        def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
            return _trace_sync(original, extra, args, kwargs)

        setattr(module, attr, sync_wrapper)


# ---------------------------------------------------------------------------
# Sync / async tracing
# ---------------------------------------------------------------------------


def _trace_sync(
    fn: Any,
    extra: dict[str, Any],
    args: tuple[Any, ...],
    kwargs: dict[str, Any],
) -> Any:
    started_at = time.time()
    is_stream = kwargs.get("stream", False)
    trace_input = _capture_input(args, kwargs)

    if is_stream:
        kwargs = dict(kwargs)
        opts = dict(kwargs.get("stream_options") or {})
        opts["include_usage"] = True
        kwargs["stream_options"] = opts

    try:
        result = fn(*args, **kwargs)
    except Exception as exc:
        _submit_trace(
            extra, trace_input,
            {"error": str(exc), "error_type": type(exc).__name__},
            started_at, time.time(), "error",
        )
        raise

    if is_stream:
        accumulator = _OpenAIAccumulator()

        def on_done() -> None:
            _submit_trace(
                extra, trace_input, accumulator.result(),
                started_at, time.time(), "success",
            )

        return _SyncStreamProxy(result, accumulator, on_done)

    output = _normalize_openai(result)
    _submit_trace(extra, trace_input, output, started_at, time.time(), "success")
    return result


async def _trace_async(
    fn: Any,
    extra: dict[str, Any],
    args: tuple[Any, ...],
    kwargs: dict[str, Any],
) -> Any:
    started_at = time.time()
    is_stream = kwargs.get("stream", False)
    trace_input = _capture_input(args, kwargs)

    if is_stream:
        kwargs = dict(kwargs)
        opts = dict(kwargs.get("stream_options") or {})
        opts["include_usage"] = True
        kwargs["stream_options"] = opts

    try:
        result = await fn(*args, **kwargs)
    except Exception as exc:
        _submit_trace(
            extra, trace_input,
            {"error": str(exc), "error_type": type(exc).__name__},
            started_at, time.time(), "error",
        )
        raise

    if is_stream:
        accumulator = _OpenAIAccumulator()

        def on_done() -> None:
            _submit_trace(
                extra, trace_input, accumulator.result(),
                started_at, time.time(), "success",
            )

        return _AsyncStreamProxy(result, accumulator, on_done)

    output = _normalize_openai(result)
    _submit_trace(extra, trace_input, output, started_at, time.time(), "success")
    return result


# ---------------------------------------------------------------------------
# Input capture — handles litellm's positional (model, messages) + kwargs
# ---------------------------------------------------------------------------

_POS_PARAMS = ("model", "messages")


def _capture_input(
    args: tuple[Any, ...], kwargs: dict[str, Any]
) -> dict[str, Any]:
    captured: dict[str, Any] = {}
    for i, name in enumerate(_POS_PARAMS):
        if name in kwargs:
            captured[name] = _safe_serialize(kwargs[name])
        elif i < len(args):
            captured[name] = _safe_serialize(args[i])

    for key, value in kwargs.items():
        if key in ("model", "messages", "stream", "stream_options", "api_key"):
            continue
        captured[key] = _safe_serialize(value)
    return captured


# ---------------------------------------------------------------------------
# Trace submission
# ---------------------------------------------------------------------------


def _submit_trace(
    extra: dict[str, Any],
    trace_input: dict[str, Any],
    output_data: dict[str, Any],
    started_at: float,
    ended_at: float,
    status: str,
) -> None:
    from carrot_ai import get_client

    parent_id = current_trace_id.get()

    metadata: dict[str, Any] = {
        "provider": "litellm",
        "source": "sdk",
        "latency_ms": round((ended_at - started_at) * 1000),
    }

    usage = output_data.get("usage")
    if usage:
        metadata.update(usage)

    model = output_data.get("model") or trace_input.get("model")
    if model:
        metadata["model"] = model

    if parent_id:
        metadata["parent_trace_id"] = parent_id

    output_payload: dict[str, Any] = {}
    for key in ("message", "embeddings_count", "error", "error_type"):
        if key in output_data:
            output_payload[key] = output_data[key]

    trace: dict[str, Any] = {
        "name": extra.get("name", "chat.completions"),
        "input": trace_input,
        "output": output_payload,
        "metadata": metadata,
        "status": status,
        "started_at": _format_time(started_at),
        "ended_at": _format_time(ended_at),
    }

    tags = extra.get("tags")
    if tags:
        trace["tags"] = list(tags)

    try:
        get_client().submit(trace)
    except Exception:
        logger.debug("carrot_ai: failed to submit litellm trace", exc_info=True)
