#!/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       - 压力测试（多连接并发测试）

选项:
    --ip         目标 IP 地址（默认: 192.168.1.100）
    --port       端口号（默认: 8080）
    --count      测试次数（默认: 10）
    --size       数据包大小（默认: 1024）

示例:
    # 测试 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
"""

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


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(f"✓ {msg}", Colors.GREEN)


def log_error(msg):
    log_msg(f"✗ {msg}", Colors.RED)


def log_info(msg):
    log_msg(f"ℹ {msg}", Colors.BLUE)


def log_warn(msg):
    log_msg(f"⚠ {msg}", Colors.YELLOW)


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


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


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

def tcp_server_test(args):
    """TCP 服务器测试（测试 CH395F TCP Client）"""
    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)
    server_socket.settimeout(args.timeout)

    try:
        server_socket.bind(('0.0.0.0', args.port))
        server_socket.listen(1)
        log_success(f"服务器启动，监听端口 {args.port}")

        for conn_idx in range(3):
            try:
                server_socket.settimeout(args.timeout)
                log_info(f"等待第 {conn_idx+1} 个连接（超时 {args.timeout}s）...")
                client_socket, client_addr = server_socket.accept()
                log_success(f"连接#{conn_idx+1} 来自: {client_addr[0]}:{client_addr[1]}")
                client_socket.settimeout(5)

                # Test 1: 接收数据
                log_info("-" * 40)
                log_info(f"[连接{conn_idx+1}] 测试 1: 接收数据")
                try:
                    data = client_socket.recv(args.size)
                    if data:
                        log_success(f"收到数据 ({len(data)} bytes): {data[:50]}...")
                except socket.timeout:
                    log_error("接收超时")

                # Test 2: PING echo
                log_info(f"[连接{conn_idx+1}] 测试 2: 双向通信 (PING)")
                ping_ok = 0
                for i in range(5):
                    try:
                        test_data = f"PING_{i}".encode()
                        client_socket.send(test_data)
                        data = client_socket.recv(args.size)
                        if data and data == test_data:
                            log_success(f"回显匹配: {test_data.decode()}")
                            ping_ok += 1
                        elif data:
                            log_error(f"回显不匹配! 收到: {data[:20]}")
                            break
                    except socket.timeout:
                        log_error("通信超时")
                        break
                    except Exception as e:
                        log_error(f"通信错误: {e}")
                        break
                log_info(f"PING 匹配: {ping_ok}/5")

                # Test 3: 大包发送
                log_info(f"[连接{conn_idx+1}] 测试 3: 大包发送")
                for i in range(args.count):
                    try:
                        test_data = generate_data(min(100, args.size))
                        client_socket.send(test_data)
                        time.sleep(0.1)
                    except Exception as e:
                        log_error(f"发送错误: {e}")
                        break
                log_success(f"发送 {args.count} 包完成")

                client_socket.close()
                log_success(f"连接#{conn_idx+1} 测试完成")
                time.sleep(0.5)

            except socket.timeout:
                log_error(f"等待连接#{conn_idx+1} 超时")
                break
            except Exception as e:
                log_error(f"连接#{conn_idx+1} 错误: {e}")
                break

        log_success("全部连接测试完成")

    except socket.timeout:
        log_error("等待连接超时")
    except Exception as e:
        log_error(f"服务器错误: {e}")
    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("连接成功!")

        # 测试发送数据
        log_info("-" * 40)
        log_info("测试 1: 发送数据")
        for i in range(args.count):
            try:
                test_data = generate_data(min(100, args.size))
                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(args.size)
                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

        # 测试大数据量
        log_info("-" * 40)
        log_info("测试 2: 大数据量传输")
        try:
            large_data = generate_data(min(1024, args.size))
            client_socket.send(large_data)
            log_success(f"发送大数据 ({len(large_data)} bytes)")

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

            if received == large_data:
                log_success("大数据传输完整!")
            else:
                log_error(f"大数据传输不完整: 发送 {len(large_data)}, 接收 {len(received)}")
        except Exception as e:
            log_error(f"大数据传输错误: {e}")

        # 测试快速收发
        log_info("-" * 40)
        log_info("测试 3: 快速收发")
        success_count = 0
        fail_count = 0
        start_time = time.time()

        for i in range(args.count):
            try:
                test_data = f"FAST_{i:04d}".encode()
                client_socket.send(test_data)
                data = client_socket.recv(100)
                if data and data == test_data:
                    success_count += 1
                else:
                    fail_count += 1
            except Exception as e:
                fail_count += 1

        elapsed = time.time() - start_time
        log_info(f"快速收发结果: 成功 {success_count}, 失败 {fail_count}, 耗时 {elapsed:.2f}s")

        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）"""
    log_info("=" * 60)
    log_info("UDP Echo 测试")
    log_info("=" * 60)

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

    try:
        # Phase A: HELLO 探针等待 STM32 就绪
        log_info(f"Phase A: 等待 CH395F 就绪（超时 {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 超时）")
            return

        sock.settimeout(5)

        # Phase B: PING Echo
        log_info("-" * 40)
        log_info("Phase B: PING Echo（5 轮）")
        ok = 0
        for i in range(5):
            msg = f"PING_{i}".encode()
            try:
                sock.sendto(msg, dest)
                data, _ = sock.recvfrom(64)
                if data == msg:
                    ok += 1
                    log_success(f"PING_{i} OK")
                else:
                    log_error(f"PING_{i} mismatch: {data[:20]}")
                    break
            except socket.timeout:
                log_error(f"PING_{i} timeout")
                break
            time.sleep(0.02)
        log_info(f"PING: {ok}/5")

        # Phase C: 大包 Echo
        log_info("-" * 40)
        log_info(f"Phase C: 大包 Echo（{args.count} 轮）")
        for i in range(args.count):
            td = generate_data(min(100, args.size))
            try:
                sock.sendto(td, dest)
                data, _ = sock.recvfrom(args.size)
                if data == td:
                    log_success(f"#{i+1} OK ({len(data)}B)")
                else:
                    log_error(f"#{i+1} mismatch")
                    break
                time.sleep(0.02)
            except socket.timeout:
                log_warn(f"#{i+1} timeout")
            except Exception as e:
                log_error(f"#{i+1}: {e}")
                break

        log_success("Phase 3 完成")

    except Exception as e:
        log_error(f"错误: {e}")
    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):
    """压力测试线程"""
    thread_name = f"Client-{client_id}"
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.settimeout(10)

    try:
        sock.connect((args.ip, args.port))
        log_success(f"[{thread_name}] 连接成功")

        success_count = 0
        fail_count = 0

        for i in range(args.count):
            try:
                # 发送
                test_data = f"T{client_id:02d}_{i:04d}".encode()
                sock.send(test_data)

                # 接收
                data = sock.recv(100)
                if data and data == test_data:
                    success_count += 1
                else:
                    fail_count += 1

            except Exception as e:
                fail_count += 1

        results[client_id] = (success_count, fail_count)
        log_info(f"[{thread_name}] 完成: 成功 {success_count}, 失败 {fail_count}")

    except Exception as e:
        log_error(f"[{thread_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 = []

    # 创建并启动线程
    for i in range(num_clients):
        thread = threading.Thread(
            target=stress_test_thread,
            args=(i, args, results)
        )
        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("压力测试有失败!")


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

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


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

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

    parser.add_argument('mode', choices=['tcp_server', 'tcp_client', 'udp_server', 'udp_client',
                                          'tcp_echo', 'udp_echo', 'stress', 'packet_size'],
                        help='测试模式')
    parser.add_argument('--ip', default='192.168.1.100', help='目标 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('--timeout', type=int, default=180, help='接受连接超时（秒，默认180）')

    args = parser.parse_args()

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

    try:
        if args.mode == 'tcp_server':
            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)
    except KeyboardInterrupt:
        log_warn("用户中断")
    except Exception as e:
        log_error(f"未预期的错误: {e}")
        import traceback
        traceback.print_exc()


if __name__ == '__main__':
    main()