# N-30 芯片级复用时序台架：主机端（与 test/board/ch395f_reuse_bench.c 共用变体表）
# 用法: python ch395f_reuse_host.py [--host IP] [--cycles N]
# 纯 TCP，无 FTP。按 idx 轮转连接 5000+idx 端口，记录收发结果与时刻。
import argparse
import socket
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 main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--host", default="192.168.1.100")
    ap.add_argument("--cycles", type=int, default=3)
    args = ap.parse_args()

    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())