"""plyra-instrument-litellm: LiteLLM instrumentation for plyra-trace."""

from __future__ import annotations

from typing import Any

_instrumented = False


def instrument(
    tracer_provider: Any = None,
    capture_content: bool = True,
    capture_usage: bool = True,
) -> None:
    """
    Instrument LiteLLM to emit plyra-trace spans.

    Works for:
    - litellm.completion() (sync)
    - litellm.acompletion() (async)
    - ALL providers through LiteLLM: openai, anthropic, groq, mistral, etc.

    Usage:
        import plyra_trace
        from plyra_instrument_litellm import instrument

        pt = plyra_trace.init(project="my-agent", endpoint="http://localhost:7700")
        instrument()
    """
    global _instrumented
    if _instrumented:
        return
    from plyra_instrument_litellm._patch import _patch_litellm

    _patch_litellm(tracer_provider, capture_content, capture_usage)
    _instrumented = True


def uninstrument() -> None:
    """Remove LiteLLM instrumentation patches."""
    global _instrumented
    if not _instrumented:
        return
    from plyra_instrument_litellm._patch import _unpatch_litellm

    _unpatch_litellm()
    _instrumented = False


__all__ = ["instrument", "uninstrument"]
