#!/usr/bin/env python3
import argparse
import pathlib
import socket
import struct
import time


LOGIN_MAGIC = b"HLDP-GHOS-LOGIN!"
COMMIT_MAGIC = b"HLDP-CODE-COMMIT"
BRANCH_MAGIC = b"HLDP-BRANCH-MOVE"
RECOVERY_MAGIC = b"HLDP-RECOVER-OS!"


def checksum(payload: bytes) -> int:
    if len(payload) % 2:
        payload += b"\0"
    words = struct.unpack(f"!{len(payload) // 2}H", payload)
    total = sum(words)
    while total >> 16:
        total = (total & 0xFFFF) + (total >> 16)
    return (~total) & 0xFFFF


def request(sequence: int, magic: bytes) -> bytes:
    payload = b"\0" * 8 + magic + magic
    packet = struct.pack("!BBHHH", 8, 0, 0, 0x4748, sequence) + payload
    return packet[:2] + struct.pack("!H", checksum(packet)) + packet[4:]


def verify_reply(packet: bytes, magic: bytes) -> None:
    if packet and packet[0] >> 4 == 4:
        header_length = (packet[0] & 0x0F) * 4
        packet = packet[header_length:]
    if len(packet) < 40 or packet[0] != 0 or packet[1] != 0:
        raise RuntimeError("unexpected ICMP reply shape")
    if packet[16:32] != magic:
        raise RuntimeError("ICMP reply does not carry the requested HLDP magic")
    if checksum(packet) != 0:
        raise RuntimeError("ICMP reply checksum failed")


def exchange(
    peer: socket.socket,
    target: tuple[str, int],
    sequence: int,
    magic: bytes,
    deadline: float,
) -> None:
    while time.monotonic() < deadline:
        peer.sendto(request(sequence, magic), target)
        try:
            packet, _ = peer.recvfrom(4096)
        except TimeoutError:
            continue
        verify_reply(packet, magic)
        return
    raise TimeoutError(f"no verified reply for sequence {sequence}")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--target", required=True)
    parser.add_argument("--receipt", required=True)
    parser.add_argument("--timeout", type=float, default=90.0)
    parser.add_argument("--resident", action="store_true")
    parser.add_argument("--login-only", action="store_true")
    args = parser.parse_args()
    if args.resident and args.login_only:
        raise SystemExit("--resident and --login-only are mutually exclusive")

    peer = socket.socket(socket.AF_INET, socket.SOCK_DGRAM, socket.IPPROTO_ICMP)
    peer.settimeout(1.0)
    target = (socket.gethostbyname(args.target), 0)
    deadline = time.monotonic() + args.timeout
    sequence = 1
    phases: list[tuple[str, bytes]] = [
        ("login", LOGIN_MAGIC),
        ("login", LOGIN_MAGIC),
        ("login", LOGIN_MAGIC),
    ]
    if not args.login_only:
        phases.extend(
            [
                ("code_commit", COMMIT_MAGIC),
                ("branch_move", BRANCH_MAGIC),
            ]
        )
    if args.resident:
        phases.extend([("resident_login", LOGIN_MAGIC)] * 10)
        phases.append(("recovery", RECOVERY_MAGIC))

    completed: list[str] = []
    for phase, magic in phases:
        exchange(peer, target, sequence, magic, deadline)
        completed.append(phase)
        sequence += 1

    receipt = pathlib.Path(args.receipt)
    receipt.write_text(
        "schema: guanghu.physical-native-icmp-peer/v1\n"
        "status: PASS_100\n"
        f"target: {args.target}\n"
        f"login_reply_count: {completed.count('login')}\n"
        f"code_commit_reply_verified: {str('code_commit' in completed).lower()}\n"
        f"branch_move_reply_verified: {str('branch_move' in completed).lower()}\n"
        f"resident_login_reply_count: {completed.count('resident_login')}\n"
        f"recovery_reply_verified: {str('recovery' in completed).lower()}\n",
        encoding="utf-8",
    )
    print(receipt.read_text(encoding="utf-8"), end="")


if __name__ == "__main__":
    main()