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

A copy of the desktop image's desktop-audio.py. 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

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
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,rate=48000,channels=2"
    " ! opusenc bitrate=96000 frame-size=20 audio-type=generic"
    " ! appsink name=out sync=false max-buffers=25 drop=true"
)

# 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


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()
            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()
                    self.send(SPEAKERS, buffer.extract_dup(0, buffer.get_size()))
            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()
