from __future__ import annotations

from contextlib import contextmanager
from contextvars import ContextVar, Token
from typing import TYPE_CHECKING, Generator, Protocol, TypeVar

if TYPE_CHECKING:
    from contextlib import _GeneratorContextManager

_T = TypeVar("_T")


class ActivateContextVarFn(Protocol[_T]):
    """Get a context manager that sets `value` on a `ContextVar` while active."""

    def __call__(self, value: _T | None = None) -> _GeneratorContextManager[_T]:
        pass


def contextvar_activator(contextvar: ContextVar[_T]) -> ActivateContextVarFn[_T]:
    """
    Create a function that sets a `ContextVar`'s value within a `with` block.

    When called with an existing `ContextVar`, this function returns another
    function that itself takes a value to set on the `ContextVar` and returns a
    context manager that will activate and deactivate the value on the
    `ContextVar`.

    Example
    -------
    >>> my_context = ContextVar('my_context', default='Not Set')
    >>> activate_my_context = contextvar_activator(my_context)

    >>> my_context.get()
    'Not Set'

    >>> with activate_my_context('Hello World!') as active_value:
    ...     assert active_value == 'Hello World!'

    >>> my_context.get()
    'Not Set'
    """

    @contextmanager
    def activate_context(value: _T | None = None) -> Generator[_T]:
        token: Token[_T] | None = None
        if value is not None:
            token = contextvar.set(value)
        try:
            yield contextvar.get()
        finally:
            if token is not None:
                contextvar.reset(token)

    activate_context.__name__ = f"activate_{contextvar.name.lower()}"
    return activate_context
