#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
CH395F BSD Socket API 测试程序
用于测试 STM32F4 + CH395F 的网络功能

使用方法:
    python ch395f_socket_test.py [模式] [选项]

模式:
    tcp_server   - TCP 回显服务器，多连接循环（测试 CH395F TCP Client）
    tcp_client   - 作为 TCP 客户端（测试 CH395F TCP Server）
    udp_server   - 作为 UDP 服务器（测试 CH395F UDP Client）
    udp_client   - 作为 UDP 客户端（测试 CH395F UDP Server）
    tcp_echo     - TCP 回显服务器（单连接，测试双向通信）
    udp_echo     - UDP 回显服务器（测试双向通信）
    stress       - 压力测试（多连接并发测试）
    tcp_bulk     - 双向大块：先收满 bulk 字节（不回声，校验 LCG），再回放（阶段 5）
    tcp_file     - 文件收发：PC 作 Client 连 MCU:8080（阶段 8），--op send/recv 收发 100KB 文件（含 CRC32）

选项:
    --ip         目标 IP 地址（默认: 192.168.1.100）
    --port       端口号（默认: 8080）
    --count      测试次数（默认: 10）
    --size       数据包大小（默认: 1024）
    --bulk       tcp_bulk 模式双向字节数（默认: 262144）
    --chunk      tcp_bulk 模式块大小，须与固件 PHASE5_CHUNK_SIZE 一致（默认: 4096）

示例:
    # 测试 CH395F TCP Server
    python ch395f_socket_test.py tcp_client --ip 192.168.1.100 --port 8080

    # 测试 CH395F TCP Client
    python ch395f_socket_test.py tcp_server --port 8080

    # 测试 CH395F UDP Client
    python ch395f_socket_test.py udp_server --port 8888

    # 测试 CH395F UDP Server
    python ch395f_socket_test.py udp_client --ip 192.168.1.100 --port 8888

    # 压力测试（多连接）
    python ch395f_socket_test.py stress --ip 192.168.1.100 --port 8080 --count 7

     # 阶段 5 双向大块（先收 256KB 再回放 256KB）
     python ch395f_socket_test.py tcp_bulk --port 8081 --bulk 262144 --chunk 4096

     # 阶段 8 文件收发（PC 连 MCU:8080）
     python ch395f_socket_test.py tcp_file --ip 192.168.1.100 --port 8080 --op recv --file rx_100k.bin --size 102400
     python ch395f_socket_test.py tcp_file --ip 192.168.1.100 --port 8080 --op send --file tx_100k.bin
     python ch395f_socket_test.py tcp_file --ip 192.168.1.100 --port 8080 --op echo   # TC-NET-803 回显（MCU 为 server）
"""

import socket
import sys
import os
import time
import argparse
import threading
import struct
import random
import string
import zlib
from datetime import datetime

# TCP 回显服务器（tcp_server 模式）参数
INTER_ACCEPT_TIMEOUT = 60   # 首个连接后，等待下一个连接的超时秒数
ECHO_IDLE_TIMEOUT = 30      # 连接内无数据的空闲超时秒数


class Colors:
    """终端颜色"""
    HEADER = '\033[95m'
    BLUE = '\033[94m'
    CYAN = '\033[96m'
    GREEN = '\033[92m'
    YELLOW = '\033[93m'
    RED = '\033[91m'
    END = '\033[0m'
    BOLD = '\033[1m'


def log_msg(msg, color=Colors.END):
    """带时间戳的日志输出"""
    timestamp = datetime.now().strftime("%H:%M:%S.%f")[:-3]
    print(f"{Colors.CYAN}[{timestamp}]{Colors.END} {color}{msg}{Colors.END}")


def log_success(msg):
    log_msg("[OK] %s" % msg, Colors.GREEN)


def log_error(msg):
    log_msg("[ERR] %s" % msg, Colors.RED)


def log_info(msg):
    log_msg("[..] %s" % msg, Colors.BLUE)


def log_warn(msg):
    log_msg("[!!] %s" % msg, Colors.YELLOW)


def generate_data(size):
    """生成测试数据"""
    return ''.join(random.choices(string.ascii_letters + string.digits, k=size)).encode()


def gen_pattern(size, offset=0):
    """生成递增模式数据 (i+offset)&0xFF，与固件 §7 自检回显期望字节流一致"""
    return bytes(((offset + i) & 0xFF) for i in range(size))


def verify_data(data, expected_len):
    """验证数据"""
    if len(data) != expected_len:
        return False, f"长度不匹配: 期望 {expected_len}, 实际 {len(data)}"
    return True, "OK"


# ==================== TCP 测试函数 ====================

def _echo_conn_handler(conn_idx, client_socket, counters, lock):
    """每个连接一个线程：收到什么回什么，直到对端关闭或空闲超时。

    线程安全：通过 lock 更新共享计数器 counters（dict: total/bytes/errors）。
    """
    rx = 0
    client_socket.settimeout(ECHO_IDLE_TIMEOUT)
    try:
        while True:
            data = client_socket.recv(4096)
            if not data:
                log_info(f"[连接{conn_idx}] 对端已关闭")
                break
            client_socket.sendall(data)
            rx += len(data)
    except socket.timeout:
        log_warn(f"[连接{conn_idx}] 空闲超时，关闭连接")
    except Exception as e:
        log_error(f"[连接{conn_idx}] 回显错误: {e}")
        with lock:
            counters['errors'] += 1
    finally:
        client_socket.close()
        with lock:
            counters['active'] = max(0, counters['active'] - 1)
            counters['total'] += 1
            counters['bytes'] += rx
            if rx == 0:
                counters['errors'] += 1
                log_warn(f"连接#{conn_idx} 未收到任何数据")
            else:
                log_success(f"连接#{conn_idx} 完成: 回显 {rx} 字节")


def tcp_server_test(args):
    """TCP 回显服务器测试（测试 CH395F TCP Client）

    与固件 TC701/702/704 配套：MCU 连接后发送数据并等待原样回显。
    每个连接由独立线程处理，支持多 Socket 并发（TC704 一次打开 8 个）。

    首个连接等待 args.timeout 秒；之后每个连接间最多等待 INTER_ACCEPT_TIMEOUT 秒；
    若指定 --max-conn（>0），达到连接数上限后停止接受并等待全部线程结束。
    无新连接则结束并汇总。返回 False 表示失败（无连接或回显出错）。
    """
    log_info("=" * 60)
    log_info("TCP 回显服务器 - 测试 CH395F TCP Client 功能（并发）")
    log_info("=" * 60)

    server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

    counters = {'total': 0, 'bytes': 0, 'errors': 0, 'active': 0}
    lock = threading.Lock()
    threads = []

    try:
        server_socket.bind(('0.0.0.0', args.port))
        server_socket.listen(16)
        log_success(f"服务器启动，监听端口 {args.port}（并发回显，backlog=16）")

        conn_idx = 0
        while True:
            if args.max_conn and args.max_conn > 0 and counters['active'] >= args.max_conn:
                log_info(f"已达 --max-conn {args.max_conn}（并发），停止接受新连接")
                break
            wait = args.timeout if conn_idx == 0 else min(args.timeout, INTER_ACCEPT_TIMEOUT)
            server_socket.settimeout(wait)
            try:
                log_info(f"等待第 {conn_idx+1} 个连接（超时 {wait}s）...")
                client_socket, client_addr = server_socket.accept()
            except socket.timeout:
                if conn_idx == 0:
                    log_error("等待连接超时")
                else:
                    log_info(f"{wait}s 内无新连接，服务器退出")
                break

            conn_idx += 1
            with lock:
                counters['active'] += 1
            log_success(f"连接#{conn_idx} 来自: {client_addr[0]}:{client_addr[1]}")
            t = threading.Thread(target=_echo_conn_handler,
                                 args=(conn_idx, client_socket, counters, lock),
                                 daemon=True)
            t.start()
            threads.append(t)

        # 等待所有连接线程结束（对端关闭或空闲超时）
        for t in threads:
            t.join()

        log_info("-" * 40)
        log_info(f"服务结束: 连接 {counters['total']} 个, 回显 {counters['bytes']} 字节, 异常 {counters['errors']} 次")
        if counters['total'] == 0:
            log_error("无任何客户端连接")
            return False
        return counters['errors'] == 0

    except Exception as e:
        log_error(f"服务器错误: {e}")
        return False
    finally:
        server_socket.close()


# ==================== 阶段 5：双向大块传输 ====================

BULK_IDLE_TIMEOUT = 10   # bulk 接收阶段单次 recv 等待超时（秒）


def _lcg_block(idx, size):
    """生成与固件 phase5_gen_block 完全一致的 LCG 块数据（uint32 回绕）"""
    seed = (0xC2D2E6F9 ^ (idx * 2654435761)) & 0xFFFFFFFF
    out = bytearray(size)
    for k in range(size):
        seed = (seed * 1664525 + 1013904223) & 0xFFFFFFFF
        out[k] = (seed >> 24) & 0xFF
    return bytes(out)


def tcp_bulk_test(args):
    """双向大块传输测试（配合固件阶段 5 硬件层）

    协议：接受连接 → 收满 bulk 字节（不回声，校验 LCG 模式）→
          按相同 LCG 模式回放 bulk 字节 → 关闭。
    MCU 侧先连续发送 256KB（本端只收），再连续接收 256KB（本端回放），
    分别测得纯 TX / 纯 RX 吞吐。
    """
    chunk = args.chunk
    bulk = (args.bulk // chunk) * chunk        # 按整块取整
    blocks = bulk // chunk
    log_info("=" * 60)
    log_info("TCP Bulk 双向大块 - 配合 CH395F 阶段 5")
    log_info(f"协议: 收满 {bulk} 字节(不回声) -> 回放 {bulk} 字节 (chunk={chunk}, blocks={blocks})")
    log_info("=" * 60)

    server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    try:
        server_socket.bind(('0.0.0.0', args.port))
        server_socket.listen(4)
        log_success(f"服务器启动，监听端口 {args.port}")

        server_socket.settimeout(args.timeout)
        try:
            client_socket, client_addr = server_socket.accept()
        except socket.timeout:
            log_error("等待连接超时")
            return False
        log_success(f"连接来自: {client_addr[0]}:{client_addr[1]}")
        client_socket.settimeout(BULK_IDLE_TIMEOUT)

        # 预生成期望/回放数据（LCG，与固件一致）
        expected = bytearray(bulk)
        for i in range(blocks):
            expected[i * chunk:(i + 1) * chunk] = _lcg_block(i, chunk)

        # ---- 阶段 A：收满 bulk 字节（不回声），校验 LCG 模式 ----
        rx_buf = bytearray()
        rx_t0 = time.time()
        try:
            while len(rx_buf) < bulk:
                data = client_socket.recv(min(65536, bulk - len(rx_buf)))
                if not data:
                    log_error(f"对端提前关闭，仅收到 {len(rx_buf)}/{bulk} 字节")
                    client_socket.close()
                    return False
                rx_buf += data
        except socket.timeout:
            log_error(f"接收空闲超时，仅收到 {len(rx_buf)}/{bulk} 字节")
            client_socket.close()
            return False
        rx_ms = int((time.time() - rx_t0) * 1000)

        ok = (bytes(rx_buf) == bytes(expected))
        if ok:
            log_success(f"阶段A 收满 {len(rx_buf)} 字节，LCG 校验一致")
        else:
            bad = -1
            for i in range(min(len(rx_buf), len(expected))):
                if rx_buf[i] != expected[i]:
                    bad = i
                    break
            if bad >= 0:
                log_error(f"阶段A 数据不一致 @第 {bad} 字节 (期望 0x{expected[bad]:02X}, "
                          f"实际 0x{rx_buf[bad]:02X})")
            else:
                log_error(f"阶段A 数据不一致（长度 {len(rx_buf)} vs {len(expected)}）")
        log_info(f"阶段A 接收 {len(rx_buf)} 字节，耗时 {rx_ms} ms "
                 f"({len(rx_buf) * 1000 // max(rx_ms, 1)} B/s)")

        # ---- 阶段 B：按 LCG 模式回放 bulk 字节 ----
        tx_t0 = time.time()
        try:
            for off in range(0, bulk, chunk):
                client_socket.sendall(expected[off:off + chunk])
        except Exception as e:
            log_error(f"阶段B 回放错误: {e}")
            ok = False
        finally:
            client_socket.close()
        tx_ms = int((time.time() - tx_t0) * 1000)
        if ok:
            log_success(f"阶段B 回放 {bulk} 字节完成，耗时 {tx_ms} ms "
                        f"({bulk * 1000 // max(tx_ms, 1)} B/s)")
        return ok

    except Exception as e:
        log_error(f"服务器错误: {e}")
        return False
    finally:
        server_socket.close()


def tcp_client_test(args):
    """TCP 客户端测试（测试 CH395F TCP Server）"""
    log_info("=" * 60)
    log_info("TCP 客户端测试 - 测试 CH395F TCP Server 功能")
    log_info("=" * 60)

    client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    client_socket.settimeout(10)

    try:
        # 连接服务器
        log_info(f"正在连接 CH395F TCP Server: {args.ip}:{args.port}")
        client_socket.connect((args.ip, args.port))
        log_success("连接成功!")

        # 测试发送数据（TC-NET-602: 64B 小包回显 ×args.count 轮）
        log_info("-" * 40)
        log_info("测试 1: 64B 小包回显 (TC-NET-602)")
        pkt_offset = 0  # 单连接内期望字节流连续累计，与固件 s_phase6_off 同步
        for i in range(args.count):
            try:
                test_data = gen_pattern(64, pkt_offset)
                pkt_offset += 64
                client_socket.send(test_data)
                log_success(f"发送数据 ({len(test_data)} bytes)")
                log_info(f"发送 hex: {' '.join(f'{b:02X}' for b in test_data[:16])}")

                # 等待回显
                data = client_socket.recv(64)
                if data:
                    log_info(f"接收 hex: {' '.join(f'{b:02X}' for b in data[:16])}")
                    log_info(f"接收长度: {len(data)}, 期望长度: {len(test_data)}")
                    if data == test_data:
                        log_success("回显数据匹配!")
                    else:
                        log_error("回显数据不匹配!")
                        for j in range(min(len(data), len(test_data))):
                            if data[j] != test_data[j]:
                                log_error(f"首个差异: 位置 {j}, 发送={test_data[j]:02X}, 接收={data[j]:02X}")
                                break
                else:
                    log_error("收到空数据")
                    break

                time.sleep(0.1)
            except Exception as e:
                log_error(f"发送/接收错误: {e}")
                break

        # 测试 1460B MSS 回显（TC-NET-603）
        log_info("-" * 40)
        log_info("测试 2: 1460B MSS 回显 (TC-NET-603)")
        try:
            test_data = gen_pattern(1460, pkt_offset)
            pkt_offset += 1460
            client_socket.send(test_data)
            log_success(f"发送 1460B ({len(test_data)} bytes)")

            received = b""
            while len(received) < len(test_data):
                chunk = client_socket.recv(min(1460, len(test_data) - len(received)))
                if not chunk:
                    break
                received += chunk

            if received == test_data:
                log_success("1460B 完全匹配!")
            else:
                log_error(f"1460B 传输不完整: 发送 {len(test_data)}, 接收 {len(received)}")
        except Exception as e:
            log_error(f"1460B 传输错误: {e}")

        # 测试 4096B 缓冲满回显（TC-NET-604）
        log_info("-" * 40)
        log_info("测试 3: 4096B 缓冲满回显 (TC-NET-604)")
        try:
            # 单 Socket 接收缓冲按 1KB block 配置为 2KB（set_recv_buf(sock,30,2)），
            # 4096B 单包会分两次 net_recv 取空，验证分次取数 + 回显 + 字节累计校验正确。
            test_data = gen_pattern(4096, pkt_offset)
            pkt_offset += 4096
            client_socket.send(test_data)
            log_success(f"发送 4096B ({len(test_data)} bytes)")

            received = b""
            while len(received) < len(test_data):
                chunk = client_socket.recv(min(4096, len(test_data) - len(received)))
                if not chunk:
                    break
                received += chunk

            if received == test_data:
                log_success("4096B 完全匹配!")
            else:
                log_error(f"4096B 传输不完整: 发送 {len(test_data)}, 接收 {len(received)}")
        except Exception as e:
            log_error(f"4096B 传输错误: {e}")

        # 测试断开重连（TC-NET-605）
        log_info("-" * 40)
        log_info("测试 4: 断开重连 (TC-NET-605)")
        try:
            client_socket.close()
            client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            client_socket.settimeout(10)
            client_socket.connect((args.ip, args.port))
            log_success("重连成功!")

            # 新连接：固件复位 s_phase6_off=0，PC 期望字节流从偏移 0 重新比对
            test_data = gen_pattern(64, 0)
            client_socket.send(test_data)
            data = client_socket.recv(64)
            if data == test_data:
                log_success("重连回显匹配!")
            else:
                log_error("重连回显不匹配!")
        except Exception as e:
            log_error(f"重连错误: {e}")

        client_socket.close()
        log_success("测试完成，连接已关闭")

    except ConnectionRefusedError:
        log_error("连接被拒绝 - CH395F TCP Server 未运行")
    except socket.timeout:
        log_error("连接超时")
    except Exception as e:
        log_error(f"客户端错误: {e}")
    finally:
        client_socket.close()


# ==================== UDP 测试函数 ====================

def udp_server_test(args):
    """UDP 服务器测试（测试 CH395F UDP Client）"""
    log_info("=" * 60)
    log_info("UDP 服务器测试 - 测试 CH395F UDP Client 功能")
    log_info("=" * 60)

    server_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server_socket.settimeout(30)

    try:
        server_socket.bind(('0.0.0.0', args.port))
        log_success(f"服务器启动，监听端口 {args.port}")
        log_info("等待 CH395F UDP Client 数据...")

        # 接收数据
        for i in range(args.count):
            try:
                data, client_addr = server_socket.recvfrom(args.size)
                log_success(f"收到来自 {client_addr[0]}:{client_addr[1]} ({len(data)} bytes)")
                log_info(f"数据: {data[:50]}...")

                # 回显数据
                server_socket.sendto(data, client_addr)
                log_info(f"回显数据到 {client_addr[0]}:{client_addr[1]}")
            except socket.timeout:
                log_warn("接收超时")
                break
            except Exception as e:
                log_error(f"接收错误: {e}")
                break

        log_success("测试完成")

    except Exception as e:
        log_error(f"服务器错误: {e}")
    finally:
        server_socket.close()


def udp_client_test(args):
    """UDP Echo 测试（测试 CH395F UDP Server，对应指南阶段 3）

    会话结构与固件 phase3_run() 约定：
      - "HELLO" 探针：回显即就绪，不计入固件轮次
      - TC-301：30 轮 64B 基础回显（--size 非默认值时覆盖单轮载荷大小）
      - TC-302：1B/64B/200B 各 --count 轮变长合规（默认 10）
    任一轮 mismatch/timeout 立即终止；结束输出汇总，失败退出码 1。
    """
    log_info("=" * 60)
    log_info("UDP Echo 测试（TC-301 基础回显 + TC-302 变长合规）")
    log_info("=" * 60)

    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    dest = (args.ip, args.port)

    try:
        # Phase A: HELLO 就绪探针（固件只回显不计数，不占用例预算）
        log_info(f"Phase A: HELLO 就绪探针（超时 {args.timeout}s）")
        sock.settimeout(2)
        t0 = time.time()
        ready = False
        while time.time() - t0 < args.timeout:
            try:
                sock.sendto(b"HELLO", dest)
                data, addr = sock.recvfrom(64)
                if data == b"HELLO":
                    log_success(f"就绪 ({time.time()-t0:.1f}s, from {addr[0]}:{addr[1]})")
                    ready = True
                    break
            except socket.timeout:
                pass
        if not ready:
            log_error(f"未就绪（{args.timeout}s 超时）")
            sys.exit(1)

        sock.settimeout(5)

        # Phase B: TC-301 基础回显 30 轮 x 64B
        size_b = 64 if args.size == 1024 else min(args.size, 200)
        rounds_b = 30
        log_info("-" * 40)
        log_info(f"Phase B: TC-301 基础回显（{rounds_b} 轮 x {size_b}B）")
        ok_b = 0
        for i in range(rounds_b):
            td = generate_data(size_b)
            try:
                sock.sendto(td, dest)
                data, _ = sock.recvfrom(size_b + 16)
                if data != td:
                    log_error(f"[TC-301] #{i+1} mismatch: {len(data)}B != {len(td)}B")
                    break
            except socket.timeout:
                log_error(f"[TC-301] #{i+1} echo timeout")
                break
            ok_b += 1
            if ok_b % 10 == 0:
                log_info(f"[TC-301] 进度 {ok_b}/{rounds_b}")
            time.sleep(0.02)

        # Phase C: TC-302 变长合规（仅 Phase B 全过后继续，否则后续必超时）
        sizes_c = [1, 64, 200]
        total_c = len(sizes_c) * args.count
        ok_c = 0
        if ok_b == rounds_b:
            log_info("-" * 40)
            log_info(f"Phase C: TC-302 变长合规（{sizes_c}B 各 {args.count} 轮）")
            for sz in sizes_c:
                ok_sz = 0
                for i in range(args.count):
                    td = generate_data(sz)
                    try:
                        sock.sendto(td, dest)
                        data, _ = sock.recvfrom(sz + 16)
                        if data != td:
                            log_error(f"[TC-302] {sz}B #{i+1} mismatch: {len(data)}B != {sz}B")
                            break
                    except socket.timeout:
                        log_error(f"[TC-302] {sz}B #{i+1} echo timeout")
                        break
                    ok_sz += 1
                    time.sleep(0.02)
                log_info(f"[TC-302] {sz}B: {ok_sz}/{args.count}")
                ok_c += ok_sz
                if ok_sz < args.count:
                    break

        # 汇总与退出码
        log_info("=" * 60)
        if ok_b == rounds_b and ok_c == total_c:
            log_success(f"PASS: TC-301 {ok_b}/{rounds_b}, TC-302 {ok_c}/{total_c}")
        else:
            log_error(f"FAIL: TC-301 {ok_b}/{rounds_b}, TC-302 {ok_c}/{total_c}")
            sys.exit(1)

    except Exception as e:
        log_error(f"错误: {e}")
        sys.exit(1)
    finally:
        sock.close()


# ==================== 回显测试 ====================

def tcp_echo_server_test(args):
    """TCP 回显服务器测试"""
    log_info("=" * 60)
    log_info("TCP 回显服务器测试 - 测试双向通信功能")
    log_info("=" * 60)

    server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server_socket.settimeout(args.timeout)

    try:
        server_socket.bind(('0.0.0.0', args.port))
        server_socket.listen(1)
        log_success(f"回显服务器启动，端口 {args.port}")
        log_info(f"等待连接（超时 {args.timeout}s）...")

        client_socket, client_addr = server_socket.accept()
        log_success(f"连接来自: {client_addr[0]}:{client_addr[1]}")
        client_socket.settimeout(5)

        # 循环回显
        echo_count = 0
        start_time = time.time()

        log_info("开始回显测试 (按 Ctrl+C 停止)...")
        while echo_count < args.count:
            try:
                data = client_socket.recv(args.size)
                if not data:
                    log_warn("客户端断开连接")
                    break

                echo_count += 1
                client_socket.send(data)

                if echo_count % 10 == 0:
                    elapsed = time.time() - start_time
                    rate = echo_count / elapsed if elapsed > 0 else 0
                    log_info(f"回显 {echo_count} 次, 速率: {rate:.1f} msg/s")

            except socket.timeout:
                continue
            except Exception as e:
                log_error(f"回显错误: {e}")
                break

        elapsed = time.time() - start_time
        log_success(f"回显测试完成: {echo_count} 次, 耗时 {elapsed:.2f}s")

        client_socket.close()

    except Exception as e:
        log_error(f"服务器错误: {e}")
    finally:
        server_socket.close()


def udp_echo_server_test(args):
    """UDP 回显服务器测试"""
    log_info("=" * 60)
    log_info("UDP 回显服务器测试 - 测试双向通信功能")
    log_info("=" * 60)

    server_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server_socket.settimeout(60)

    try:
        server_socket.bind(('0.0.0.0', args.port))
        log_success(f"UDP 回显服务器启动，端口 {args.port}")
        log_info("等待数据...")

        echo_count = 0
        start_time = time.time()
        clients = set()

        while echo_count < args.count:
            try:
                data, client_addr = server_socket.recvfrom(args.size)
                if not data:
                    continue

                clients.add(client_addr)
                echo_count += 1

                # 回显数据
                server_socket.sendto(data, client_addr)

                if echo_count % 10 == 0:
                    elapsed = time.time() - start_time
                    rate = echo_count / elapsed if elapsed > 0 else 0
                    log_info(f"回显 {echo_count} 次, 客户端数: {len(clients)}, 速率: {rate:.1f} msg/s")

            except socket.timeout:
                continue
            except Exception as e:
                log_error(f"回显错误: {e}")
                break

        elapsed = time.time() - start_time
        log_success(f"回显测试完成: {echo_count} 次, 耗时 {elapsed:.2f}s")

    except Exception as e:
        log_error(f"服务器错误: {e}")
    finally:
        server_socket.close()


# ==================== 压力测试 ====================

def stress_test_thread(client_id, args, results, barrier):
    """压力测试线程"""
    thread_name = f"Client-{client_id}"
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.settimeout(5)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

    try:
        t0 = time.time()
        sock.connect((args.ip, args.port))
        local = sock.getsockname()
        remote = sock.getpeername()
        t_conn = (time.time() - t0) * 1000
        log_success(f"[{thread_name}] 连接成功 "
                    f"(local={local[0]}:{local[1]}, "
                    f"remote={remote[0]}:{remote[1]}, "
                    f"{t_conn:.0f}ms)")

        barrier.wait()

        success_count = 0
        fail_count = 0

        for i in range(args.count):
            try:
                test_data = f"T{client_id:02d}_{i:04d}".encode()
                t_send = time.time()
                sock.send(test_data)
                data = sock.recv(100)
                t_rtt = (time.time() - t_send) * 1000
                if data:
                    if data == test_data:
                        success_count += 1
                        log_info(f"[{thread_name}] #{i}: OK ({t_rtt:.0f}ms)")
                    else:
                        fail_count += 1
                        log_warn(f"[{thread_name}] #{i}: MISMATCH "
                                 f"want={len(test_data)}B/{test_data[:20].decode(errors='replace')!r} "
                                 f"got={len(data)}B/{data[:20].decode(errors='replace')!r}"
                                 f" ({t_rtt:.0f}ms)")
                else:
                    fail_count += 1
                    log_warn(f"[{thread_name}] #{i}: recv returned empty ({t_rtt:.0f}ms)")

            except socket.timeout:
                fail_count += 1
                t_elapsed = (time.time() - t_send) * 1000
                log_error(f"[{thread_name}] #{i}: TIMEOUT ({t_elapsed:.0f}ms)")
            except ConnectionResetError:
                fail_count += 1
                t_elapsed = (time.time() - t_send) * 1000
                log_error(f"[{thread_name}] #{i}: RESET by peer ({t_elapsed:.0f}ms)")
            except Exception as e:
                fail_count += 1
                t_elapsed = (time.time() - t_send) * 1000
                log_error(f"[{thread_name}] #{i}: {type(e).__name__}: {e} ({t_elapsed:.0f}ms)")

        results[client_id] = (success_count, fail_count)
        t_total = (time.time() - t0) * 1000
        log_info(f"[{thread_name}] 完成: 成功 {success_count}, 失败 {fail_count} ({t_total:.0f}ms)")

    except Exception as e:
        log_error(f"[{thread_name}] 连接失败: {type(e).__name__}: {e}")
        results[client_id] = (0, args.count)
    finally:
        sock.close()


def stress_test(args):
    """压力测试（多连接并发）"""
    log_info("=" * 60)
    log_info("压力测试 - 多连接并发")
    log_info("=" * 60)

    num_clients = min(args.count, 7)  # CH395F 最多支持 7 个并发连接
    log_info(f"创建 {num_clients} 个并发连接...")

    results = {}
    threads = []
    barrier = threading.Barrier(num_clients)

    # 创建并启动线程
    for i in range(num_clients):
        thread = threading.Thread(
            target=stress_test_thread,
            args=(i, args, results, barrier)
        )
        threads.append(thread)
        thread.start()
        time.sleep(0.1)  # 错开连接时间

    # 等待所有线程完成
    for thread in threads:
        thread.join()

    # 汇总结果
    total_success = sum(r[0] for r in results.values())
    total_fail = sum(r[1] for r in results.values())

    log_info("-" * 40)
    log_info("压力测试结果:")
    log_info(f"  并发连接数: {num_clients}")
    log_info(f"  每连接测试次数: {args.count}")
    log_info(f"  总成功: {total_success}")
    log_info(f"  总失败: {total_fail}")
    log_info(f"  成功率: {total_success/(total_success+total_fail)*100:.1f}%")

    if total_fail == 0:
        log_success("压力测试通过!")
    else:
        log_error("压力测试有失败!")


# ==================== DHCP 测试 ====================

DHCP_SERVER_PORT = 67
DHCP_CLIENT_PORT = 68
DHCP_DISCOVER = 1
DHCP_OFFER = 2
DHCP_REQUEST = 3
DHCP_ACK = 5
OPT_MESSAGE_TYPE = 53
DHCP_ANNOUNCE_PORT = 60001    # MCU Socket6 广播宣告端口（固件 phase4_run）
OPT_SERVER_ID = 54
OPT_SUBNET_MASK = 1
OPT_ROUTER = 3
OPT_DNS = 6
OPT_LEASE_TIME = 51
OPT_END = 255


def _ip_to_bytes(ip):
    return bytes(map(int, ip.split('.')))


def _build_dhcp_pkt(op, xid, client_mac, yiaddr, siaddr, options):
    op_b = bytes([op])
    htype = b'\x01'
    hlen = b'\x06'
    hops = b'\x00'
    secs = b'\x00\x00'
    flags = b'\x00\x00'
    ciaddr = b'\x00\x00\x00\x00'
    giaddr = b'\x00\x00\x00\x00'
    chaddr = client_mac + b'\x00' * 10
    sname = b'\x00' * 64
    file = b'\x00' * 128
    magic = b'\x63\x82\x53\x63'
    return op_b + htype + hlen + hops + xid + secs + flags + ciaddr + yiaddr + siaddr + giaddr + chaddr + sname + file + magic + options


def _announce_listener(results):
    """后台监听 MCU 的 4B IP 广播宣告（UDP :60001，TC-602 步1 的 PC 侧观测）"""
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    try:
        s.bind(('0.0.0.0', DHCP_ANNOUNCE_PORT))
    except OSError as e:
        results['announce_err'] = f'绑定 {DHCP_ANNOUNCE_PORT} 失败: {e}'
        return
    s.settimeout(45)
    try:
        data, addr = s.recvfrom(16)
        if len(data) == 4:
            ip = '.'.join(str(b) for b in data)
            log_success(f"收到 IP 宣告: {ip} (from {addr[0]}:{addr[1]})")
            results['announce'] = ip
        else:
            results['announce_err'] = f'宣告长度异常 {len(data)}B（预期 4B）'
    except socket.timeout:
        results['announce_err'] = '45s 未收到宣告'
    finally:
        s.close()


def dhcp_server_test(args):
    """DHCP Server 模式（DISCOVER → OFFER → REQUEST → ACK）"""
    log_info("=" * 60)
    log_info("DHCP Server 测试")
    log_info("=" * 60)

    server_ip = args.dhcp_iface
    offer_ip = args.ip
    subnet_mask = args.dhcp_mask
    router = args.dhcp_gw
    lease_time = 3600

    log_info(f"DHCP 服务器 IP: {server_ip}")
    log_info(f"待分配 IP: {offer_ip}")
    log_info(f"子网掩码: {subnet_mask}")
    log_info(f"网关: {router}")
    print()

    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(('0.0.0.0', DHCP_SERVER_PORT))
    sock.settimeout(60)

    log_info("DHCP 服务器已启动，等待 CH395F DISCOVER...")
    log_info("（请确保 MCU 端已启用 ENABLE_PHASE4_TESTS 并复位）")
    print()

    # 宣告监听线程：MCU 在 ACK 后立即广播 4B IP 到 :60001
    results = {}
    listener = threading.Thread(target=_announce_listener, args=(results,), daemon=True)
    listener.start()

    try:
        # Phase 1: DISCOVER → OFFER
        data, addr = sock.recvfrom(4096)
        if len(data) <= 250:
            log_error("收到无效 DHCP 包")
            return

        msg_type = data[242]
        if msg_type != DHCP_DISCOVER:
            log_error(f"预期 DISCOVER(1)，收到 {msg_type}")
            return

        xid = data[4:8]
        client_mac = data[28:34]
        mac_str = ':'.join(f'{b:02X}' for b in client_mac)
        log_success(f"收到 DHCP DISCOVER from {mac_str} (xid={xid.hex()})")

        # Build OFFER options
        options = b''
        options += struct.pack('BBB', OPT_MESSAGE_TYPE, 1, DHCP_OFFER)
        options += struct.pack('BB', OPT_SUBNET_MASK, 4) + _ip_to_bytes(subnet_mask)
        options += struct.pack('BB', OPT_ROUTER, 4) + _ip_to_bytes(router)
        options += struct.pack('BB', OPT_DNS, 4) + _ip_to_bytes(router)
        options += struct.pack('BB', OPT_SERVER_ID, 4) + _ip_to_bytes(server_ip)
        options += bytes([OPT_LEASE_TIME, 4]) + struct.pack('!I', lease_time)
        options += bytes([OPT_END])

        offer = _build_dhcp_pkt(2, xid, client_mac,
                                _ip_to_bytes(offer_ip),
                                _ip_to_bytes(server_ip),
                                options)

        # Try both global broadcast and subnet broadcast
        for bcast in ['255.255.255.255', '192.168.1.255']:
            sock.sendto(offer, (bcast, DHCP_CLIENT_PORT))
        log_success("发送 DHCP OFFER")

        # Phase 2: REQUEST → ACK（最多重试 8 次）
        sock.settimeout(5)
        phase2_ok = False
        for retry in range(8):
            try:
                data, addr = sock.recvfrom(4096)
                if len(data) <= 250:
                    continue
                msg_type = data[242]
                if msg_type == DHCP_REQUEST:
                    rx_xid = data[4:8]
                    rx_mac = data[28:34]
                    log_success(f"收到 DHCP REQUEST from {':'.join(f'{b:02X}' for b in rx_mac)}")

                    ack_options = b''
                    ack_options += struct.pack('BBB', OPT_MESSAGE_TYPE, 1, DHCP_ACK)
                    ack_options += struct.pack('BB', OPT_SUBNET_MASK, 4) + _ip_to_bytes(subnet_mask)
                    ack_options += struct.pack('BB', OPT_ROUTER, 4) + _ip_to_bytes(router)
                    ack_options += struct.pack('BB', OPT_DNS, 4) + _ip_to_bytes(router)
                    ack_options += struct.pack('BB', OPT_SERVER_ID, 4) + _ip_to_bytes(server_ip)
                    ack_options += bytes([OPT_LEASE_TIME, 4]) + struct.pack('!I', lease_time)
                    ack_options += bytes([OPT_END])

                    ack = _build_dhcp_pkt(2, rx_xid, rx_mac,
                                          _ip_to_bytes(offer_ip),
                                          _ip_to_bytes(server_ip),
                                          ack_options)
                    for bcast in ['255.255.255.255', '192.168.1.255']:
                        sock.sendto(ack, (bcast, DHCP_CLIENT_PORT))
                    log_success(f"发送 DHCP ACK → {offer_ip}")
                    log_success("DHCP 分配完成！")
                    phase2_ok = True
                    break
                elif msg_type == DHCP_DISCOVER:
                    # 重传 DISCOVER，重发 OFFER
                    xid = data[4:8]
                    client_mac = data[28:34]
                    options = b''
                    options += struct.pack('BBB', OPT_MESSAGE_TYPE, 1, DHCP_OFFER)
                    options += struct.pack('BB', OPT_SUBNET_MASK, 4) + _ip_to_bytes(subnet_mask)
                    options += struct.pack('BB', OPT_ROUTER, 4) + _ip_to_bytes(router)
                    options += struct.pack('BB', OPT_DNS, 4) + _ip_to_bytes(router)
                    options += struct.pack('BB', OPT_SERVER_ID, 4) + _ip_to_bytes(server_ip)
                    options += bytes([OPT_LEASE_TIME, 4]) + struct.pack('!I', lease_time)
                    options += bytes([OPT_END])
                    offer = _build_dhcp_pkt(2, xid, client_mac,
                                            _ip_to_bytes(offer_ip),
                                            _ip_to_bytes(server_ip),
                                            options)
                    for bcast in ['255.255.255.255', '192.168.1.255']:
                        sock.sendto(offer, (bcast, DHCP_CLIENT_PORT))
                    log_info(f"重发 DHCP OFFER (retry={retry+1})")
            except socket.timeout:
                log_warn(f"Phase 2 等待超时 (retry={retry+1})")
        if not phase2_ok:
            log_warn("DHCP REQUEST 未收到")

        print()

        # 汇总宣告监听结果（TC-602 步 1 PC 侧观测；join 延迟由固件 30s HELLO 窗覆盖）
        listener.join(timeout=10)
        if results.get('announce'):
            if results['announce'] != offer_ip:
                log_error(f"IP 宣告 {results['announce']} != 租约 {offer_ip}！"
                          "芯片固件可能回读了未提交的旧 IP 配置")
        elif 'announce_err' in results:
            log_warn(f"IP 宣告未观测到：{results['announce_err']}")

        print()

        # Phase 3: 连通性验证（MCU Phase 6 内置 HELLO 端口 = 60000）
        hello_port = args.port if args.port != 8080 else 60000
        log_info(f"验证 {offer_ip} 连通性（UDP HELLO → port {hello_port}）")
        sock.close()

        time.sleep(3)
        verify = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        verify.settimeout(2)
        phase3_ok = False
        for retry in range(5):
            try:
                verify.sendto(b"HELLO", (offer_ip, hello_port))
                data, addr = verify.recvfrom(64)
                if data == b"HELLO":
                    log_success(f"连通性验证通过：{addr[0]}:{addr[1]} 响应正常")
                    phase3_ok = True
                    break
                time.sleep(1)
            except (socket.timeout, ConnectionResetError, OSError):
                log_info(f"等待 UDP 响应... (retry={retry+1})")
                time.sleep(1)
        if not phase3_ok:
            log_warn("连通性验证超时（MCU 的 UDP 服务可能未运行或端口不正确）")
            log_info("提示：CH395F UDP echo 默认端口为 60000，请使用 --port 60000")

    except socket.timeout:
        log_error("等待 DHCP DISCOVER 超时（60s）")
    finally:
        sock.close()



# ==================== 数据包大小测试 ====================

def packet_size_test(args):
    """不同数据包大小测试"""
    log_info("=" * 60)
    log_info("数据包大小测试")
    log_info("=" * 60)

    client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    client_socket.settimeout(5)

    try:
        client_socket.connect((args.ip, args.port))
        log_success("连接成功")

        # 测试不同的数据包大小
        test_sizes = [1, 10, 100, 512, 1024, 2048, 4096]

        for size in test_sizes:
            try:
                test_data = generate_data(size)
                start_time = time.time()
                client_socket.send(test_data)

                received = b""
                while len(received) < size:
                    chunk = client_socket.recv(size - len(received))
                    if not chunk:
                        break
                    received += chunk

                elapsed = time.time() - start_time

                if received == test_data:
                    rate = size / elapsed / 1024 if elapsed > 0 else 0
                    log_success(f"大小 {size:5d} bytes: 成功, 速率 {rate:.1f} KB/s")
                else:
                    log_error(f"大小 {size:5d} bytes: 数据不匹配")
            except Exception as e:
                log_error(f"大小 {size:5d} bytes: 错误 - {e}")

            time.sleep(0.1)

    except Exception as e:
        log_error(f"测试错误: {e}")
    finally:
        client_socket.close()


# ==================== 文件收发测试（配合固件阶段 8 / PHASE8_DIRECTION） ====================

def tcp_file_test(args):
    """文件收发测试：PC 作 TCP Client 连接 MCU:8080（net_test_task 监听）

    协议：
      --op send : 读取 --file（或生成 --size 字节递增模式）发往 MCU，
                  末尾附 4 字节 CRC32（小端，zlib.crc32），供 MCU 整包校验。
      --op recv : 从 MCU 接收 --size 字节写 --file，逐字节比对递增模式并报告 CRC32。
      --op echo : PC 作 Client 连 MCU:8080，把 MCU 发来的数据原样回显，直到对端关闭。

    对应固件 PHASE8_DIRECTION：
      SEND(1)  -> PC 用 --op recv 收 MCU 发来的文件
      RECV(2)  -> PC 用 --op send 发文件（+CRC32 尾）给 MCU
      ECHO(0)  -> PC 用 --op echo 连 MCU:8080 回显（MCU 为 server，先发后收）
    """
    op = args.op
    size = args.size
    ip, port = args.ip, args.port

    if op not in ('send', 'recv', 'echo'):
        log_error("tcp_file: --op 必须为 send / recv / echo")
        return False

    client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    client_socket.settimeout(args.timeout)
    try:
        log_info(f"正在连接 MCU TCP Server: {ip}:{port}")
        client_socket.connect((ip, port))
        log_success(f"已连接 {ip}:{port}")
    except Exception as e:
        log_error(f"连接失败: {e}")
        return False

    try:
        if op == 'send':
            # ---- 准备数据 ----
            if args.file and os.path.exists(args.file):
                with open(args.file, 'rb') as f:
                    data = f.read()
                if len(data) != size:
                    log_warn(f"文件大小 {len(data)} != --size {size}，以实际文件大小 {len(data)} 发送")
                size = len(data)
            else:
                log_info(f"无 --file，生成 {size} 字节递增模式数据")
                data = bytes((i & 0xFF) for i in range(size))
            crc = zlib.crc32(data) & 0xFFFFFFFF
            footer = struct.pack('<I', crc)

            # ---- 发送数据 + CRC32 尾 ----
            t0 = time.time()
            sent = 0
            while sent < size:
                n = client_socket.send(data[sent:size])
                if n <= 0:
                    log_error("发送中断")
                    return False
                sent += n
            sent_f = 0
            while sent_f < 4:
                n = client_socket.send(footer[sent_f:4])
                if n <= 0:
                    log_error("CRC 尾发送失败")
                    return False
                sent_f += n
            ms = int((time.time() - t0) * 1000)
            rate = (size / 1024.0 / (ms / 1000.0)) if ms else 0.0
            log_success(f"已发送 {size} 字节 -> {args.file or '(pattern)'}，"
                        f"CRC32=0x{crc:08X}，耗时 {ms}ms ({rate:.1f} KB/s)")
            return True

        elif op == 'recv':
            rx = bytearray()
            t0 = time.time()
            while len(rx) < size:
                chunk = client_socket.recv(min(65536, size - len(rx)))
                if not chunk:
                    log_error(f"对端提前关闭，仅收到 {len(rx)}/{size}")
                    return False
                rx += chunk
            ms = int((time.time() - t0) * 1000)
            rate = (size / 1024.0 / (ms / 1000.0)) if ms else 0.0

            ok = all(rx[i] == (i & 0xFF) for i in range(size))
            crc = zlib.crc32(bytes(rx)) & 0xFFFFFFFF
            if args.file:
                with open(args.file, 'wb') as f:
                    f.write(rx)
                log_info(f"已写入文件 {args.file}")
            log_success(f"已接收 {size} 字节，耗时 {ms}ms ({rate:.1f} KB/s)，"
                        f"递增模式校验={'通过' if ok else '失败'}，CRC32=0x{crc:08X}")
            return ok

        elif op == 'echo':
            # MCU 为 server：先发 100KB，PC 原样回显，MCU 收齐后关闭连接
            log_info("回显模式：将把 MCU 发来的数据原样返回，直到对端关闭")
            t0 = time.time()
            echoed = 0
            while True:
                data = client_socket.recv(65536)
                if not data:
                    break
                client_socket.sendall(data)
                echoed += len(data)
            ms = int((time.time() - t0) * 1000)
            rate = (echoed / 1024.0 / (ms / 1000.0)) if ms else 0.0
            log_success(f"回显完成，共 {echoed} 字节，耗时 {ms}ms ({rate:.1f} KB/s)")
            return True

    except socket.timeout:
        log_error("传输超时")
        return False
    except Exception as e:
        log_error(f"错误: {e}")
        return False
    finally:
        client_socket.close()


# ==================== 主程序 ====================

def main():
    parser = argparse.ArgumentParser(
        description='CH395F BSD Socket API 测试程序',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
  %(prog)s tcp_client --ip 192.168.1.100 --port 8080
  %(prog)s tcp_server --port 8080
  %(prog)s stress --ip 192.168.1.100 --port 8080 --count 5
  %(prog)s dhcp_server --dhcp-iface 192.168.1.2 --ip 192.168.1.100
  %(prog)s tcp_bulk --port 8081 --bulk 262144 --chunk 4096
        """
    )

    parser.add_argument('mode', choices=['tcp_server', 'tcp_client', 'udp_server', 'udp_client',
                                          'tcp_echo', 'udp_echo', 'stress', 'packet_size',
                                          'dhcp_server', 'tcp_bulk', 'tcp_file'],
                         help='测试模式')
    parser.add_argument('--ip', default='192.168.1.100', help='目标 IP 地址 / DHCP 待分配 IP')
    parser.add_argument('--port', type=int, default=8080, help='端口号')
    parser.add_argument('--count', type=int, default=10, help='测试次数')
    parser.add_argument('--size', type=int, default=1024, help='数据包大小')
    parser.add_argument('--bulk', type=int, default=262144, help='tcp_bulk 模式双向字节数')
    parser.add_argument('--chunk', type=int, default=4096, help='tcp_bulk 模式块大小，须与固件 PHASE5_CHUNK_SIZE 一致')
    parser.add_argument('--timeout', type=int, default=180, help='接受连接超时（秒，默认180）')
    parser.add_argument('--op', choices=['send', 'recv', 'echo'], default='recv',
                        help='tcp_file 方向: send=读文件发往 MCU, recv=从 MCU 收文件落盘, echo=连 MCU 回显')
    parser.add_argument('--file', default=None, help='tcp_file 文件路经（send: 读取; recv: 写入）')
    parser.add_argument('--max-conn', type=int, default=0, help='最大并发连接数（默认0=不限；按并发计数，非累计，TC704 并发测试用 8）')
    parser.add_argument('--dhcp-iface', default='192.168.1.2', help='DHCP 服务器网卡 IP')
    parser.add_argument('--dhcp-mask', default='255.255.255.0', help='DHCP 分配的子网掩码')
    parser.add_argument('--dhcp-gw', default='192.168.1.1', help='DHCP 分配的网关')

    args = parser.parse_args()

    # GBK 控制台容错：不可编码字符替换而非崩溃
    for stream in (sys.stdout, sys.stderr):
        if stream is not None and hasattr(stream, 'reconfigure'):
            stream.reconfigure(errors='replace')

    if sys.platform == 'win32':
        os.system('')  # 启用 Windows 控制台 VT 序列，避免 ANSI 颜色码乱码

    print()
    print(Colors.BOLD + "=" * 60 + Colors.END)
    print(Colors.BOLD + "  CH395F BSD Socket API 测试程序" + Colors.END)
    print(Colors.BOLD + "=" * 60 + Colors.END)
    print()

    ret = None
    try:
        if args.mode == 'tcp_server':
            ret = tcp_server_test(args)
        elif args.mode == 'tcp_client':
            tcp_client_test(args)
        elif args.mode == 'udp_server':
            udp_server_test(args)
        elif args.mode == 'udp_client':
            udp_client_test(args)
        elif args.mode == 'tcp_echo':
            tcp_echo_server_test(args)
        elif args.mode == 'udp_echo':
            udp_echo_server_test(args)
        elif args.mode == 'stress':
            stress_test(args)
        elif args.mode == 'packet_size':
            packet_size_test(args)
        elif args.mode == 'tcp_bulk':
            ret = tcp_bulk_test(args)
        elif args.mode == 'tcp_file':
            ret = tcp_file_test(args)
        elif args.mode == 'dhcp_server':
            dhcp_server_test(args)
    except KeyboardInterrupt:
        log_warn("用户中断")
        sys.exit(130)
    except Exception as e:
        log_error(f"未预期的错误: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(2)

    if ret is False:
        sys.exit(1)


if __name__ == '__main__':
    main()