import asyncio
import logging
from typing import Dict
from queues import queue_names

logger = logging.getLogger("s2s_server")

class QueueRegistry:
    def __init__(self, session_id: str = ""):
        self.session_id = session_id
        self._queues: Dict[str, asyncio.Queue] = {}
        
        # Initialize all queues
        for name in [
            queue_names.AUDIO,
            queue_names.TRANSCRIPT,
            queue_names.LLM,
            queue_names.TTS,
            queue_names.TOOL,
            queue_names.MEMORY,
            queue_names.RAG,
            queue_names.SEARCH,
            queue_names.METRICS,
        ]:
            self._queues[name] = asyncio.Queue()

    def get(self, name: str) -> asyncio.Queue:
        if name not in self._queues:
            # Dynamically initialize if missing to prevent failure
            logger.warning(f"[{self.session_id}] Queue '{name}' not found. Initializing dynamically.")
            self._queues[name] = asyncio.Queue()
        return self._queues[name]

    def clear_all(self):
        logger.debug(f"[{self.session_id}] QueueRegistry: Flushing all queues.")
        for name, q in self._queues.items():
            while not q.empty():
                try:
                    q.get_nowait()
                except (asyncio.QueueEmpty, ValueError):
                    break