"""Correlation ID context variable for distributed tracing.

This module provides a :class:`contextvars.ContextVar` that stores the correlation
ID for the current execution context. Background tasks inherit or explicitly set
the correlation ID to enable log correlation across request/task boundaries.

Usage in background tasks::

    from app.utils.correlation import get_correlation_id, set_correlation_id

    async def my_background_task(correlation_id: str) -> None:
        token = set_correlation_id(correlation_id)
        try:
            log.info("task_started")  # Logs include correlation_id
        finally:
            reset_correlation_id(token)

For APScheduler jobs, pass the correlation ID through kwargs and call
:func:`set_correlation_id` at the start of the async callback.
"""

from __future__ import annotations

from contextvars import ContextVar
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from contextvars import Token

#: Default correlation ID used when no parent request context exists.
_DEFAULT_BG_CORRELATION_ID: str = "bg-task"

#: ContextVar holding the current correlation ID.
correlation_id_var: ContextVar[str] = ContextVar(
    "correlation_id",
    default=_DEFAULT_BG_CORRELATION_ID,
)


def get_correlation_id() -> str:
    """Return the current correlation ID from context.

    Returns:
        The current correlation ID string. Returns ``"bg-task"`` if no
        correlation ID has been set in the current context.
    """
    return correlation_id_var.get()


def set_correlation_id(correlation_id: str) -> Token[str]:
    """Set the correlation ID for the current context.

    Args:
        correlation_id: The correlation ID to bind (typically passed from
            the triggering request).

    Returns:
        A :class:`contextvars.Token` that must be passed to
        :func:`reset_correlation_id` when the task completes.
    """
    return correlation_id_var.set(correlation_id)


def reset_correlation_id(token: Token[str]) -> None:
    """Reset the correlation ID to its previous value.

    Args:
        token: The token returned by :func:`set_correlation_id`.
    """
    correlation_id_var.reset(token)