# relisten 沉降段原语化受阻：芯片级复用时序台架主机端（与 test/board/ch395f_reuse_bench.c 共用变体表）
# 用法: python ch395f_reuse_host.py [--host IP] [--cycles N] [--mode legacy|e4]
# legacy: 按 idx 轮转连接 5000+idx 端口，记录收发结果与时刻。
# e4:     E4 sock_reset A/B/C 矩阵配套（2026-10-13）：每轮两次连接——
#         ①污染：连上发 3000B 后 RST（SO_LINGER 0，TRAP-NET-18 中止形态）；
#         ②载荷：立即重连发 4096B 图案（pat[i]=(seed+i)&0xFF, seed=(rg*97+13)&0xFF）后 FIN。
#         端口 5100+rg，rg=0..(3*rounds-1)，变体顺序 A/B/C（与固件 bench 表一致）。
import argparse
import socket
import struct
import sys
import time

# {idx: (tx_len, close_mode)}  与固件表一致
VARS = {
    0: (100, 0),   # 0ms control
    1: (100, 0),   # 300ms control
    2: (2000, 0),  # 2KB control
    3: (100, 0),   # 0ms + ghost-wait fix
    4: (100, 0),   # 300ms + ghost-wait fix
    5: (2000, 0),  # 2KB + ghost-wait fix
}


def ts():
    return time.strftime("%H:%M:%S") + ".%03d" % (int(time.time() * 1000) % 1000)


def run_round(host, idx, tx_len, close_mode):
    port = 5000 + idx
    # 等端口监听（芯片还没 open 时连接被拒，重试）
    s = None
    t0 = time.time()
    while time.time() - t0 < 12.0:
        try:
            s = socket.create_connection((host, port), timeout=5)
            break
        except OSError:
            time.sleep(0.2)
    if s is None:
        print("%s r%u CONNECT timeout" % (ts(), idx))
        return False

    ok = True
    try:
        s.settimeout(5)
        got = 0
        while got < tx_len:
            chunk = s.recv(min(1024, tx_len - got))
            if not chunk:
                break
            got += len(chunk)
        note = ""
        if got != tx_len:
            note = "  <<<< SWALLOWED (want %d)" % tx_len
            ok = False
        if close_mode == 1:
            # 等芯片 FIN（mode1 芯片先关）
            try:
                fin = s.recv(1)
                if fin:
                    note += "  unexpected data after tx: %r" % fin
            except socket.timeout:
                note += "  <<<< no FIN from chip (mode1)"
                ok = False
        print("%s r%u connected got=%d/%d%s" % (ts(), idx, got, tx_len, note))
    except socket.timeout:
        print("%s r%u recv timeout  <<<< STALL" % (ts(), idx))
        ok = False
    except OSError as e:
        print("%s r%u sock error %r" % (ts(), idx, e))
        ok = False
    finally:
        s.close()  # 主机关闭（mode0/2 主机先关）
    return ok


def connect_retry(host, port, budget_s):
    """connect with retry until budget_s elapses; None on give-up"""
    t0 = time.time()
    while time.time() - t0 < budget_s:
        try:
            return socket.create_connection((host, port), timeout=5)
        except OSError:
            time.sleep(0.2)
    return None


def e4_round(host, rg, rounds_per_variant, first_budget):
    """one E4 round: contaminate (send 3000B then RST) + payload (4096B pattern then FIN)"""
    port = 5100 + rg
    seed = (rg * 97 + 13) & 0xFF
    budget = first_budget if rg == 0 else 30
    s = connect_retry(host, port, budget)
    if s is None:
        print("%s E4 r%2d port%5d  <<<< contam CONNECT timeout" % (ts(), rg, port))
        return False
    try:
        s.sendall(b"\x33" * 3000)
    except OSError as e:
        print("%s E4 r%2d contam send err %r" % (ts(), rg, e))
    try:
        s.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER, struct.pack("ii", 1, 0))
    except OSError:
        pass
    s.close()  # RST: abort mid-session like ftp_abort_probe (TRAP-NET-18 form)
    time.sleep(0.05)

    s = connect_retry(host, port, 20)
    if s is None:
        print("%s E4 r%2d port%5d  <<<< payload CONNECT timeout" % (ts(), rg, port))
        return False
    payload = bytes(((seed + i) & 0xFF) for i in range(4096))
    try:
        s.sendall(payload)
    except OSError as e:
        print("%s E4 r%2d payload send err %r" % (ts(), rg, e))
        s.close()
        return False
    s.close()  # FIN: full payload delivered, chip drains
    print("%s E4 r%2d var=%s port%5d contam+payload sent (seed=%02X)"
          % (ts(), rg, "ABC"[rg // rounds_per_variant], port, seed))
    return True


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--host", default="192.168.1.100")
    ap.add_argument("--cycles", type=int, default=3)
    ap.add_argument("--mode", choices=["legacy", "e4"], default="legacy")
    ap.add_argument("--rounds", type=int, default=10, help="e4: rounds per variant")
    ap.add_argument("--first-budget", type=int, default=300,
                    help="e4: connect budget for the first round (bench runs E2/E3 first)")
    args = ap.parse_args()

    if args.mode == "e4":
        fails = []
        total = 3 * args.rounds
        for rg in range(total):
            if not e4_round(args.host, rg, args.rounds, args.first_budget):
                fails.append(rg)
        print("E4 RESULT: %d/%d host-side ok, failures %s" % (total - len(fails), total, fails))
        return 1 if fails else 0

    fails = []
    for c in range(args.cycles):
        for idx in sorted(VARS):
            tx_len, mode = VARS[idx]
            print("%s === cycle%u r%u (tx=%d mode=%d) ===" % (ts(), c, idx, tx_len, mode))
            if not run_round(args.host, idx, tx_len, mode):
                fails.append((c, idx))
            time.sleep(0.3)
    print("RESULT: %d failures %s" % (len(fails), fails))
    return 1 if fails else 0


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