#!/usr/bin/env python3
"""Receives the TV's RTP/L16 audio and plays it into a local sound device.

This is the piece that makes HyperHDR work. HyperHDR's sound-reactive effects
read a *local* capture device, so the TV's audio has to arrive as one. This
script takes the RTP stream and writes it to a playback device; pair it with
a loopback (host/install-loopback.sh) and HyperHDR sees a normal input.

    ./lgtv-audiocap-receiver.py                       # auto-detect an output
    ./lgtv-audiocap-receiver.py --output pacat --device lgtv-audio
    ./lgtv-audiocap-receiver.py --output aplay --device hw:Loopback,0
    ./lgtv-audiocap-receiver.py --output - > /tmp/tv.raw

Standard library only, so it runs on anything with Python 3.6 and either
PulseAudio/PipeWire (pacat) or ALSA (aplay) installed.
"""

import argparse
import array
import errno
import os
import shutil
import signal
import socket
import struct
import subprocess
import sys
import time

RTP_HEADER_BYTES = 12
RTP_PAYLOAD_TYPE = 96  # matches the TV sink
DEFAULT_PORT = 5004


def log(message):
    sys.stderr.write(message + "\n")
    sys.stderr.flush()


# ---------------------------------------------------------------------------
# RTP
# ---------------------------------------------------------------------------

class RtpPacket(object):
    __slots__ = ("sequence", "timestamp", "ssrc", "payload", "payload_type")

    def __init__(self, sequence, timestamp, ssrc, payload_type, payload):
        self.sequence = sequence
        self.timestamp = timestamp
        self.ssrc = ssrc
        self.payload_type = payload_type
        self.payload = payload


def parse_rtp(data):
    """Returns an RtpPacket, or None if this is not RTP we can use."""
    if len(data) < RTP_HEADER_BYTES:
        return None

    byte0, byte1, sequence, timestamp, ssrc = struct.unpack("!BBHII",
                                                            data[:RTP_HEADER_BYTES])
    if (byte0 >> 6) != 2:  # version
        return None

    offset = RTP_HEADER_BYTES + (byte0 & 0x0F) * 4  # CSRC list
    if byte0 & 0x10:  # extension header
        if len(data) < offset + 4:
            return None
        ext_words = struct.unpack("!H", data[offset + 2:offset + 4])[0]
        offset += 4 + ext_words * 4

    end = len(data)
    if byte0 & 0x20:  # padding: the last byte counts the padding bytes
        pad = data[-1] if isinstance(data[-1], int) else ord(data[-1])
        if pad and pad <= end - offset:
            end -= pad

    if offset >= end:
        return None

    return RtpPacket(sequence, timestamp, ssrc, byte1 & 0x7F, data[offset:end])


def to_native_pcm(payload):
    """L16 is big-endian; sound devices want the host's order."""
    samples = array.array("h")
    if len(payload) % 2:
        payload = payload[:-1]
    samples.frombytes(payload)
    if sys.byteorder == "little":
        samples.byteswap()
    return samples.tobytes()


# ---------------------------------------------------------------------------
# Output
# ---------------------------------------------------------------------------

def detect_output():
    """Pick a player. PulseAudio/PipeWire first: it needs no card set up."""
    if shutil.which("pacat"):
        try:
            subprocess.check_output(["pactl", "info"], stderr=subprocess.DEVNULL,
                                    timeout=3)
            return "pacat"
        except Exception:
            pass
    if shutil.which("aplay"):
        return "aplay"
    if shutil.which("pacat"):
        return "pacat"
    if shutil.which("ffplay"):
        return "ffplay"
    return "-"


def build_command(kind, device, rate, channels, latency_ms):
    if kind == "pacat":
        cmd = ["pacat", "--playback", "--format=s16le",
               "--rate=%d" % rate, "--channels=%d" % channels,
               "--stream-name=LG TV Audio Cap",
               "--latency-msec=%d" % latency_ms]
        if device:
            cmd += ["--device=%s" % device]
        return cmd

    if kind == "aplay":
        cmd = ["aplay", "-t", "raw", "-f", "S16_LE",
               "-r", str(rate), "-c", str(channels), "-q",
               # aplay's default buffer is far larger than we want in a chain
               # that already has a jitter buffer in front of it.
               "--buffer-time=%d" % (latency_ms * 1000)]
        if device:
            cmd += ["-D", device]
        return cmd

    if kind == "ffplay":
        return ["ffplay", "-hide_banner", "-loglevel", "error", "-nodisp",
                "-autoexit", "-f", "s16le", "-ar", str(rate),
                "-ac", str(channels), "-i", "pipe:0"]

    raise ValueError("unknown output %r" % kind)


class Output(object):
    """Where the audio goes. Restarts the player if it dies."""

    def __init__(self, kind, device, rate, channels, latency_ms):
        self.kind = kind
        self.device = device
        self.rate = rate
        self.channels = channels
        self.latency_ms = latency_ms
        self.process = None
        self.stream = None
        self.restarts = 0
        self._open()

    def _open(self):
        if self.kind == "-":
            self.stream = getattr(sys.stdout, "buffer", sys.stdout)
            return
        cmd = build_command(self.kind, self.device, self.rate, self.channels,
                            self.latency_ms)
        log("playing into: %s" % " ".join(cmd))
        self.process = subprocess.Popen(cmd, stdin=subprocess.PIPE)
        self.stream = self.process.stdin

    def write(self, data):
        try:
            self.stream.write(data)
            self.stream.flush()
            return True
        except (IOError, OSError, ValueError) as exc:
            if getattr(exc, "errno", None) == errno.EINTR:
                return True
            if self.kind == "-":
                raise
            log("output died (%s); restarting" % exc)
            self.close()
            self.restarts += 1
            time.sleep(0.5)
            self._open()
            return False

    def close(self):
        if self.process:
            try:
                if self.process.stdin:
                    self.process.stdin.close()
            except Exception:
                pass
            try:
                self.process.terminate()
                self.process.wait(timeout=2)
            except Exception:
                pass
            self.process = None
        self.stream = None


# ---------------------------------------------------------------------------
# Receiver
# ---------------------------------------------------------------------------

def open_socket(bind, port, group, iface, rcvbuf):
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    if hasattr(socket, "SO_REUSEPORT"):
        try:
            sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
        except OSError:
            pass
    try:
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, rcvbuf)
    except OSError:
        pass

    sock.bind(("" if group else bind, port))

    if group:
        # Joining on the wildcard interface lets the kernel choose; an explicit
        # one is needed on hosts with several networks.
        local = socket.inet_aton(iface) if iface else struct.pack("=I", socket.INADDR_ANY)
        sock.setsockopt(socket.IPPROTO_IP, socket.IP_ADD_MEMBERSHIP,
                        socket.inet_aton(group) + local)
        log("joined multicast group %s" % group)

    return sock


def run(args):
    frame_bytes = 2 * args.channels
    sock = open_socket(args.bind, args.port, args.multicast, args.iface,
                       args.rcvbuf)
    sock.settimeout(0.2)

    out = Output(args.output, args.device, args.rate, args.channels,
                 args.latency_ms)

    # A block of silence sized to roughly one packet, reused for gap filling.
    gap_frames = max(1, int(args.rate * 0.01))
    silence = b"\x00" * (gap_frames * frame_bytes)

    # Prime the device so the first real packet is not chasing an empty buffer.
    prime = int(args.rate * args.prebuffer_ms / 1000.0)
    if prime and args.output != "-":
        out.write(b"\x00" * (prime * frame_bytes))

    stats = {"packets": 0, "bytes": 0, "lost": 0, "late": 0, "resets": 0}
    expected = None
    ssrc = None
    last_packet = time.time()
    last_stats = time.time()
    running = [True]

    def stop(signum, frame):
        running[0] = False

    signal.signal(signal.SIGINT, stop)
    signal.signal(signal.SIGTERM, stop)

    log("listening on %s:%d for %d Hz %d-channel L16"
        % (args.multicast or args.bind, args.port, args.rate, args.channels))

    while running[0]:
        try:
            data, sender = sock.recvfrom(4096)
        except socket.timeout:
            now = time.time()
            # Keep the device fed while the TV is quiet or gone, otherwise the
            # player underruns and HyperHDR's effect freezes on the last frame
            # instead of fading out.
            if args.fill_silence and args.output != "-" and expected is not None:
                out.write(silence)
            if expected is not None and now - last_packet > args.reset_after:
                log("no audio for %.0f s; waiting for the stream to come back"
                    % args.reset_after)
                expected = None
            continue
        except OSError as exc:
            if exc.errno == errno.EINTR:
                continue
            raise

        packet = parse_rtp(data)
        if packet is None or packet.payload_type != args.payload_type:
            continue

        if ssrc is None or packet.ssrc != ssrc:
            if ssrc is not None:
                log("stream restarted (new SSRC from %s)" % sender[0])
                stats["resets"] += 1
            else:
                log("stream started from %s" % sender[0])
            ssrc = packet.ssrc
            expected = packet.sequence

        # 16-bit sequence numbers wrap; compare in that space.
        delta = (packet.sequence - expected) & 0xFFFF
        if delta == 0:
            pass
        elif delta < args.max_gap:
            # Lost packets. Substitute silence so playback keeps its timing
            # rather than jumping forward.
            missing = delta
            stats["lost"] += missing
            payload_frames = len(packet.payload) // frame_bytes
            if payload_frames:
                out.write(b"\x00" * (payload_frames * frame_bytes) * missing)
        else:
            # Either a very late packet or a huge jump. Late ones would play
            # out of order, so drop them and resynchronise on a jump.
            if delta > 0xFFFF - args.max_gap:
                stats["late"] += 1
                continue
            log("sequence jumped by %d; resynchronising" % delta)
            expected = packet.sequence

        out.write(to_native_pcm(packet.payload))
        expected = (packet.sequence + 1) & 0xFFFF
        stats["packets"] += 1
        stats["bytes"] += len(packet.payload)
        last_packet = time.time()

        if args.stats and last_packet - last_stats >= args.stats:
            last_stats = last_packet
            seconds = stats["bytes"] / float(frame_bytes * args.rate)
            log("%d packets, %.1f s audio, %d lost, %d late, %d restarts"
                % (stats["packets"], seconds, stats["lost"], stats["late"],
                   stats["resets"] + out.restarts))

    log("stopping")
    out.close()
    sock.close()
    return 0


def main():
    parser = argparse.ArgumentParser(
        description=__doc__.split("\n")[0],
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__[__doc__.index("This is the piece"):])

    parser.add_argument("--port", type=int, default=DEFAULT_PORT,
                        help="UDP port to listen on (default %d)" % DEFAULT_PORT)
    parser.add_argument("--bind", default="0.0.0.0",
                        help="address to bind (default all interfaces)")
    parser.add_argument("--multicast", default=None,
                        help="multicast group to join, if the TV sends to one")
    parser.add_argument("--iface", default=None,
                        help="local address to join the multicast group on")

    parser.add_argument("--rate", type=int, default=48000,
                        help="sample rate the TV is sending (default 48000)")
    parser.add_argument("--channels", type=int, default=2,
                        help="channel count the TV is sending (default 2)")
    parser.add_argument("--payload-type", type=int, default=RTP_PAYLOAD_TYPE,
                        help="RTP payload type to accept (default %d)"
                             % RTP_PAYLOAD_TYPE)

    parser.add_argument("--output", default="auto",
                        choices=["auto", "pacat", "aplay", "ffplay", "-"],
                        help="how to play the audio; '-' writes raw PCM to stdout")
    parser.add_argument("--device", default=None,
                        help="sink or PCM to play into, e.g. lgtv-audio or hw:Loopback,0")

    parser.add_argument("--latency-ms", type=int, default=80,
                        help="playback buffer to ask the device for (default 80)")
    parser.add_argument("--prebuffer-ms", type=int, default=60,
                        help="silence written before the first packet (default 60)")
    parser.add_argument("--max-gap", type=int, default=200,
                        help="packets of loss to paper over before resynchronising")
    parser.add_argument("--reset-after", type=float, default=5.0,
                        help="seconds of silence before the stream is considered gone")
    parser.add_argument("--no-fill-silence", dest="fill_silence",
                        action="store_false",
                        help="do not write silence while no packets arrive")
    parser.add_argument("--rcvbuf", type=int, default=1 << 20,
                        help="socket receive buffer in bytes")
    parser.add_argument("--stats", type=float, default=30.0,
                        help="seconds between statistics lines, 0 to disable")

    args = parser.parse_args()

    if args.output == "auto":
        args.output = detect_output()
        if args.output == "-":
            log("no pacat, aplay or ffplay found; writing raw PCM to stdout")

    if args.output == "-" and os.isatty(sys.stdout.fileno()):
        parser.error("refusing to write raw PCM to a terminal; redirect stdout")

    return run(args)


if __name__ == "__main__":
    sys.exit(main())