import asyncio
import socket
import logging
import numpy as np
from typing import AsyncGenerator

class NetworkAudioReceiverUncoded:
    """
    Receives PCM audio over UDP and provides an async generator interface for uncoded PCM frames.
    Combines network receiving and input logic for use with Auracast streamer.
    """
    def __init__(self, port: int = 50007, samplerate: int = 16000, channels: int = 1, chunk_size: int = 1024):
        self.port = port
        self.samplerate = samplerate
        self.channels = channels
        self.chunk_size = chunk_size
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.sock.bind(('0.0.0.0', self.port))
        self.sock.setblocking(False)
        self._running = False
        # Reduce queue size for lower latency (less buffering)
        self._queue = asyncio.Queue(maxsize=2)  # Was 20

    async def receive(self):
        self._running = True
        logging.info(f"NetworkAudioReceiver listening on UDP port {self.port}")
        try:
            while self._running:
                try:
                    data, _ = await asyncio.get_event_loop().sock_recvfrom(self.sock, self.chunk_size * 2)
                    await self._queue.put(data)
                except Exception:
                    await asyncio.sleep(0.01)
        finally:
            self.sock.close()
            logging.info("NetworkAudioReceiver stopped.")

    def stop(self):
        self._running = False

    async def open(self):
        # Dummy PCM format object
        class PCMFormat:
            channels = self.channels
            sample_type = 'int16'
            sample_rate = self.samplerate
        return PCMFormat()

    def rewind(self):
        pass  # Not supported for live network input

    async def frames(self, samples_per_frame: int) -> AsyncGenerator[np.ndarray, None]:
        bytes_per_frame = samples_per_frame * 2 * self.channels  # 2 bytes for int16
        buf = bytearray()
        while True:
            data = await self._queue.get()
            # Optional: log queue size for latency debugging
            # logging.debug(f'NetworkAudioReceiver queue size: {self._queue.qsize()}')
            if data is None:
                break
            buf.extend(data)
            while len(buf) >= bytes_per_frame:
                frame = np.frombuffer(buf[:bytes_per_frame], dtype=np.int16).reshape(-1, self.channels)
                # Optional: log when a frame is yielded
                # logging.debug(f'Yielding frame of shape {frame.shape}')
                yield frame
                buf = buf[bytes_per_frame:]
