#!/usr/bin/env python3
"""Minimal UDP listener for the camera stream, used to verify the firmware.

Not the real client - just enough to confirm datagrams arrive, that they
reassemble into a valid JPEG, and to measure the frame rate independently of
the Slint app.
"""
import socket, sys, time, collections

HOST = "192.168.1.139"
PORT = 5600
DURATION = float(sys.argv[1]) if len(sys.argv) > 1 else 10.0

# Must match protocol/src/lib.rs: MAGIC, HEADER_LEN, MAX_FRAGMENT_PAYLOAD,
# DEFAULT_PORT. Duplicated here on purpose - this is a throwaway diagnostic, and
# making it depend on the Rust crate would defeat the point of it being an
# independent check of the wire format.
HEADER = 20
MAGIC = 0x4D564331
FRAG = 1400

sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
sock.bind(("", PORT))

# Announce ourselves, exactly as the Slint client will.
sock.sendto(b"\x00", (HOST, PORT))
sock.settimeout(0.5)

frag_bytes = FRAG
frames = {}
stats = collections.Counter()
sizes = []
start = time.time()
last_print = start

while time.time() - start < DURATION:
    try:
        data, _ = sock.recvfrom(2048)
    except socket.timeout:
        continue

    stats["datagrams"] += 1

    if len(data) < HEADER:
        stats["short"] += 1
        continue

    magic = int.from_bytes(data[0:4], "little")
    if magic != MAGIC:
        stats["bad_magic"] += 1
        continue

    frame_id = int.from_bytes(data[4:8], "little")
    frag_index = int.from_bytes(data[8:10], "little")
    frag_count = int.from_bytes(data[10:12], "little")
    total_len = int.from_bytes(data[12:16], "little")
    width = int.from_bytes(data[16:18], "little")
    height = int.from_bytes(data[18:20], "little")

    payload = data[HEADER:]

    if frame_id not in frames:
        if frames:
            stats["superseded"] += 1
        frames[frame_id] = {
            "buf": bytearray(total_len),
            "seen": 0,
            "count": frag_count,
            "total": total_len,
            "w": width,
            "h": height,
            "t0": time.time(),
        }

    f = frames[frame_id]
    off = frag_index * frag_bytes
    if off + len(payload) <= total_len:
        f["buf"][off:off + len(payload)] = payload
        f["seen"] += 1
    else:
        stats["overflow"] += 1

    if f["seen"] >= f["count"]:
        jpeg = bytes(f["buf"])
        if jpeg[:2] == b"\xff\xd8" and jpeg[-2:] == b"\xff\xd9":
            stats["good"] += 1
            sizes.append((f["total"], f["w"], f["h"]))
        else:
            stats["corrupt"] += 1
        del frames[frame_id]

    now = time.time()
    if now - last_print >= 2.0:
        elapsed = now - last_print
        print(
            f"  {stats['good']/elapsed:5.1f} fps  "
            f"datagrams={stats['datagrams']:5d}  "
            f"superseded={stats['superseded']:4d}  corrupt={stats['corrupt']:3d}  "
            f"good={stats['good']:4d}"
        )
        last_print = now

sock.close()
print()
print("totals:", dict(stats))
if sizes:
    avg = sum(s[0] for s in sizes) / len(sizes)
    print(f"avg frame: {avg:.0f} bytes, {sizes[0][1]}x{sizes[0][2]}")
    print(f"throughput: {sum(s[0] for s in sizes)/DURATION/1024:.1f} kB/s")