#!/usr/bin/python3
"""The desktop add-on's sound, to and from the viewer.

Listens on a Unix socket that devgrail-desktop puts behind its /audio
WebSocket, so the viewer can hear the desktop and speak into it. Each
connection carries frames both ways: a type byte, the payload's length as two
big-endian bytes, and the payload. The bridge does not keep WebSocket messages
apart on the socket, so the length is what tells one frame from the next.

  0x01 CONFIG     viewer -> desktop  one byte: bit 0 set sends the speakers
  0x02 MIC        viewer -> desktop  one Opus packet, 48 kHz mono
  0x81 SPEAKERS   desktop -> viewer  one Opus packet, 48 kHz stereo, 20 ms
  0x82 SILENCE    desktop -> viewer  empty: no speakers frames until sound

The speakers are what apps play into desktop_out, read off its monitor; the
microphone is played into devgrail_mic_in, whose monitor apps record as
devgrail_mic (desktop-audio.pa). Both are encoded and decoded with GStreamer,
in this process, so neither direction starts a program per connection.

A connection that sends anything else is closed: everything on it comes from a
browser, and the viewer's code is the only client that should be here.
"""

import os
import socket
import struct
import sys
import threading
import time

import gi

gi.require_version("Gst", "1.0")
from gi.repository import GLib, Gst  # noqa: E402

SOCKET_DIR = "/tmp/devgrail-audio"
SOCKET_PATH = os.path.join(SOCKET_DIR, "stream")

CONFIG = 0x01
MIC = 0x02
SPEAKERS = 0x81
SILENCE = 0x82
HEADER = struct.Struct(">BH")

# The largest Opus packet: 120 ms of 20 ms frames at the codec's per-frame
# ceiling of 1275 bytes, plus framing. The viewer sends one 20 ms frame.
MAX_PACKET = 6 * 1275 + 8

# A desktop has one viewer at a time; the tab and the desktop window hand it
# over between them, and for a moment both can be connected. The rest is room
# for connections that died without closing: a laptop that slept or changed
# networks leaves one here until TCP gives up on it, which can take many
# minutes, and a cap that small would refuse the viewer's own reconnects.
MAX_CONNECTIONS = 16

# How long to wait before trying a pipeline again that failed, typically
# because the sound server has not started yet.
RETRY_SECONDS = 1.0

SPEAKERS_PIPELINE = (
    "pulsesrc device=desktop_out.monitor buffer-time=40000 latency-time=10000"
    ' client-name="DevGrail viewer speakers"'
    " ! audioconvert ! audioresample ! audio/x-raw,format=S16LE,rate=48000,channels=2"
    " ! opusenc name=enc bitrate=96000 frame-size=20 audio-type=generic"
    " ! appsink name=out sync=false max-buffers=25 drop=true"
)

# A silent desktop would still cost the viewer 96 kbit/s, so the speakers
# pause while it is quiet: after SILENCE_HOLD_FRAMES of samples no further
# from zero than SILENCE_PEAK, one SILENCE frame, then nothing until there is
# sound again. The peak leaves room for dither, about -72 dBFS, far below what
# anyone hears. The hold outlasts the encoder's delay, so the end of the sound
# reaches the viewer, and the silent packets it sends first leave both codecs
# at rest, so the next packet picks up cleanly after the gap.
SILENCE_PEAK = 8
SILENCE_HOLD_FRAMES = 48000 // 2


def is_silent(data):
    """Whether 16-bit little-endian samples all lie within SILENCE_PEAK of zero."""
    # Sound almost always shows in the high bytes, which are cheap to check:
    # a quiet sample's is 0x00, or 0xff below zero.
    if data[1::2].translate(None, b"\x00\xff"):
        return False
    samples = memoryview(data).cast("h")
    return len(samples) == 0 or (max(samples) <= SILENCE_PEAK and min(samples) >= -SILENCE_PEAK)


class SilenceGate:
    """Which of the speakers' packets one viewer is sent.

    hear() sees the sound on its way into the encoder, on GStreamer's thread;
    frame() is asked of each packet that comes out, on the speakers' thread.
    The sound is heard before its packet comes out, so a pause ends ahead of
    the packet that holds the sound that ended it.
    """

    def __init__(self):
        self.silent_frames = 0
        self.paused = False

    def hear(self, data):
        if is_silent(data):
            self.silent_frames += len(data) // 4
        else:
            self.silent_frames = 0

    def frame(self, packet):
        """The frame to send for this packet as (type, payload), or None."""
        if self.silent_frames < SILENCE_HOLD_FRAMES:
            self.paused = False
            return SPEAKERS, packet
        if self.paused:
            return None
        self.paused = True
        return SILENCE, b""

# sync=false plays each packet as it arrives: the viewer's network already
# decided its timing, and waiting on timestamps would only add latency. That
# also means nothing ever drains a backlog: packets that arrive in a burst
# (after a network stall, or a browser tab that was held up) queue in front of
# the sound server, and the voice lags by that much for as long as the
# connection lasts, as it does a little more over time when the viewer's
# microphone clock runs faster than the desktop's. So appsrc keeps at most
# MIC_QUEUE_PACKETS of 20 ms and drops the oldest past that; opusdec's packet
# loss concealment covers the gap.
MIC_QUEUE_PACKETS = 4
MIC_PIPELINE = (
    "appsrc name=src is-live=true format=time do-timestamp=true"
    f" max-buffers={MIC_QUEUE_PACKETS} leaky-type=downstream"
    ' caps="audio/x-opus,rate=48000,channels=1,channel-mapping-family=0"'
    " ! opusdec plc=true ! audioconvert ! audioresample"
    " ! audio/x-raw,rate=48000,channels=1"
    " ! pulsesink device=devgrail_mic_in sync=false buffer-time=80000 latency-time=20000"
    ' client-name="DevGrail viewer microphone"'
)


def pack_frame(kind, payload):
    return HEADER.pack(kind, len(payload)) + payload


def read_frame(reader):
    """The next frame as (type, payload), or None once the connection ends."""
    header = reader.read(HEADER.size)
    if len(header) < HEADER.size:
        return None
    kind, length = HEADER.unpack(header)
    payload = reader.read(length)
    if len(payload) < length:
        return None
    return kind, payload


def log(message):
    print(f"devgrail-desktop audio: {message}", file=sys.stderr, flush=True)


def pipeline_error(bus):
    message = bus.pop_filtered(Gst.MessageType.ERROR | Gst.MessageType.EOS)
    if message is None:
        return None
    if message.type == Gst.MessageType.EOS:
        return "stream ended"
    error, _ = message.parse_error()
    return error.message


def listen(_pad, info, gate):
    """Hands the gate each buffer of sound on its way into the encoder."""
    buffer = info.get_buffer()
    if buffer is not None:
        gate.hear(buffer.extract_dup(0, buffer.get_size()))
    return Gst.PadProbeReturn.OK


class Speakers(threading.Thread):
    """Sends the desktop's sound to one viewer until stopped or it goes away."""

    def __init__(self, send):
        super().__init__(daemon=True)
        self.send = send
        self.stopping = threading.Event()

    def stop(self):
        self.stopping.set()

    def run(self):
        while not self.stopping.is_set():
            try:
                pipeline = Gst.parse_launch(SPEAKERS_PIPELINE)
            except GLib.Error as error:
                # A missing element: logged, and tried again like any other
                # failure rather than ending this viewer's speakers for good.
                log(f"speakers: {error.message}")
                self.stopping.wait(RETRY_SECONDS)
                continue
            sink = pipeline.get_by_name("out")
            bus = pipeline.get_bus()
            gate = SilenceGate()
            pipeline.get_by_name("enc").get_static_pad("sink").add_probe(Gst.PadProbeType.BUFFER, listen, gate)
            pipeline.set_state(Gst.State.PLAYING)
            try:
                while not self.stopping.is_set():
                    error = pipeline_error(bus)
                    if error:
                        log(f"speakers: {error}")
                        break
                    sample = sink.emit("try-pull-sample", 100 * Gst.MSECOND)
                    if sample is None:
                        continue
                    buffer = sample.get_buffer()
                    frame = gate.frame(buffer.extract_dup(0, buffer.get_size()))
                    if frame is not None:
                        self.send(*frame)
            except OSError:
                # The viewer went away; the reader notices too and cleans up.
                self.stopping.set()
            finally:
                pipeline.set_state(Gst.State.NULL)
            self.stopping.wait(RETRY_SECONDS)


class Microphone:
    """Plays one viewer's microphone into the desktop's microphone source."""

    def __init__(self):
        self.pipeline = None
        self.retry_at = 0.0

    def push(self, packet):
        if self.pipeline is not None:
            error = pipeline_error(self.bus)
            if error:
                log(f"microphone: {error}")
                self.close()
                self.retry_at = time.monotonic() + RETRY_SECONDS
        if self.pipeline is None:
            if time.monotonic() < self.retry_at:
                return
            try:
                self.pipeline = Gst.parse_launch(MIC_PIPELINE)
            except GLib.Error as error:
                log(f"microphone: {error.message}")
                self.retry_at = time.monotonic() + RETRY_SECONDS
                return
            self.source = self.pipeline.get_by_name("src")
            self.bus = self.pipeline.get_bus()
            self.pipeline.set_state(Gst.State.PLAYING)
        self.source.emit("push-buffer", Gst.Buffer.new_wrapped(packet))

    def close(self):
        if self.pipeline is not None:
            self.pipeline.set_state(Gst.State.NULL)
            self.pipeline = None


def serve(conn):
    send_lock = threading.Lock()

    def send(kind, payload):
        with send_lock:
            conn.sendall(pack_frame(kind, payload))

    speakers = None
    microphone = Microphone()
    reader = conn.makefile("rb")
    try:
        while True:
            frame = read_frame(reader)
            if frame is None:
                return
            kind, payload = frame
            length = len(payload)
            if kind == CONFIG and length == 1:
                wanted = bool(payload[0] & 1)
                if wanted and speakers is None:
                    speakers = Speakers(send)
                    speakers.start()
                elif not wanted and speakers is not None:
                    speakers.stop()
                    speakers = None
            elif kind == MIC and 0 < length <= MAX_PACKET:
                microphone.push(payload)
            else:
                log(f"closing a connection that sent frame type {kind:#04x} of {length} bytes")
                return
    except OSError:
        return
    finally:
        if speakers is not None:
            speakers.stop()
        microphone.close()
        reader.close()
        conn.close()


def main():
    Gst.init(None)
    os.makedirs(SOCKET_DIR, mode=0o700, exist_ok=True)
    try:
        os.unlink(SOCKET_PATH)
    except FileNotFoundError:
        pass
    server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    server.bind(SOCKET_PATH)
    os.chmod(SOCKET_PATH, 0o600)
    server.listen()
    slots = threading.BoundedSemaphore(MAX_CONNECTIONS)
    log(f"listening on {SOCKET_PATH}")
    while True:
        conn, _ = server.accept()
        if not slots.acquire(blocking=False):
            log("refusing a connection: too many viewers")
            conn.close()
            continue

        def run(conn=conn):
            try:
                serve(conn)
            finally:
                slots.release()

        threading.Thread(target=run, daemon=True).start()


if __name__ == "__main__":
    main()
