"""Host-local execution metadata for Cloud planner accounting."""

from __future__ import annotations

from contextlib import contextmanager
from contextvars import ContextVar
from dataclasses import dataclass
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from collections.abc import Iterator

    from cloud.internal_api.models import AssignmentModel


@dataclass(frozen=True)
class PlannerExecutionContext:
    task_id: str
    attempt: int
    lease_id: str


_context: ContextVar[PlannerExecutionContext | None] = ContextVar(
    "host_agent_planner_execution_context",
    default=None,
)


def current_planner_execution_context() -> PlannerExecutionContext | None:
    return _context.get()


@contextmanager
def bind_planner_execution_context(
    assignment: "AssignmentModel",
) -> "Iterator[None]":
    token = _context.set(
        PlannerExecutionContext(
            task_id=assignment.task_id,
            attempt=assignment.attempt,
            lease_id=assignment.lease_id,
        )
    )
    try:
        yield
    finally:
        _context.reset(token)