from __future__ import annotations
import asyncio
from dataclasses import dataclass
from typing import Callable, Awaitable, Optional, Dict

@dataclass
class PhaseState:
    name: str
    started_at: float

PhaseHook = Callable[[str], Awaitable[None]]  # on_phase_change(new_phase)

class PhaseMachine:
    """
    Phase transitions with safe cancellation.
    Similar to your "phase task" approach.
    """
    def __init__(self, initial: str = "lobby"):
        self.phase = initial
        self._task: Optional[asyncio.Task] = None
        self._lock = asyncio.Lock()

    async def set(self, phase: str, on_change: Optional[PhaseHook] = None):
        async with self._lock:
            self.phase = phase
            if on_change:
                await on_change(phase)

    async def schedule_transition(self, seconds: float, next_phase: str, on_change: Optional[PhaseHook] = None):
        async with self._lock:
            if self._task and not self._task.done():
                self._task.cancel()
            self._task = asyncio.create_task(self._run_timer(seconds, next_phase, on_change))

    async def cancel(self):
        async with self._lock:
            if self._task and not self._task.done():
                self._task.cancel()
            self._task = None

    async def _run_timer(self, seconds: float, next_phase: str, on_change: Optional[PhaseHook]):
        try:
            await asyncio.sleep(max(0.0, seconds))
            await self.set(next_phase, on_change=on_change)
        except asyncio.CancelledError:
            return