1"""
  2# Overview
  3
  4Audio support is very much a work in progress. The radio uses SBC audio encoding, there doesn't exist
  5yet a good way to decode SBC in Python. At some point I'm looking into creating bindings for google's
  6libsbc, but that's a bit of a ways off. (Issue thread [here](https://github.com/khusmann/benlink/issues/11))
  7
  8In the meantime I have a hacky approach for decoding via pyav (ffmpeg bindings for python). I have
  9two proofs of concept, one for receiving audio, the other for sending. To run the receiving audio
 10POC, run:
 11
 12```
 13python -m benlink.examples.audiomonitor <UUID> <CHANNEL>
 14```
 15
 16Where `<UUID>` is the UUID of the device you want to connect to, and `<CHANNEL>` is the RFCOMM audio channel.
 17When the radio receives audio, it will play the audio to the default audio output device using pyaudio.
 18
 19Similarly, to run the sending audio POC, run:
 20
 21```
 22python -m benlink.examples.audiotransmit <UUID> <CHANNEL>
 23```
 24
 25This example uses pyaudio to record audio from the default audio input device and sends it to the radio.
 26"""
 27
 28from __future__ import annotations
 29import typing as t
 30import asyncio
 31from .link import AudioLink, RfcommAudioLink
 32from . import protocol as p
 33
 34
 35class AudioConnection:
 36    _link: AudioLink
 37    _handlers: list[t.Callable[[AudioMessage], None]]
 38
 39    def is_connected(self) -> bool:
 40        return self._link.is_connected()
 41
 42    def __init__(
 43        self,
 44        link: AudioLink,
 45    ):
 46        self._link = link
 47        self._handlers = []
 48
 49    @classmethod
 50    def new_rfcomm(cls, device_uuid: str, channel: int | t.Literal["auto"] = "auto") -> AudioConnection:
 51        return AudioConnection(
 52            RfcommAudioLink(device_uuid, channel)
 53        )
 54
 55    def add_event_handler(self, handler: t.Callable[[AudioEvent], None]) -> t.Callable[[], None]:
 56        def on_message(msg: AudioMessage):
 57            if isinstance(msg, AudioEvent):
 58                handler(msg)
 59        return self._add_message_handler(on_message)
 60
 61    def _add_message_handler(self, handler: t.Callable[[AudioMessage], None]) -> t.Callable[[], None]:
 62        def remove_handler():
 63            self._handlers.remove(handler)
 64
 65        self._handlers.append(handler)
 66
 67        return remove_handler
 68
 69    async def send_message(self, msg: AudioMessage) -> None:
 70        await self._link.send(audio_message_to_protocol(msg))
 71
 72    async def send_message_expect_reply(self, msg: AudioMessage, reply: t.Type[AudioMessageT]) -> AudioMessageT:
 73        queue: asyncio.Queue[AudioMessageT] = asyncio.Queue()
 74
 75        def on_rcv(msg: AudioMessage):
 76            if isinstance(msg, reply):
 77                queue.put_nowait(msg)
 78
 79        remove_handler = self._add_message_handler(on_rcv)
 80
 81        await self.send_message(msg)
 82
 83        out = await queue.get()
 84
 85        remove_handler()
 86
 87        return out
 88
 89    async def connect(self) -> None:
 90        def on_msg(msg: p.AudioMessage):
 91            for handler in self._handlers:
 92                handler(audio_message_from_protocol(msg))
 93        await self._link.connect(on_msg)
 94
 95    async def disconnect(self) -> None:
 96        await self._link.disconnect()
 97
 98    # Audio API
 99
100    async def send_audio_data(self, sbc_data: bytes) -> None:
101        # Radio does not send an ack for audio data
102        await self.send_message(AudioData(sbc_data))
103
104    async def send_audio_end(self) -> None:
105        # Radio does not send an ack for audio end
106        await self.send_message(AudioEnd())
107
108    # Async Context Manager
109    async def __aenter__(self):
110        await self.connect()
111        return self
112
113    async def __aexit__(
114        self,
115        exc_type: t.Any,
116        exc_value: t.Any,
117        traceback: t.Any,
118    ) -> None:
119        # Send extra audio end message to ensure radio stops transmitting
120        await self.send_audio_end()
121        # Wait for the audio end message to be fully sent
122        # before disconnecting, otherwise radio
123        # gets stuck in transmit mode (no ack from radio, unfortunately)
124        await asyncio.sleep(1.5)
125        await self.disconnect()
126
127
128class AudioData(t.NamedTuple):
129    sbc_data: bytes
130
131
132class AudioEnd:
133    pass
134
135
136class AudioAck:
137    pass
138
139
140class AudioUnknown(t.NamedTuple):
141    type: int
142    data: bytes
143
144
145AudioEvent = AudioData | AudioEnd | AudioUnknown
146
147AudioMessage = AudioEvent | AudioAck
148
149AudioMessageT = t.TypeVar("AudioMessageT", bound=AudioMessage)
150
151
152def audio_message_from_protocol(proto: p.AudioMessage) -> AudioMessage:
153    match proto:
154        case p.AudioData(sbc_data=sbc_data):
155            return AudioData(sbc_data)
156        case p.AudioEnd():
157            return AudioEnd()
158        case p.AudioAck():
159            return AudioAck()
160        case p.AudioUnknown(type=type, data=data):
161            return AudioUnknown(type, data)
162
163
164def audio_message_to_protocol(msg: AudioMessage) -> p.AudioMessage:
165    match msg:
166        case AudioData(sbc_data=sbc_data):
167            return p.AudioData(sbc_data)
168        case AudioEnd():
169            return p.AudioEnd()
170        case AudioAck():
171            return p.AudioAck()
172        case AudioUnknown(type=type, data=data):
173            return p.AudioUnknown(type, data)