#!/usr/bin/env python3
"""独立验证装甲板血量上报的诊断脚本。

不依赖 Rust，也不依赖 protocol / armor-protocol crate——常量是手写的。正因如此它
是线路格式的**独立**检查：如果 Rust 侧的编码写错了，这个脚本仍按文档解析，会立刻
报出不一致。

用法：

    python3 scripts/armor_probe.py [秒数]

它监听 UDP 5601（装甲板主板的广播目标），把收到 `ArmorHealth` 数据报按
`(team, index)` 汇总，打印血量变化与心跳间隔。板子每 200 ms 广播四条，所以
正常情况下每秒每个 `(team, index)` 大约 5 条。

用它可以确认三件事，都不需要打开裁判端界面：

  1. 板子在广播、IP 与 MAC 派生 id 是多少
  2. 四条 index 的血量是否一致（协议约定：整车一条血，四格同值）
  3. 血量有没有随受击下降（每击 −3，100 起）

常量与 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"
ARMOR_PORT = 5601     # ARMOR_PORT
MSG_ARMOR_HEALTH = 5  # MsgType::ArmorHealth

TEAM_NAMES = {0: "队伍一", 1: "队伍二"}


def decode(data):
    """解析一个数据报，返回 (team, index, hp, sender_id, version) 或 None。"""
    if len(data) < HEADER + 4:
        return None

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

    if data[4] != MSG_ARMOR_HEALTH:
        return None

    version = int.from_bytes(data[5:7], "little")
    sender_id = int.from_bytes(data[7:11], "little")
    team = data[HEADER]
    index = data[HEADER + 1]
    hp = int.from_bytes(data[HEADER + 2 : HEADER + 4], "little")

    return team, index, hp, sender_id, version


def main():
    seconds = float(sys.argv[1]) if len(sys.argv) > 1 else 10.0

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

    print(f"listening on UDP {ARMOR_PORT} for {seconds:g}s ...")

    # (team, index) -> 最后一次 (hp, 收到时刻)
    latest = {}
    # (team, index) -> 收到过的 hp 序列（去重后按时间）
    history = defaultdict(list)
    # (team, index) -> 包数
    counts = defaultdict(int)
    senders = set()
    versions = set()

    deadline = time.time() + seconds
    while time.time() < deadline:
        try:
            data, addr = sock.recvfrom(2048)
        except socket.timeout:
            continue

        parsed = decode(data)
        if parsed is None:
            continue

        team, index, hp, sender_id, version = parsed
        now = time.time()
        key = (team, index)
        counts[key] += 1
        senders.add((sender_id, f"{addr[0]}:{addr[1]}"))
        versions.add(version)

        previous = latest.get(key)
        if previous is None or previous[0] != hp:
            history[key].append((round(now, 3), hp))
            print(
                f"  {TEAM_NAMES.get(team, f'team{team}')} plate {index}: "
                f"hp {hp}"
                + (f"  (was {previous[0]})" if previous else "  (first seen)")
            )
        latest[key] = (hp, now)

    print()

    if not counts:
        print("没有收到装甲板上报。检查：")
        print("  - 板子上电了吗（UART0 上应有 'reporting to 255.255.255.255:5601'）")
        print("  - 板子和这台机器在同一个局域网、同一个子网吗")
        print("  - 防火墙放行 UDP 5601 了吗")
        return 1

    print(f"发送方 ({len(senders)}):")
    for sender_id, addr in sorted(senders):
        print(f"  0x{sender_id:08X}  {addr}")

    print(f"协议版本: {sorted(versions)}")
    print()

    by_team = defaultdict(dict)
    for (team, index), count in sorted(counts.items()):
        by_team[team][index] = count

    elapsed = max(seconds, 0.001)

    print("每个 (队伍, index) 的包数 / 心率:")
    for team, plates in sorted(by_team.items()):
        cells = []
        for index in range(4):
            n = plates.get(index, 0)
            cells.append(f"index{index}={n} ({n / elapsed:.1f}/s)")
        print(f"  {TEAM_NAMES.get(team, f'team{team}')}: " + "  ".join(cells))

    # 协议约定：整车一条血 → 四格同值；四条都在新鲜期内才等于"这条队伍在线"。
    print()
    print("四条 index 的血量:")
    for team, plates in sorted(by_team.items()):
        hps = []
        for index in range(4):
            reading = latest.get((team, index))
            hps.append(str(reading[0]) if reading else "-")
        same = len(set(hps)) == 1 and "-" not in hps
        print(
            f"  {TEAM_NAMES.get(team, f'team{team}')}: [{', '.join(hps)}]"
            + ("  <- 一致（整车一条血量）" if same else "  <- 不一致：检查 index 或丢包")
        )

    if history:
        print()
        print("血量变化:")
        for key in sorted(history):
            team, index = key
            events = history[key]
            if len(events) == 1:
                continue
            line = "  ".join(f"t={t:+.1f}s hp={hp}" for t, hp in events[:12])
            if len(events) > 12:
                line += f"  ... (+{len(events) - 12} more)"
            print(f"  {TEAM_NAMES.get(team, f'team{team}')} plate {index}: {line}")

    return 0


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