import contextlib
import contextvars
import os

namespace_context_var = contextvars.ContextVar('namespace_context_var', default=None)


@contextlib.contextmanager
def namespace_provider(namespace: str):
    # noinspection PyTypeChecker
    token = namespace_context_var.set(namespace)
    yield
    namespace_context_var.reset(token)


@contextlib.contextmanager
def namespace_context():
    # noinspection PyTypeChecker
    namespace = namespace_context_var.get()
    if namespace is None:
        namespace = os.environ.get("EVENTIX_NAMESPACE", "default")
    yield namespace


delay_tasks_context_var = contextvars.ContextVar('delay_tasks_context_var', default=True)


@contextlib.contextmanager
def delay_tasks(delay_tasks: bool = True):
    # noinspection PyTypeChecker
    token = delay_tasks_context_var.set(delay_tasks)
    yield
    delay_tasks_context_var.reset(token)


@contextlib.contextmanager
def delay_tasks_context():
    # noinspection PyTypeChecker
    delay_tasks = delay_tasks_context_var.get()
    yield delay_tasks


worker_id_context_var = contextvars.ContextVar('worker_id_context_var', default=None)


@contextlib.contextmanager
def worker_id_provider(worker_id: str):
    # noinspection PyTypeChecker
    token = worker_id_context_var.set(worker_id)
    yield
    worker_id_context_var.reset(token)


@contextlib.contextmanager
def worker_id_context():
    # noinspection PyTypeChecker
    worker_id = worker_id_context_var.get()
    if worker_id is None:
        worker_id = "default"
    yield worker_id
