main.py

import asyncio
import ctypes
import ctypes.util
import struct
import wave
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional, Tuple


# ============================================================
# Exceptions
# ============================================================

class OpusError(RuntimeError):
    pass


# ============================================================
# Config
# ============================================================

@dataclass(frozen=True)
class AudioConfig:
    host: str = "0.0.0.0"
    port: int = 9999
    wav_file: str = "wall-e.wav"
    uplink_output_dir: str = "recordings"

    sample_rate: int = 16000
    channels: int = 1
    sample_width: int = 2  # PCM16 mono

    frame_ms: int = 20

    application_voip: int = 2048
    opus_ok: int = 0

    opus_set_bitrate_request: int = 4002
    opus_set_complexity_request: int = 4010
    opus_set_vbr_request: int = 4006
    opus_set_signal_request: int = 4024
    opus_set_packet_loss_perc_request: int = 4014
    opus_set_inband_fec_request: int = 4012
    opus_signal_voice: int = 3001

    target_bitrate: int = 12000
    encoder_complexity: int = 5
    expected_packet_loss_perc: int = 8
    use_vbr: int = 0
    use_inband_fec: int = 1

    opus_magic: bytes = b"OPUS"
    pcm_magic: bytes = b"PCM!"
    version: int = 1
    flag_end: int = 0x01

    end_repeat_count: int = 5
    end_repeat_interval: float = 0.02
    log_every_n_packets: int = 200

    @property
    def frame_samples(self) -> int:
        return self.sample_rate * self.frame_ms // 1000

    @property
    def pcm_bytes_per_frame(self) -> int:
        return self.frame_samples * self.channels * self.sample_width

    @property
    def send_interval_sec(self) -> float:
        return self.frame_ms / 1000.0


# ============================================================
# Packet formats
# ============================================================

class BasePacketFormat:
    HEADER_STRUCT = struct.Struct("!4sBBIIHH")

    def __init__(self, magic: bytes, version: int, flag_end: int):
        self.magic = magic
        self.version = version
        self.flag_end = flag_end

    def pack_audio_packet(
        self,
        seq: int,
        pts_samples: int,
        frame_samples: int,
        payload: bytes,
    ) -> bytes:
        header = self.HEADER_STRUCT.pack(
            self.magic,
            self.version,
            0,
            seq,
            pts_samples,
            frame_samples,
            len(payload),
        )
        return header + payload

    def pack_end_packet(self, seq: int, pts_samples: int) -> bytes:
        return self.HEADER_STRUCT.pack(
            self.magic,
            self.version,
            self.flag_end,
            seq,
            pts_samples,
            0,
            0,
        )

    def unpack_packet(self, data: bytes) -> dict:
        if len(data) < self.HEADER_STRUCT.size:
            raise ValueError(
                f"Packet too short: got {len(data)}, expected at least {self.HEADER_STRUCT.size}"
            )

        magic, version, flags, seq, pts_samples, frame_samples, payload_len = (
            self.HEADER_STRUCT.unpack(data[: self.HEADER_STRUCT.size])
        )

        payload = data[self.HEADER_STRUCT.size :]

        return {
            "magic": magic,
            "version": version,
            "flags": flags,
            "seq": seq,
            "pts_samples": pts_samples,
            "frame_samples": frame_samples,
            "payload_len": payload_len,
            "payload": payload,
        }

    def validate_packet(self, packet: dict) -> None:
        if packet["magic"] != self.magic:
            raise ValueError(f"Invalid magic: {packet['magic']!r}")

        if packet["version"] != self.version:
            raise ValueError(f"Invalid version: {packet['version']}")

        expected_total = self.HEADER_STRUCT.size + packet["payload_len"]
        actual_total = self.HEADER_STRUCT.size + len(packet["payload"])

        if actual_total != expected_total:
            raise ValueError(
                f"Packet size mismatch: actual={actual_total}, expected={expected_total}"
            )

    def is_end_packet(self, packet: dict) -> bool:
        return (packet["flags"] & self.flag_end) != 0


class OpusPacketFormat(BasePacketFormat):
    def __init__(self, config: AudioConfig):
        super().__init__(config.opus_magic, config.version, config.flag_end)


class PcmPacketFormat(BasePacketFormat):
    def __init__(self, config: AudioConfig):
        super().__init__(config.pcm_magic, config.version, config.flag_end)


# ============================================================
# Base libopus loader
# ============================================================

class LibOpusBase:
    def __init__(self, config: AudioConfig):
        self.config = config
        self.lib = self._load_libopus()

    def _load_libopus(self):
        candidates = []
        found = ctypes.util.find_library("opus")
        if found:
            candidates.append(found)

        candidates.extend(
            [
                "libopus.so",
                "libopus.so.0",
                "libopus.dylib",
                "opus.dll",
                "libopus-0.dll",
            ]
        )

        last_error = None
        for name in candidates:
            try:
                return ctypes.CDLL(name)
            except OSError as exc:
                last_error = exc

        raise OpusError(f"Could not load libopus: {last_error}")


# ============================================================
# Encoder
# ============================================================

class LibOpusEncoder(LibOpusBase):
    def __init__(self, config: AudioConfig):
        super().__init__(config)
        self.encoder = None
        self._configure_signatures()
        self._create_encoder()
        self._configure_encoder()

    def _configure_signatures(self):
        self.lib.opus_encoder_create.argtypes = [
            ctypes.c_int,
            ctypes.c_int,
            ctypes.c_int,
            ctypes.POINTER(ctypes.c_int),
        ]
        self.lib.opus_encoder_create.restype = ctypes.c_void_p

        self.lib.opus_encode.argtypes = [
            ctypes.c_void_p,
            ctypes.POINTER(ctypes.c_int16),
            ctypes.c_int,
            ctypes.POINTER(ctypes.c_ubyte),
            ctypes.c_int,
        ]
        self.lib.opus_encode.restype = ctypes.c_int

        self.lib.opus_encoder_destroy.argtypes = [ctypes.c_void_p]
        self.lib.opus_encoder_destroy.restype = None

        self.lib.opus_encoder_ctl.restype = ctypes.c_int

    def _create_encoder(self):
        err = ctypes.c_int(0)
        self.encoder = self.lib.opus_encoder_create(
            self.config.sample_rate,
            self.config.channels,
            self.config.application_voip,
            ctypes.byref(err),
        )
        if not self.encoder or err.value != self.config.opus_ok:
            raise OpusError(f"opus_encoder_create failed: {err.value}")

    def _ctl(self, request: int, value: int):
        result = self.lib.opus_encoder_ctl(self.encoder, request, value)
        if result != self.config.opus_ok:
            raise OpusError(f"opus_encoder_ctl({request}, {value}) failed: {result}")

    def _configure_encoder(self):
        self._ctl(self.config.opus_set_bitrate_request, self.config.target_bitrate)
        self._ctl(self.config.opus_set_complexity_request, self.config.encoder_complexity)
        self._ctl(self.config.opus_set_vbr_request, self.config.use_vbr)
        self._ctl(self.config.opus_set_signal_request, self.config.opus_signal_voice)
        self._ctl(
            self.config.opus_set_packet_loss_perc_request,
            self.config.expected_packet_loss_perc,
        )
        self._ctl(
            self.config.opus_set_inband_fec_request,
            self.config.use_inband_fec,
        )

    def encode_pcm16_frame(self, pcm16_frame: bytes) -> bytes:
        expected_len = self.config.frame_samples * 2
        if len(pcm16_frame) != expected_len:
            raise ValueError(
                f"Invalid PCM frame size: got {len(pcm16_frame)}, expected {expected_len}"
            )

        pcm_array = (ctypes.c_int16 * self.config.frame_samples).from_buffer_copy(
            pcm16_frame
        )
        max_packet_bytes = 512
        out_buf = (ctypes.c_ubyte * max_packet_bytes)()

        encoded_len = self.lib.opus_encode(
            self.encoder,
            pcm_array,
            self.config.frame_samples,
            out_buf,
            max_packet_bytes,
        )
        if encoded_len < 0:
            raise OpusError(f"opus_encode failed: {encoded_len}")

        return bytes(out_buf[:encoded_len])

    def close(self):
        if self.encoder:
            self.lib.opus_encoder_destroy(self.encoder)
            self.encoder = None

    def __del__(self):
        try:
            self.close()
        except Exception:
            pass


# ============================================================
# Preencoded stream container
# ============================================================

@dataclass
class PreencodedStream:
    packets: List[bytes]
    end_packets: List[bytes]
    total_audio_packets: int


# ============================================================
# WAV -> preencoded Opus packets
# ============================================================

class WavPreencoder:
    def __init__(self, config: AudioConfig, packet_format: OpusPacketFormat):
        self.config = config
        self.packet_format = packet_format

    def build_stream(self, wav_path: Path) -> PreencodedStream:
        if not wav_path.exists():
            raise FileNotFoundError(f"WAV file not found: {wav_path}")

        encoder = LibOpusEncoder(self.config)
        packets: List[bytes] = []

        try:
            with wave.open(str(wav_path), "rb") as wf:
                self._validate_wav(wf)

                seq = 0
                pts_samples = 0

                while True:
                    pcm = wf.readframes(self.config.frame_samples)
                    if not pcm:
                        break

                    if len(pcm) < self.config.pcm_bytes_per_frame:
                        pcm += b"\x00" * (self.config.pcm_bytes_per_frame - len(pcm))

                    opus_payload = encoder.encode_pcm16_frame(pcm)
                    packet = self.packet_format.pack_audio_packet(
                        seq=seq,
                        pts_samples=pts_samples,
                        frame_samples=self.config.frame_samples,
                        payload=opus_payload,
                    )
                    packets.append(packet)

                    if seq < 8 or seq % self.config.log_every_n_packets == 0:
                        print(
                            f"[SERVER] preencoded seq={seq} pts={pts_samples} "
                            f"opus_bytes={len(opus_payload)} udp_bytes={len(packet)}"
                        )

                    seq += 1
                    pts_samples += self.config.frame_samples

            end_packet = self.packet_format.pack_end_packet(seq=seq, pts_samples=pts_samples)
            end_packets = [end_packet for _ in range(self.config.end_repeat_count)]

            print("[SERVER] Pre-encode complete:")
            print(f"  frame_ms={self.config.frame_ms}")
            print(f"  frame_samples={self.config.frame_samples}")
            print(f"  bitrate={self.config.target_bitrate}")
            print(f"  packet_loss_perc={self.config.expected_packet_loss_perc}")
            print(f"  inband_fec={self.config.use_inband_fec}")
            print(f"  audio_packets={len(packets)}")
            print(f"  send_interval_sec={self.config.send_interval_sec:.4f}")
            print(f"  end_repeat_count={self.config.end_repeat_count}")

            return PreencodedStream(
                packets=packets,
                end_packets=end_packets,
                total_audio_packets=len(packets),
            )
        finally:
            encoder.close()

    def _validate_wav(self, wf: wave.Wave_read) -> None:
        channels = wf.getnchannels()
        sample_width = wf.getsampwidth()
        sample_rate = wf.getframerate()
        comptype = wf.getcomptype()

        print("[SERVER] WAV info:")
        print(f"  channels={channels}")
        print(f"  sample_width={sample_width}")
        print(f"  sample_rate={sample_rate}")
        print(f"  compression={comptype}")

        if channels != self.config.channels:
            raise ValueError("WAV must be mono")
        if sample_width != self.config.sample_width:
            raise ValueError("WAV must be 16-bit PCM")
        if sample_rate != self.config.sample_rate:
            raise ValueError("WAV must be 16 kHz")
        if comptype != "NONE":
            raise ValueError("WAV must be uncompressed PCM")


# ============================================================
# PCM session recorder
# ============================================================

class PcmSessionRecorder:
    def __init__(self, config: AudioConfig):
        self.config = config
        self.output_dir = Path(config.uplink_output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)

        self.wave_file: Optional[wave.Wave_write] = None
        self.current_path: Optional[Path] = None
        self.started = False
        self.finished = False
        self.expected_seq = 0
        self.received_packets = 0
        self.received_bytes = 0
        self.session_addr: Optional[Tuple[str, int]] = None

    def _build_session_path(self, addr: Tuple[str, int]) -> Path:
        stamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
        host, port = addr
        safe_host = host.replace(".", "_").replace(":", "_")
        return self.output_dir / f"uplink_{stamp}_{safe_host}_{port}.wav"

    def start_session(self, addr: Tuple[str, int], first_seq: int) -> None:
        self.close_session()

        self.current_path = self._build_session_path(addr)
        self.wave_file = wave.open(str(self.current_path), "wb")
        self.wave_file.setnchannels(self.config.channels)
        self.wave_file.setsampwidth(self.config.sample_width)
        self.wave_file.setframerate(self.config.sample_rate)

        self.started = True
        self.finished = False
        self.expected_seq = first_seq
        self.received_packets = 0
        self.received_bytes = 0
        self.session_addr = addr

        print(f"[PCM RX] New session from {addr}")
        print(f"[PCM RX] Writing WAV to: {self.current_path}")

    def write_packet(self, addr: Tuple[str, int], packet: dict) -> None:
        if not self.started or self.wave_file is None or self.finished:
            self.start_session(addr, packet["seq"])

        if self.session_addr != addr:
            print(f"[PCM RX] New sender detected: {addr}, closing previous session")
            self.start_session(addr, packet["seq"])

        sample_count = packet["frame_samples"]
        payload_len = packet["payload_len"]
        payload = packet["payload"]

        expected_payload_len = sample_count * self.config.sample_width
        if payload_len != expected_payload_len:
            raise ValueError(
                f"PCM payload mismatch: payload_len={payload_len}, "
                f"expected={expected_payload_len}"
            )

        if packet["seq"] != self.expected_seq:
            print(
                f"[PCM RX] sequence gap: expected={self.expected_seq} "
                f"got={packet['seq']}"
            )
            self.expected_seq = packet["seq"]

        self.wave_file.writeframes(payload)
        self.received_packets += 1
        self.received_bytes += len(payload)
        self.expected_seq += 1

        seq = packet["seq"]
        if seq < 8 or seq % self.config.log_every_n_packets == 0:
            print(
                f"[PCM RX] wrote seq={seq} pts={packet['pts_samples']} "
                f"samples={sample_count} pcm_bytes={payload_len}"
            )

    def finish_session(self, packet: dict) -> None:
        if not self.started:
            return

        print(
            f"[PCM RX] END received seq={packet['seq']} "
            f"pts={packet['pts_samples']}"
        )
        self.close_session()

    def close_session(self) -> None:
        if self.wave_file is not None:
            try:
                self.wave_file.close()
            except Exception:
                pass

            print("[PCM RX] WAV closed")
            print(f"[PCM RX] received_packets={self.received_packets}")
            print(f"[PCM RX] received_bytes={self.received_bytes}")
            if self.current_path is not None:
                print(f"[PCM RX] saved_to={self.current_path}")

        self.wave_file = None
        self.current_path = None
        self.started = False
        self.finished = False
        self.expected_seq = 0
        self.received_packets = 0
        self.received_bytes = 0
        self.session_addr = None


# ============================================================
# Unified UDP bridge
# ============================================================

class UDPAudioBridgeProtocol(asyncio.DatagramProtocol):
    def __init__(
        self,
        config: AudioConfig,
        preencoded_stream: PreencodedStream,
        opus_packet_format: OpusPacketFormat,
        pcm_packet_format: PcmPacketFormat,
    ):
        self.config = config
        self.preencoded_stream = preencoded_stream
        self.opus_packet_format = opus_packet_format
        self.pcm_packet_format = pcm_packet_format
        self.transport = None
        self.client_tasks: Dict[Tuple[str, int], asyncio.Task] = {}
        self.pcm_recorder = PcmSessionRecorder(config)

    def connection_made(self, transport):
        self.transport = transport
        print("[BRIDGE] UDP audio bridge ready")
        print(f"[BRIDGE] Listening on udp://{self.config.host}:{self.config.port}")
        print(f"[BRIDGE] PCM sessions dir: {Path(self.config.uplink_output_dir).resolve()}")

    def datagram_received(self, data, addr):
        if data == b"START":
            print(f"[BRIDGE] START from {addr}")
            self.restart_stream_for_client(addr)
            return

        if len(data) >= 4 and data[:4] == self.config.pcm_magic:
            self.handle_pcm_packet(data, addr)
            return

        if len(data) >= 4 and data[:4] == self.config.opus_magic:
            print(f"[BRIDGE] Unexpected inbound OPUS packet from {addr}, ignoring")
            return

        try:
            message = data.decode(errors="ignore").strip()
        except Exception:
            message = "<binary>"
        print(f"[BRIDGE] Unknown packet from {addr}: {message[:80]}")

    def restart_stream_for_client(self, addr: Tuple[str, int]) -> None:
        old_task = self.client_tasks.get(addr)
        if old_task is not None and not old_task.done():
            print(f"[BRIDGE] Stream already active for {addr}, ignoring duplicate START")
            return

        task = asyncio.create_task(self.stream_preencoded(addr))
        self.client_tasks[addr] = task

    async def stream_preencoded(self, addr: Tuple[str, int]) -> None:
        try:
            print(f"[BRIDGE] Starting Opus stream for {addr}")

            loop = asyncio.get_running_loop()
            start_time = loop.time()

            for seq, packet in enumerate(self.preencoded_stream.packets):
                self.transport.sendto(packet, addr)

                if seq < 8 or seq % self.config.log_every_n_packets == 0:
                    print(f"[BRIDGE] sent OPUS seq={seq} udp_bytes={len(packet)}")

                next_send_time = start_time + ((seq + 1) * self.config.send_interval_sec)
                delay = next_send_time - loop.time()
                if delay > 0:
                    await asyncio.sleep(delay)

            for i, end_packet in enumerate(self.preencoded_stream.end_packets, start=1):
                self.transport.sendto(end_packet, addr)
                print(f"[BRIDGE] sent OPUS END repeat {i}/{len(self.preencoded_stream.end_packets)}")
                if i < self.config.end_repeat_count:
                    await asyncio.sleep(self.config.end_repeat_interval)

            print(f"[BRIDGE] Opus stream finished for {addr}")

        except asyncio.CancelledError:
            print(f"[BRIDGE] Opus stream cancelled for {addr}")
            raise
        except Exception as exc:
            print(f"[BRIDGE] Opus stream error for {addr}: {exc}")
        finally:
            current_task = self.client_tasks.get(addr)
            if current_task is asyncio.current_task():
                self.client_tasks.pop(addr, None)

    def handle_pcm_packet(self, data: bytes, addr: Tuple[str, int]) -> None:
        try:
            packet = self.pcm_packet_format.unpack_packet(data)
            self.pcm_packet_format.validate_packet(packet)

            if self.pcm_packet_format.is_end_packet(packet):
                self.pcm_recorder.finish_session(packet)
                return

            self.pcm_recorder.write_packet(addr, packet)

        except Exception as exc:
            print(f"[PCM RX] receive/write error from {addr}: {exc}")

    def connection_lost(self, exc):
        for task in self.client_tasks.values():
            task.cancel()
        self.client_tasks.clear()
        self.pcm_recorder.close_session()
        print("[BRIDGE] UDP bridge closed")


# ============================================================
# Application
# ============================================================

class AudioBridgeApp:
    def __init__(self, config: AudioConfig):
        self.config = config
        self.opus_packet_format = OpusPacketFormat(config)
        self.pcm_packet_format = PcmPacketFormat(config)
        self.preencoder = WavPreencoder(config, self.opus_packet_format)

    async def run_bridge(self):
        wav_path = Path(self.config.wav_file)
        preencoded_stream = self.preencoder.build_stream(wav_path)

        loop = asyncio.get_running_loop()
        transport, _protocol = await loop.create_datagram_endpoint(
            lambda: UDPAudioBridgeProtocol(
                self.config,
                preencoded_stream,
                self.opus_packet_format,
                self.pcm_packet_format,
            ),
            local_addr=(self.config.host, self.config.port),
        )

        try:
            await asyncio.Future()
        finally:
            transport.close()


# ============================================================
# CLI
# ============================================================

def build_app() -> AudioBridgeApp:
    return AudioBridgeApp(AudioConfig())


async def async_main():
    app = build_app()
    await app.run_bridge()


if __name__ == "__main__":
    asyncio.run(async_main())