#!/usr/bin/env python3
"""Checks the TV's RTP sink against the host receiver.

The C sink writes a known ramp; this binds the port, collects the datagrams and
decodes them with the very functions host/lgtv-audiocap-receiver.py uses. If
the two ever disagree about the header layout or the sample byte order, the
ramp comes back wrong and this fails.

Also checks the SAP/SDP announcement, since PulseAudio's module-rtp-recv builds
its source purely from that text.

    python3 test/verify_rtp.py [path-to-rtp_send]
"""

import os
import re
import socket
import struct
import subprocess
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, os.pardir, "host"))

# The receiver's filename is not an identifier, so load it by path.
try:
    import importlib.util

    spec = importlib.util.spec_from_file_location(
        "audiocap_receiver",
        os.path.join(HERE, os.pardir, "host", "lgtv-audiocap-receiver.py"))
    receiver = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(receiver)
except Exception as exc:  # pragma: no cover - only when the file is missing
    print("cannot load the host receiver: %s" % exc)
    sys.exit(1)

BLOCKS = 8
FRAMES = 512
CHANNELS = 2
RATE = 48000

checks = 0
failures = 0


def check(name, condition, detail=None):
    global checks, failures
    checks += 1
    if condition:
        print("  ok   %s" % name)
    else:
        failures += 1
        print("  FAIL %s%s" % (name, "" if detail is None else " — %s" % detail))


def eq(name, actual, expected):
    check(name, actual == expected, "got %r, wanted %r" % (actual, expected))


def sample_at(index):
    """Must match sample_at() in test/rtp_send.c."""
    value = (index * 251) % 65536 - 32768
    return value


def main():
    binary = sys.argv[1] if len(sys.argv) > 1 else os.path.join(HERE, "rtp_send")
    if not os.path.exists(binary):
        print("build test/rtp_send.c first (run-tests.sh does it for you)")
        return 1

    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 1 << 20)
    sock.bind(("127.0.0.1", 0))
    port = sock.getsockname()[1]
    sock.settimeout(2.0)

    proc = subprocess.run([binary, str(port), str(BLOCKS), str(FRAMES)],
                          stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    if proc.returncode != 0:
        print("rtp_send failed: %s" % proc.stderr.decode("utf-8", "replace"))
        return 1
    expected_samples = int(proc.stdout.decode().strip())

    packets = []
    try:
        while True:
            data, _ = sock.recvfrom(4096)
            packets.append(data)
    except socket.timeout:
        pass
    sock.close()

    print("wire format")
    check("packets arrived", len(packets) > 0, "%d packets" % len(packets))
    if not packets:
        return 1

    first = receiver.parse_rtp(packets[0])
    check("receiver parses the header", first is not None)
    eq("payload type", first.payload_type, 96)
    eq("version 2, no CSRCs, no extension", packets[0][0], 0x80)
    eq("marker bit clear", packets[0][1] >> 7, 0)

    # Every packet must stay inside a 1500-byte MTU with room for the IP and
    # UDP headers, or the stream fragments and loss goes from bad to total.
    largest = max(len(p) for p in packets)
    check("no packet exceeds the MTU budget", largest <= 1472, "%d bytes" % largest)

    print("sequencing")
    parsed = [receiver.parse_rtp(p) for p in packets]
    check("all packets parse", all(p is not None for p in parsed))
    ssrcs = set(p.ssrc for p in parsed)
    eq("one SSRC for the run", len(ssrcs), 1)

    sequences = [p.sequence for p in parsed]
    expected_sequences = [(sequences[0] + i) & 0xFFFF for i in range(len(sequences))]
    eq("sequence numbers increment by one", sequences, expected_sequences)

    frame_bytes = 2 * CHANNELS
    stamps = [p.timestamp for p in parsed]
    steps = set((stamps[i + 1] - stamps[i]) & 0xFFFFFFFF for i in range(len(stamps) - 1))
    frames_per_packet = set(len(p.payload) // frame_bytes for p in parsed[:-1])
    eq("timestamp advances by the frame count", steps, frames_per_packet)

    print("payload")
    pcm = b"".join(receiver.to_native_pcm(p.payload) for p in parsed)
    samples = struct.unpack("<%dh" % (len(pcm) // 2), pcm)
    eq("every sample arrived", len(samples), expected_samples)
    wrong = [i for i, v in enumerate(samples) if v != sample_at(i)]
    check("the ramp survives the round trip", not wrong,
          "%d samples differ, first at %s" % (len(wrong), wrong[:1]))

    # A wrong byte order still produces "audio", just noise; check explicitly
    # that the payload really is big-endian on the wire.
    raw_be = struct.unpack(">%dh" % (len(parsed[0].payload) // 2), parsed[0].payload)
    eq("payload is big-endian on the wire", raw_be[0], sample_at(0))

    print("SAP announcement")
    sdp = capture_sap()
    if sdp is None:
        print("  SKIP: no announcement seen (multicast on loopback is often"
              " blocked); the SDP text itself is unchecked")
    else:
        check("SDP names an L16 stream", "L16/%d/%d" % (RATE, CHANNELS) in sdp, sdp)
        check("SDP carries a media line", re.search(r"m=audio \d+ RTP/AVP 96", sdp)
              is not None, sdp)
        check("SDP is recvonly", "a=recvonly" in sdp, sdp)

    print("\n%d/%d checks passed" % (checks - failures, checks))
    return 1 if failures else 0


def capture_sap(timeout=1.5):
    """Listens for one SAP announcement from a second, SAP-enabled run."""
    binary = sys.argv[1] if len(sys.argv) > 1 else os.path.join(HERE, "rtp_send")
    sap = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sap.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    try:
        sap.bind(("", 9875))
        # Join on whichever interface the kernel picks: the announcement leaves
        # by the default route, so that is where it can loop back from.
        membership = socket.inet_aton("224.0.0.56") + struct.pack("=I", socket.INADDR_ANY)
        sap.setsockopt(socket.IPPROTO_IP, socket.IP_ADD_MEMBERSHIP, membership)
    except OSError:
        sap.close()
        return None
    sap.settimeout(timeout)

    subprocess.run([binary, "9999", "2", "512", "sap"], stdout=subprocess.DEVNULL,
                   stderr=subprocess.DEVNULL)
    try:
        data, _ = sap.recvfrom(2048)
    except socket.timeout:
        return None
    finally:
        sap.close()

    # RFC 2974: 4-byte header, 4-byte source, NUL-terminated MIME type.
    body = data[8:]
    end = body.find(b"\x00")
    return body[end + 1:].decode("utf-8", "replace") if end >= 0 else None


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