"""Carrot AI Python SDK — automatic tracing for LLM calls."""

from __future__ import annotations

import os
from typing import Any

from carrot_ai._client import CarrotClient
from carrot_ai._litellm import patch_litellm
from carrot_ai._trace import trace
from carrot_ai._wrap import wrap

__all__ = ["init", "wrap", "trace", "patch_litellm", "get_client", "flush"]
__version__ = "0.2.1"

_client: CarrotClient | None = None

_DEFAULT_BASE_URL = "https://api.carrotlabs.ai"


def init(
    *,
    api_key: str | None = None,
    base_url: str | None = None,
    flush_interval: float = 5.0,
    batch_size: int = 50,
) -> None:
    """Initialize the Carrot SDK.

    Must be called before ``wrap()`` or ``@trace`` capture any data.
    Alternatively, set the ``CARROT_API_KEY`` environment variable for
    auto-initialization on first use.

    Args:
        api_key: Your Carrot API key (``sk-...``).  Falls back to
            ``CARROT_API_KEY`` env var.
        base_url: Override the API endpoint.  Falls back to
            ``CARROT_BASE_URL`` env var, then ``https://api.carrotlabs.ai``.
        flush_interval: Seconds between background flushes (default 5).
        batch_size: Max traces per batch POST (default 50).
    """
    global _client

    if api_key is None:
        api_key = os.environ.get("CARROT_API_KEY")
    if api_key is None:
        raise ValueError(
            "api_key is required — pass it to carrot_ai.init() "
            "or set the CARROT_API_KEY environment variable"
        )
    if base_url is None:
        base_url = os.environ.get("CARROT_BASE_URL", _DEFAULT_BASE_URL)

    _client = CarrotClient(
        api_key=api_key,
        base_url=base_url,
        flush_interval=flush_interval,
        batch_size=batch_size,
    )


def get_client() -> CarrotClient:
    """Return the global ``CarrotClient``, auto-initializing if possible."""
    global _client
    if _client is None:
        api_key = os.environ.get("CARROT_API_KEY")
        if api_key:
            init(api_key=api_key)
        else:
            raise RuntimeError(
                "carrot_ai.init() has not been called and CARROT_API_KEY is not set"
            )
    return _client


def flush() -> None:
    """Flush any pending traces to the server immediately."""
    if _client is not None:
        _client.flush()
