from __future__ import annotations

import uuid
from contextvars import ContextVar
from typing import Optional

_trace_id: ContextVar[Optional[str]] = ContextVar("_trace_id", default=None)
_step_index: ContextVar[int] = ContextVar("_step_index", default=0)
_budget_limit: ContextVar[Optional[float]] = ContextVar("_budget_limit", default=None)
_budget_spent: ContextVar[float] = ContextVar("_budget_spent", default=0.0)


def get_trace_id() -> Optional[str]:
    return _trace_id.get()


def get_step_index() -> int:
    return _step_index.get()


def increment_step() -> int:
    current = _step_index.get()
    _step_index.set(current + 1)
    return current


def get_budget_limit() -> Optional[float]:
    return _budget_limit.get()


def get_budget_spent() -> float:
    return _budget_spent.get()


def add_budget_spent(amount: float) -> None:
    _budget_spent.set(_budget_spent.get() + amount)


class trace:
    """Context manager that groups all LLM calls under a single trace ID.

    Usage::

        with art.trace("user-session-abc"):
            response1 = client.chat.completions.create(...)
            response2 = client.chat.completions.create(...)
    """

    def __init__(self, trace_id: Optional[str] = None) -> None:
        self.trace_id = trace_id or str(uuid.uuid4())
        self._tokens: list = []

    def __enter__(self) -> "trace":
        self._tokens = [
            _trace_id.set(self.trace_id),
            _step_index.set(0),
        ]
        return self

    def __exit__(self, *_) -> None:
        for token in self._tokens:
            token.var.reset(token)


class cost_guard:
    """Context manager that raises BudgetExceededError if a request would exceed budget.

    Usage::

        with art.cost_guard(max_usd=0.05):
            response = client.chat.completions.create(...)
    """

    def __init__(self, max_usd: float) -> None:
        self.max_usd = max_usd
        self._tokens: list = []

    def __enter__(self) -> "cost_guard":
        self._tokens = [
            _budget_limit.set(self.max_usd),
            _budget_spent.set(0.0),
        ]
        return self

    def __exit__(self, *_) -> None:
        for token in self._tokens:
            token.var.reset(token)
