from __future__ import annotations

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

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


def get_current_turn_id() -> Optional[str]:
    return _current_turn_id.get()


def set_current_turn_id(turn_id: Optional[str]) -> None:
    _current_turn_id.set(turn_id)


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