"""
Context management using Python's contextvars.

This allows automatic context propagation without manual passing of IDs.
"""
from contextvars import ContextVar
from typing import Optional

from .models import ExecutionRun, TaskRun, StepRun


_current_execution_run: ContextVar[Optional[ExecutionRun]] = ContextVar(
    "current_execution_run",
    default=None
)

_current_task_run: ContextVar[Optional[TaskRun]] = ContextVar(
    "current_task_run",
    default=None
)

_current_step_run: ContextVar[Optional[StepRun]] = ContextVar(
    "current_step_run",
    default=None
)


def get_current_execution_run() -> Optional[ExecutionRun]:
    """Get the current execution run from context.

    Returns:
        Optional[ExecutionRun]: Current execution run or None.
    """
    return _current_execution_run.get()


def set_current_execution_run(execution_run: Optional[ExecutionRun]) -> None:
    """Set the current execution run in context.

    Args:
        execution_run (Optional[ExecutionRun]): Execution run to set.
    """
    _current_execution_run.set(execution_run)


def get_current_task_run() -> Optional[TaskRun]:
    """Get the current task run from context.

    Returns:
        Optional[TaskRun]: Current task run or None.
    """
    return _current_task_run.get()


def set_current_task_run(task_run: Optional[TaskRun]) -> None:
    """Set the current task run in context.

    Args:
        task_run (Optional[TaskRun]): Task run to set.
    """
    _current_task_run.set(task_run)


def get_current_step_run() -> Optional[StepRun]:
    """Get the current step run from context.

    Returns:
        Optional[StepRun]: Current step run or None.
    """
    return _current_step_run.get()


def set_current_step_run(step_run: Optional[StepRun]) -> None:
    """Set the current step run in context.

    Args:
        step_run (Optional[StepRun]): Step run to set.
    """
    _current_step_run.set(step_run)
