from __future__ import annotations

from contextvars import ContextVar, Token
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from agenx.trace.spans.span import LiveSpan

_current_span: ContextVar["LiveSpan | None"] = ContextVar(
    "_current_span", default=None
)

def get_current_span() -> "LiveSpan | None":
    return _current_span.get()

def push_span(live_span: "LiveSpan") -> "Token[LiveSpan | None]":
    return _current_span.set(live_span)

def pop_span(token: "Token[LiveSpan | None]") -> None:
    return _current_span.reset(token)
