"""Native context propagation for span parent-child linking."""
from __future__ import annotations

from contextvars import ContextVar, Token
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from .provider import Span

_current_span: ContextVar[Span | None] = ContextVar("_staso_span", default=None)
session_id_var: ContextVar[str] = ContextVar("_staso_session", default="")
user_id_var: ContextVar[str] = ContextVar("_staso_user_id", default="")


def get_current_span() -> Span | None:
    """Return the active span in the current context, or None."""
    return _current_span.get()


def set_current_span(span: Span) -> Token:
    """Set *span* as the active span. Returns a token to restore the previous value."""
    return _current_span.set(span)


def reset_current_span(token: Token) -> None:
    """Restore the previous span from *token*."""
    _current_span.reset(token)
