import asyncio
from enum import Enum


class SessionState(str, Enum):
    IDLE = "idle"
    USER_SPEAKING = "user_speaking"
    THINKING = "thinking"
    AVATAR_SPEAKING = "avatar_speaking"


class Arbitrator:
    def __init__(self) -> None:
        self.state = SessionState.IDLE
        self._lock = asyncio.Lock()
        self.cancel_avatar_event = asyncio.Event()

    async def on_speech_start(self) -> None:
        async with self._lock:
            if self.state == SessionState.AVATAR_SPEAKING:
                self.cancel_avatar_event.set()
            self.state = SessionState.USER_SPEAKING

    async def on_speech_end(self) -> None:
        async with self._lock:
            self.cancel_avatar_event.clear()
            self.state = SessionState.THINKING

    async def on_avatar_start(self) -> None:
        async with self._lock:
            self.state = SessionState.AVATAR_SPEAKING

    async def on_avatar_done(self) -> None:
        async with self._lock:
            self.state = SessionState.IDLE