from __future__ import annotations

from contextvars import ContextVar, Token
from dataclasses import dataclass
from typing import Any

_CURRENT_CONTEXT: ContextVar["AgentRunContext | None"] = ContextVar("agent_run_context", default=None)


@dataclass(slots=True)
class AgentRunContext:
    thread_id: str
    subagent_mode: bool
    research_mode: bool
    has_media_attachment: bool
    include_cli: bool
    quick_panel: bool
    dashboard_surface: bool
    force_plain_chat: bool
    policy_context: Any | None = None
    iteration_budget: Any | None = None


def set_current_run_context(ctx: AgentRunContext | None) -> Token:
    return _CURRENT_CONTEXT.set(ctx)


def reset_current_run_context(token: Token) -> None:
    _CURRENT_CONTEXT.reset(token)


def get_current_run_context() -> AgentRunContext | None:
    return _CURRENT_CONTEXT.get()

