audio_engine/priority_manager.py

from typing import Dict, Optional
from dataclasses import dataclass, field
from datetime import datetime, timezone


@dataclass
class MicrophoneState:
    mic_id: str
    name: str
    priority: int = 5  # 1–10
    is_chairman: bool = False
    is_muted: bool = False
    is_speech_active: bool = False
    db_level: float = -100.0
    speech_probability: float = 0.0
    last_speech_time: Optional[datetime] = field(default=None)
    gain: float = 1.0


class PriorityManager:
    """
    Manages microphone priorities and speech-collision prevention.

    When multiple microphones detect speech simultaneously:
    - Chairman microphone always wins.
    - Otherwise the highest-priority (then loudest) microphone wins and all
      others are flagged for suppression.
    """

    def __init__(self):
        self._microphones: Dict[str, MicrophoneState] = {}
        self._active_speaker: Optional[str] = None

    # ------------------------------------------------------------------
    # Microphone registration
    # ------------------------------------------------------------------

    def register_microphone(
        self,
        mic_id: str,
        name: str,
        priority: int = 5,
        is_chairman: bool = False,
    ) -> MicrophoneState:
        state = MicrophoneState(
            mic_id=mic_id,
            name=name,
            priority=priority,
            is_chairman=is_chairman,
        )
        self._microphones[mic_id] = state
        return state

    def remove_microphone(self, mic_id: str) -> None:
        self._microphones.pop(mic_id, None)
        if self._active_speaker == mic_id:
            self._active_speaker = None

    # ------------------------------------------------------------------
    # State updates
    # ------------------------------------------------------------------

    def update_state(self, mic_id: str, vad_result: dict, gain: float) -> None:
        """Apply a VAD result to the stored microphone state."""
        if mic_id not in self._microphones:
            return
        mic = self._microphones[mic_id]
        mic.db_level = vad_result.get("db_level", -100.0)
        mic.speech_probability = vad_result.get("speech_probability", 0.0)
        mic.is_speech_active = vad_result.get("is_speech", False)
        mic.gain = gain
        if mic.is_speech_active:
            mic.last_speech_time = datetime.now(timezone.utc)

    # ------------------------------------------------------------------
    # Collision resolution
    # ------------------------------------------------------------------

    def resolve_collisions(self) -> Dict[str, bool]:
        """
        Determine which microphones should be suppressed due to collision.

        Returns {mic_id: should_suppress} for every registered microphone.
        """
        active_speakers = [
            mic
            for mic in self._microphones.values()
            if mic.is_speech_active and not mic.is_muted
        ]

        # No collision with 0 or 1 active speaker
        if len(active_speakers) <= 1:
            if active_speakers:
                self._active_speaker = active_speakers[0].mic_id
            return {mid: False for mid in self._microphones}

        # Chairman always wins
        chairman_speaking = any(m.is_chairman for m in active_speakers)
        if chairman_speaking:
            chairman_id = next(
                m.mic_id for m in active_speakers if m.is_chairman
            )
            self._active_speaker = chairman_id
            return {
                mid: (mic.is_speech_active and not mic.is_chairman and not mic.is_muted)
                for mid, mic in self._microphones.items()
            }

        # No chairman — highest priority (then loudest) wins
        active_speakers.sort(key=lambda m: (-m.priority, -m.db_level))
        winner = active_speakers[0]
        self._active_speaker = winner.mic_id

        return {
            mid: (mic.is_speech_active and mic.mic_id != winner.mic_id and not mic.is_muted)
            for mid, mic in self._microphones.items()
        }

    # ------------------------------------------------------------------
    # Accessors / mutators
    # ------------------------------------------------------------------

    def get_active_speaker(self) -> Optional[str]:
        return self._active_speaker

    def get_all_states(self) -> Dict[str, dict]:
        return {
            mid: {
                "mic_id": mic.mic_id,
                "name": mic.name,
                "priority": mic.priority,
                "is_chairman": mic.is_chairman,
                "is_muted": mic.is_muted,
                "is_speech_active": mic.is_speech_active,
                "db_level": mic.db_level,
                "speech_probability": mic.speech_probability,
                "gain": mic.gain,
                "last_speech_time": (
                    mic.last_speech_time.isoformat()
                    if mic.last_speech_time
                    else None
                ),
            }
            for mid, mic in self._microphones.items()
        }

    def set_chairman(self, mic_id: str) -> None:
        """Designate a single microphone as chairman; demote all others."""
        for mic in self._microphones.values():
            mic.is_chairman = mic.mic_id == mic_id

    def set_muted(self, mic_id: str, muted: bool) -> None:
        if mic_id in self._microphones:
            self._microphones[mic_id].is_muted = muted

    def set_priority(self, mic_id: str, priority: int) -> None:
        if mic_id in self._microphones:
            self._microphones[mic_id].priority = max(1, min(10, priority))