#!/usr/bin/env python3
"""独立验证相机 UDP 图传线路的诊断脚本。

不依赖 Rust，也不依赖 protocol crate——正因为如此，它能当作线路格式的
**独立**检查：如果 protocol crate 的编码有错，这个脚本（手写的常量）仍然
按文档解析，会立刻报出不一致。

用法：

    python3 scripts/udp_probe.py [秒数]

它做三件事：
  1. 广播一次 Subscribe，触发板子开始推流
  2. 解析收到的数据报，按 sender_id 重组
  3. 报告帧率、丢帧、以及是否发现了板子的 Hello

常量是手写的，与 protocol/src/lib.rs 保持一致。改了协议记得改这里——
这正是它作为交叉检查的价值所在。
"""
import socket
import sys
import time
from collections import defaultdict

# ---- 与 protocol/src/lib.rs 一致 ----
HEADER = 27          # magic4 + type1 + ver2 + sender4 + frame4 + idx2 + cnt2 + len4 + w2 + h2
MAGIC = 0x4D564331   # "MVC1"
FRAG = 1400          # MAX_FRAGMENT_PAYLOAD
PORT = 5600          # DEFAULT_PORT
BROADCAST = "255.255.255.255"

MSG_VIDEO = 1
MSG_SUBSCRIBE = 2
MSG_HELLO = 3
MSG_HEARTBEAT = 4

TYPE_NAMES = {
    MSG_VIDEO: "video",
    MSG_SUBSCRIBE: "subscribe",
    MSG_HELLO: "hello",
    MSG_HEARTBEAT: "heartbeat",
}


def decode(data):
    """把一个数据报解析成 (msg_type, header_dict, payload)。"""
    if len(data) < HEADER:
        return None

    magic = int.from_bytes(data[0:4], "little")
    if magic != MAGIC:
        return None

    msg_type = data[4]
    header = {
        "version": int.from_bytes(data[5:7], "little"),
        "sender_id": int.from_bytes(data[7:11], "little"),
        "frame_id": int.from_bytes(data[11:15], "little"),
        "frag_index": int.from_bytes(data[15:17], "little"),
        "frag_count": int.from_bytes(data[17:19], "little"),
        "total_len": int.from_bytes(data[19:23], "little"),
        "width": int.from_bytes(data[23:25], "little"),
        "height": int.from_bytes(data[25:27], "little"),
    }
    return msg_type, header, data[HEADER:]


def main():
    duration = float(sys.argv[1]) if len(sys.argv) > 1 else 10.0
    host = sys.argv[2] if len(sys.argv) > 2 else BROADCAST

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

    # 订阅：板子看到它就开始往本机推流。
    #
    # role 字节紧跟在 HEADER 之后，也就是偏移 HEADER 处。之前这里多填了 12 个
    # 零字节，role 落到了偏移 39——板子的 Subscribe::decode 严格按偏移读，于是
    # 认不出 role 而拒收，表现为"板子在线但 0 clients subscribed"。
    # 这里刻意手写而不是复用 protocol crate：这个脚本的意义就在于独立验证，
    # role 的位置写错时它会给出和固件一致的结论，而不是把错误藏起来。
    header = (
        MAGIC.to_bytes(4, "little")
        + bytes([MSG_SUBSCRIBE])
        + (1).to_bytes(2, "little")           # version
        + (0x0C110001).to_bytes(4, "little")  # client_id
        + b"\x00" * 4                          # frame_id
        + b"\x00" * 4                          # frag_index + frag_count
        + b"\x00" * 4                          # total_len
        + b"\x00" * 4                          # width + height
    )
    sub = header + bytes([2])                  # role = referee，紧跟 27 字节头
    sock.sendto(sub, (host, PORT))
    print(f"已向 {host}:{PORT} 广播订阅（role=referee），监听 {duration:.0f} 秒…")

    # 每个发送方一份重组状态：两块板的 frame_id 各自独立，混在一起会互相顶掉。
    pending = {}
    boards = {}
    counters = defaultdict(int)
    sizes = []
    start = time.time()
    last_print = start

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

        parsed = decode(data)
        if parsed is None:
            counters["无法解析"] += 1
            continue

        msg_type, h, payload = parsed
        counters[f"type:{TYPE_NAMES.get(msg_type, msg_type)}"] += 1

        if msg_type in (MSG_HELLO, MSG_HEARTBEAT):
            boards[h["sender_id"]] = addr
            continue

        if msg_type != MSG_VIDEO:
            continue

        sid = h["sender_id"]

        # 与 Fragment::decode 相同的自洽性检查。
        if h["frag_count"] == 0 or h["frag_index"] >= h["frag_count"]:
            counters["分片索引非法"] += 1
            continue
        if len(payload) > h["total_len"]:
            counters["载荷超过总长"] += 1
            continue

        state = pending.setdefault(
            sid,
            {"buf": bytearray(h["total_len"]), "seen": 0, "expect": h["frag_count"],
             "have": set(), "total": h["total_len"], "w": h["width"], "h": h["height"],
             "started": time.time()},
        )

        # 重复分片直接跳过且不计数：否则一个丢失的分片会被副本凑数，
        # seen 达标而缓冲区留空洞，交出去一张尾部是 0 的坏图。
        if h["frag_index"] in state["have"]:
            counters["重复分片"] += 1
            continue
        state["have"].add(h["frag_index"])

        off = h["frag_index"] * FRAG
        state["buf"][off:off + len(payload)] = payload
        state["seen"] += 1

        tail = (h["total_len"] - 1) // FRAG
        if state["seen"] >= state["expect"] and tail in state["have"]:
            jpeg = bytes(state["buf"][:state["total"]])
            if jpeg[:2] == b"\xff\xd8" and jpeg[-2:] == b"\xff\xd9":
                counters["完整帧"] += 1
                sizes.append((state["total"], state["w"], state["h"]))
            else:
                counters["JPEG 不完整"] += 1
            del pending[sid]

        now = time.time()
        if now - last_print >= 2.0:
            elapsed = now - last_print
            good = counters["完整帧"]
            print(
                f"  {good / elapsed:5.1f} fps  "
                f"视频包={counters['type:video']:5d}  "
                f"完整帧={good:4d}  丢弃={counters['JPEG 不完整']:3d}"
            )
            last_print = now

    sock.close()

    total_time = time.time() - start
    print()
    print("板子:", ", ".join(f"0x{k:08X} @ {v[0]}" for k, v in boards.items()) or "（未发现）")
    print("完整帧:", counters["完整帧"], " 累计视频包:", counters["type:video"])
    if counters["完整帧"]:
        avg = sum(s[0] for s in sizes) / len(sizes)
        print(f"平均帧: {avg:.0f} 字节  {sizes[0][1]}x{sizes[0][2]}")
        print(f"吞吐: {sum(s[0] for s in sizes) / total_time / 1024:.1f} kB/s")
    bad = {k: v for k, v in counters.items() if k.startswith(("无法", "JPEG", "分片", "载荷"))}
    print("异常:", bad or "无")


if __name__ == "__main__":
    main()