from __future__ import annotations
from dataclasses import dataclass, field
import unittest

TERMINAL = {"SUCCEEDED", "FAILED", "CANCELLED"}
ALLOWED = {
    "CREATED": {"RUNNING", "CANCELLED"},
    "RUNNING": {"WAITING_APPROVAL", "SUCCEEDED", "FAILED", "CANCELLED"},
    "WAITING_APPROVAL": {"RUNNING", "CANCELLED"},
    "SUCCEEDED": set(), "FAILED": set(), "CANCELLED": set(),
}

@dataclass
class Run:
    status: str = "CREATED"
    steps: int = 0
    max_steps: int = 3
    completed_actions: set = field(default_factory=set)
    pending_action: str | None = None
    events: list = field(default_factory=list)

    def transition(self, new_status):
        if new_status not in ALLOWED[self.status]:
            raise ValueError(f"invalid transition: {self.status} -> {new_status}")
        self.events.append(("TRANSITION", self.status, new_status))
        self.status = new_status

    def propose(self, action_id, requires_approval=False):
        if self.status != "RUNNING":
            raise RuntimeError("run is not executing")
        if self.steps >= self.max_steps:
            self.transition("FAILED")
            return "BUDGET_EXHAUSTED"
        self.steps += 1
        if action_id in self.completed_actions:
            self.events.append(("DEDUPLICATED", action_id))
            return "ALREADY_RECORDED"
        if requires_approval:
            self.pending_action = action_id
            self.transition("WAITING_APPROVAL")
            return "APPROVAL_REQUIRED"
        self.completed_actions.add(action_id)
        self.events.append(("ACTION_RECORDED", action_id))
        return "RECORDED"

    def approve(self, approved):
        if self.status != "WAITING_APPROVAL" or not self.pending_action:
            raise RuntimeError("no pending approval")
        action_id = self.pending_action
        self.pending_action = None
        if not approved:
            self.transition("CANCELLED")
            return "REJECTED"
        self.transition("RUNNING")
        self.completed_actions.add(action_id)
        self.events.append(("ACTION_RECORDED", action_id))
        return "APPROVED_AND_RECORDED"

class ExecutionTests(unittest.TestCase):
    def test_approval_pause_resume(self):
        r = Run(); r.transition("RUNNING")
        self.assertEqual(r.propose("send:1", True), "APPROVAL_REQUIRED")
        self.assertEqual(r.status, "WAITING_APPROVAL")
        self.assertEqual(r.approve(True), "APPROVED_AND_RECORDED")
        self.assertIn("send:1", r.completed_actions)

    def test_rejection(self):
        r = Run(); r.transition("RUNNING"); r.propose("write:1", True)
        self.assertEqual(r.approve(False), "REJECTED")
        self.assertEqual(r.status, "CANCELLED")

    def test_deduplication(self):
        r = Run(); r.transition("RUNNING")
        self.assertEqual(r.propose("search:1"), "RECORDED")
        self.assertEqual(r.propose("search:1"), "ALREADY_RECORDED")

    def test_budget_and_terminal(self):
        r = Run(max_steps=1); r.transition("RUNNING")
        r.propose("search:1")
        self.assertEqual(r.propose("search:2"), "BUDGET_EXHAUSTED")
        self.assertEqual(r.status, "FAILED")
        with self.assertRaises(ValueError):
            r.transition("RUNNING")

if __name__ == "__main__":
    unittest.main(verbosity=2)