from contextvars import ContextVar
from dataclasses import dataclass


@dataclass
class SkyvernContext:
    request_id: str | None = None
    organization_id: str | None = None
    task_id: str | None = None
    workflow_id: str | None = None
    workflow_run_id: str | None = None
    max_steps_override: int | None = None

    def __repr__(self) -> str:
        return f"SkyvernContext(request_id={self.request_id}, organization_id={self.organization_id}, task_id={self.task_id}, workflow_id={self.workflow_id}, workflow_run_id={self.workflow_run_id}, max_steps_override={self.max_steps_override})"

    def __str__(self) -> str:
        return self.__repr__()


_context: ContextVar[SkyvernContext | None] = ContextVar(
    "Global context",
    default=None,
)


def current() -> SkyvernContext | None:
    """
    Get the current context

    Returns:
        The current context, or None if there is none
    """
    return _context.get()


def ensure_context() -> SkyvernContext:
    """
    Get the current context, or raise an error if there is none

    Returns:
        The current context if there is one

    Raises:
        RuntimeError: If there is no current context
    """
    context = current()
    if context is None:
        raise RuntimeError("No skyvern context")
    return context


def set(context: SkyvernContext) -> None:
    """
    Set the current context

    Args:
        context: The context to set

    Returns:
        None
    """
    _context.set(context)


def reset() -> None:
    """
    Reset the current context

    Returns:
        None
    """
    _context.set(None)
