import asyncio
import time
import audioop

class UdpClientProtocol(asyncio.DatagramProtocol):
    def __init__(self):
        self.transport = None
        self.future = None

    def connection_made(self, transport):
        self.transport = transport

    def datagram_received(self, data, addr):
        if self.future and not self.future.done():
            self.future.set_result(data)

class VocoderEngine:
    def __init__(self, mode="emulator", host="127.0.0.1"):
        self.host = host
        self.ports = list(range(2470, 2480))
        self.pool = {}         # porta -> (transport, protocol, lock)
        self.stream_alloc = {} # stream_id -> porta
        self.last_seen = {}    # stream_id -> timestamp ultimo pacchetto

    async def avvia(self):
        loop = asyncio.get_running_loop()
        for p in self.ports:
            transport, protocol = await loop.create_datagram_endpoint(
                lambda: UdpClientProtocol(),
                remote_addr=(self.host, p)
            )
            # Aggiungiamo un Lock esclusivo per ogni porta UDP
            self.pool[p] = (transport, protocol, asyncio.Lock())
        print(f"[Vocoder] Pool UDP Multi-Porta Sincronizzato: {self.ports}")

    def pulisci_stream_inattivi(self):
        """Libera automaticamente le porte dei Talkgroup/SIP muti da più di 2.5 secondi"""
        ora = time.time()
        inattivi = [s for s, t in self.last_seen.items() if ora - t > 2.5]
        for s in inattivi:
            if s in self.stream_alloc:
                del self.stream_alloc[s]
            del self.last_seen[s]

    def assegna_porta(self, stream_id):
        self.pulisci_stream_inattivi()
        self.last_seen[stream_id] = time.time()

        if stream_id in self.stream_alloc:
            return self.stream_alloc[stream_id]

        porte_usate = set(self.stream_alloc.values())
        porte_libere = [p for p in self.ports if p not in porte_usate]

        if porte_libere:
            porta_scelta = porte_libere[0]
        else:
            # Se tutte le 3 porte sono saturate contemporaneamente, usiamo la prima come fallback
            porta_scelta = self.ports[0]

        self.stream_alloc[stream_id] = porta_scelta
        return porta_scelta

    def rilascia_porta(self, stream_id):
        if stream_id in self.stream_alloc:
            del self.stream_alloc[stream_id]
        if stream_id in self.last_seen:
            del self.last_seen[stream_id]

    async def _scambia_dati(self, porta, dati):
        transport, protocol, lock = self.pool[porta]
        
        # Garantisce che un solo pacchetto alla volta usi questa porta UDP!
        async with lock:
            loop = asyncio.get_running_loop()
            protocol.future = loop.create_future()
            transport.sendto(dati)
            try:
                return await asyncio.wait_for(protocol.future, timeout=0.08)
            except asyncio.TimeoutError:
                return b""
            finally:
                protocol.future = None

    async def converti_ambe_in_pcm(self, stream_id, ambe_bytes):
        porta = self.assegna_porta(stream_id)
        # Invio dati RAW
        return await self._scambia_dati(porta, ambe_bytes)

    async def converti_rtp_in_ambe(self, stream_id, pcm_bytes):
        porta = self.assegna_porta(stream_id)
        
        # Se riceviamo 160 byte dal SIP, li convertiamo in PCM lineare 16-bit (320 byte)
        if len(pcm_bytes) == 160:
            # Uso alaw2lin per PCMA (A-Law, standard europeo).
            # (Se il tuo centralino usa u-Law, ti basterà cambiare in ulaw2lin)
            pcm_bytes = audioop.alaw2lin(pcm_bytes, 2)
            
        return await self._scambia_dati(porta, pcm_bytes)