from __future__ import annotations
from dataclasses import dataclass, field
from typing import Any, Dict
import time
import uuid

@dataclass
class RunState:
    steps: int = 0
    tool_calls: int = 0
    cost_usd: float = 0.0

@dataclass
class SentinelContext:
    correlation_id: str = field(default_factory=lambda: str(uuid.uuid4()))
    run_id: str = field(default_factory=lambda: str(uuid.uuid4()))
    started_at_epoch_ms: int = field(default_factory=lambda: int(time.time() * 1000))
    metadata: Dict[str, Any] = field(default_factory=dict)
    state: RunState = field(default_factory=RunState)

    def bump_step(self) -> None:
        self.state.steps += 1

    def bump_tool_call(self) -> None:
        self.state.tool_calls += 1

    def add_cost(self, amount_usd: float) -> None:
        self.state.cost_usd += float(amount_usd)

    def snapshot(self) -> Dict[str, Any]:
        return {
            "correlation_id": self.correlation_id,
            "run_id": self.run_id,
            "started_at_epoch_ms": self.started_at_epoch_ms,
            "metadata": self.metadata,
            "state": {
                "steps": self.state.steps,
                "tool_calls": self.state.tool_calls,
                "cost_usd": self.state.cost_usd,
            },
        }
