from __future__ import annotations

from contextvars import ContextVar, Token


_job_id_var: ContextVar[str | None] = ContextVar("connecthub_job_id", default=None)
_job_log_id_var: ContextVar[int | None] = ContextVar("connecthub_job_log_id", default=None)


JobContextTokens = tuple[Token[str | None], Token[int | None]]


def set_job_context(*, job_id: str | None, job_log_id: int | None) -> JobContextTokens:
    """
    设置当前执行上下文（用于日志隔离）。
    返回 tokens，便于在 finally 中 reset 回原值。
    """
    t1 = _job_id_var.set(job_id)
    t2 = _job_log_id_var.set(job_log_id)
    return (t1, t2)


def clear_job_context(tokens: JobContextTokens | None = None) -> None:
    """
    清理当前执行上下文。
    - 若提供 tokens：reset 回进入上下文前的值（推荐）
    - 否则：直接置空
    """
    if tokens is not None:
        _job_id_var.reset(tokens[0])
        _job_log_id_var.reset(tokens[1])
        return
    _job_id_var.set(None)
    _job_log_id_var.set(None)


def get_job_id() -> str | None:
    return _job_id_var.get()


def get_job_log_id() -> int | None:
    return _job_log_id_var.get()