"""Per-run execution context shared across agent subsystems."""
from __future__ import annotations

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


@dataclass
class RunContext:
    """Execution context for the current agent run."""

    trace_id: str | None = None
    request_id: str | None = None
    session_id: str | None = None
    workspace_id: str | None = None
    task_type: str | None = None
    user_segment: str | None = None
    prompt_text: str | None = None
    prompt_name: str | None = None
    prompt_version: str | None = None
    prompt_variant: str | None = None
    experiment_id: str | None = None
    routing_hint: str | None = None
    criticality: str | None = None
    budget_preference: str | None = None
    fail_on_budget_exhausted: bool = False
    task_budget: Any | None = None
    team_budget: Any | None = None
    pipeline_ctx: Any | None = None


_current_run_context: ContextVar[RunContext | None] = ContextVar("logos_adk_run_context", default=None)


def get_run_context() -> RunContext | None:
    """Return the active per-run context if present."""
    return _current_run_context.get()


def set_run_context(context: RunContext) -> Token:
    """Install a run context and return the reset token."""
    return _current_run_context.set(context)


def reset_run_context(token: Token) -> None:
    """Reset a previously installed run context."""
    _current_run_context.reset(token)


@contextmanager
def use_run_context(context: RunContext):
    """Install a per-run context for the current coroutine/task."""
    token: Token = _current_run_context.set(context)
    try:
        yield context
    finally:
        _current_run_context.reset(token)


__all__ = [
    "RunContext",
    "get_run_context",
    "set_run_context",
    "reset_run_context",
    "use_run_context",
]
