from __future__ import annotations

from contextlib import contextmanager
from contextvars import ContextVar
from typing import Iterator


_current_turn_id: ContextVar[str | None] = ContextVar("current_turn_id", default=None)


def set_current_turn_id(turn_id: str | None) -> None:
    _current_turn_id.set(turn_id)


def get_current_turn_id() -> str | None:
    return _current_turn_id.get()


@contextmanager
def use_turn_id(turn_id: str | None) -> Iterator[None]:
    token = _current_turn_id.set(turn_id)
    try:
        yield
    finally:
        _current_turn_id.reset(token)
