#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <openssl/ssl.h>
#include <openssl/err.h>
#include <getopt.h>
#include <time.h>
#include <stdarg.h>
#include <errno.h>
#include <sys/time.h>

#define BUFFER_SIZE 4096

// 初始化OpenSSL
void init_openssl() {
    SSL_load_error_strings();
    OpenSSL_add_ssl_algorithms();
    SSL_library_init();
}

// 清理OpenSSL
void cleanup_openssl() {
    EVP_cleanup();
    ERR_free_strings();
}

// 创建SSL上下文
SSL_CTX *create_context(int tls_version) {
    const SSL_METHOD *method;
    SSL_CTX *ctx;

    method = TLS_client_method();
    ctx = SSL_CTX_new(method);
    if (!ctx) {
        perror("无法创建SSL上下文");
        ERR_print_errors_fp(stderr);
        exit(EXIT_FAILURE);
    }

    // 设置TLS版本
    switch (tls_version) {
        case 1:  // TLS 1.2
            SSL_CTX_set_min_proto_version(ctx, TLS1_2_VERSION);
            SSL_CTX_set_max_proto_version(ctx, TLS1_2_VERSION);
            break;
        case 2:  // TLS 1.3
            SSL_CTX_set_min_proto_version(ctx, TLS1_3_VERSION);
            SSL_CTX_set_max_proto_version(ctx, TLS1_3_VERSION);
            break;
        case 0:  // 自动协商（默认）
        default:
            // 不设置版本限制，让OpenSSL自动协商
            break;
    }

    // 跳过证书验证（仅用于测试）
    SSL_CTX_set_verify(ctx, SSL_VERIFY_NONE, NULL);

    return ctx;
}

// 时间戳日志函数
void log_with_timestamp(const char *format, ...) {
    time_t now;
    struct tm *tm_info;
    char timestamp[64];
    va_list args;
    
    time(&now);
    tm_info = localtime(&now);
    strftime(timestamp, sizeof(timestamp), "%Y-%m-%d %H:%M:%S", tm_info);
    
    printf("[%s] ", timestamp);
    va_start(args, format);
    vprintf(format, args);
    va_end(args);
}

// 完整测试函数（包含重试逻辑）
int perform_test(struct sockaddr_in *server_addr, SSL_CTX *ssl_ctx, const char *message, 
                int byte_count, int max_retries, int *retry_count_out) {
    int socket_fd;
    SSL *ssl = NULL;
    char buffer[BUFFER_SIZE];
    int bytes_received;
    int current_retry = 0;
    int test_success = 0;
    
    *retry_count_out = 0;
    
    log_with_timestamp("进入perform_test函数\n");
    
    // 重试循环
    for (current_retry = 0; current_retry <= max_retries; current_retry++) {
        log_with_timestamp("重试循环: 第 %d/%d 次尝试\n", current_retry, max_retries);
        
        if (current_retry > 0) {
            // 指数退避：2^retry_count 秒，最大60秒
            int backoff_time = (1 << current_retry); // 2的幂：2, 4, 8, 16...
            if (backoff_time > 60) {
                backoff_time = 60; // 最大60秒
            }
            log_with_timestamp("开始第 %d 次重试，等待%d秒（指数退避）...\n", current_retry, backoff_time);
            sleep(backoff_time);
            *retry_count_out = current_retry;
        }
        
        // 创建套接字
        log_with_timestamp("开始创建套接字...\n");
        socket_fd = socket(AF_INET, SOCK_STREAM, 0);
        if (socket_fd < 0) {
            log_with_timestamp("创建套接字失败: %s\n", strerror(errno));
            continue;
        }
        log_with_timestamp("套接字创建成功，fd=%d\n", socket_fd);
            
            // 设置连接超时（30秒）
            struct timeval connect_timeout;
            connect_timeout.tv_sec = 30;
            connect_timeout.tv_usec = 0;
            
            if (setsockopt(socket_fd, SOL_SOCKET, SO_SNDTIMEO, &connect_timeout, sizeof(connect_timeout)) < 0) {
                log_with_timestamp("设置连接超时失败: %s\n", strerror(errno));
            }
            
            // 连接到服务器
            log_with_timestamp("开始连接到服务器...\n");
            if (connect(socket_fd, (struct sockaddr *)server_addr, sizeof(*server_addr)) < 0) {
                log_with_timestamp("连接失败: %s\n", strerror(errno));
                close(socket_fd);
                continue;
            }
            
            log_with_timestamp("TCP连接成功\n");
            
            // 设置套接字超时（30秒）
            struct timeval timeout;
            timeout.tv_sec = 30;
            timeout.tv_usec = 0;
            
            if (setsockopt(socket_fd, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout)) < 0) {
                log_with_timestamp("设置接收超时失败\n");
            }
            
            if (setsockopt(socket_fd, SOL_SOCKET, SO_SNDTIMEO, &timeout, sizeof(timeout)) < 0) {
                log_with_timestamp("设置发送超时失败\n");
            }
            
            // 创建SSL连接
            ssl = SSL_new(ssl_ctx);
            if (!ssl) {
                log_with_timestamp("SSL对象创建失败\n");
                close(socket_fd);
                continue;
            }
            
            SSL_set_fd(ssl, socket_fd);
            
            if (SSL_connect(ssl) <= 0) {
                log_with_timestamp("SSL握手失败\n");
                ERR_print_errors_fp(stderr);
                SSL_free(ssl);
                close(socket_fd);
                continue;
            }
            
            log_with_timestamp("SSL握手成功，TLS版本: %s\n", SSL_get_version(ssl));
            
            // 发送消息
            int write_result = SSL_write(ssl, message, strlen(message));
            if (write_result <= 0) {
                int ssl_error = SSL_get_error(ssl, write_result);
                log_with_timestamp("发送消息失败，SSL错误: %d\n", ssl_error);
                if (ssl_error == SSL_ERROR_SYSCALL) {
                    log_with_timestamp("系统调用错误: %s\n", strerror(errno));
                }
                SSL_free(ssl);
                close(socket_fd);
                continue;
            }
            
            if (byte_count > 0) {
                log_with_timestamp("字节数据已发送 (长度: %zu)\n", strlen(message));
            } else {
                log_with_timestamp("消息已发送 (长度: %zu)\n", strlen(message));
            }
            
            // 等待一小段时间，确保服务器处理完成
            usleep(100000); // 等待100ms
            
            // 接收MD5结果
            bytes_received = SSL_read(ssl, buffer, BUFFER_SIZE - 1);
            if (bytes_received <= 0) {
                int ssl_error = SSL_get_error(ssl, bytes_received);
                log_with_timestamp("接收MD5结果失败，SSL错误: %d\n", ssl_error);
                
                // 详细的错误诊断
                switch (ssl_error) {
                    case SSL_ERROR_ZERO_RETURN:
                        log_with_timestamp("服务器关闭连接\n");
                        break;
                    case SSL_ERROR_WANT_READ:
                        log_with_timestamp("SSL需要读取数据（可能超时）\n");
                        break;
                    case SSL_ERROR_WANT_WRITE:
                        log_with_timestamp("SSL需要写入数据\n");
                        break;
                    case SSL_ERROR_SYSCALL:
                        if (errno == EAGAIN || errno == EWOULDBLOCK) {
                            log_with_timestamp("接收超时（30秒无响应）\n");
                        } else {
                            log_with_timestamp("系统调用错误: %s\n", strerror(errno));
                        }
                        break;
                    case SSL_ERROR_SSL:
                        log_with_timestamp("SSL协议错误\n");
                        ERR_print_errors_fp(stderr);
                        break;
                    default:
                        log_with_timestamp("未知SSL错误\n");
                        break;
                }
                
                SSL_free(ssl);
                close(socket_fd);
                continue;
            }
            
            buffer[bytes_received] = '\0';
            log_with_timestamp("收到MD5结果: %s\n", buffer);
            
            // 清理连接
            SSL_shutdown(ssl);
            SSL_free(ssl);
            close(socket_fd);
            
            log_with_timestamp("连接已关闭\n");
            
        test_success = 1;
        break;
    }
    
    return test_success;
}

int main(int argc, char *argv[]) {
    char *host = "localhost";
    int port = 8443;
    int socket_fd;
    struct sockaddr_in server_addr;
    SSL_CTX *ssl_ctx;
    char message[BUFFER_SIZE];
    int test_count = 1;        // 默认测试1次
    int test_interval = 1;      // 默认间隔1秒
    int tls_version = 0;        // 默认自动协商TLS版本
    int byte_count = 0;         // 字节数参数，0表示使用消息内容
    int tls_compare_mode = 0;   // TLS版本对比测试模式，0=关闭，1=开启
    int c;
    
    // 统计变量
    int success_count = 0;      // 成功次数
    int failure_count = 0;      // 失败次数
    int total_retry_count = 0;  // 总重试次数
    time_t start_time, end_time; // 开始和结束时间

    // 解析命令行参数
    while ((c = getopt(argc, argv, "m:h:p:c:i:t:s:C")) != -1) {
        switch (c) {
            case 'm':
                strcpy(message, optarg);
                break;
            case 'h':
                host = optarg;
                break;
            case 'p':
                port = atoi(optarg);
                if (port <= 0 || port > 65535) {
                    printf("错误: 端口号必须在1-65535之间\n");
                    exit(EXIT_FAILURE);
                }
                break;
            case 'c':
                test_count = atoi(optarg);
                if (test_count <= 0) {
                    printf("错误: 测试次数必须大于0\n");
                    exit(EXIT_FAILURE);
                }
                break;
            case 'i':
                test_interval = atoi(optarg);
                if (test_interval < 0) {
                    printf("错误: 测试间隔不能为负数\n");
                    exit(EXIT_FAILURE);
                }
                break;
            case 't':
                if (strcmp(optarg, "1.2") == 0) {
                    tls_version = 1;  // TLS 1.2
                } else if (strcmp(optarg, "1.3") == 0) {
                    tls_version = 2;  // TLS 1.3
                } else if (strcmp(optarg, "auto") == 0) {
                    tls_version = 0;  // 自动协商
                } else {
                    printf("错误: 无效的TLS版本 '%s'，支持: 1.2, 1.3, auto\n", optarg);
                    exit(EXIT_FAILURE);
                }
                break;
            case 's':
                byte_count = atoi(optarg);
                if (byte_count <= 0) {
                    printf("错误: 字节数必须大于0\n");
                    exit(EXIT_FAILURE);
                }
                break;
            case 'C':
                tls_compare_mode = 1;
                break;
            case '?':
                printf("使用 -h 查看帮助信息\n");
                exit(EXIT_FAILURE);
            default:
                abort();
        }
    }

    // 如果没有通过-m参数指定消息，检查位置参数
    if (strlen(message) == 0 && optind < argc) {
        strcpy(message, argv[optind]);
    }

    // 检查必需参数：必须有消息内容或字节数
    if (strlen(message) == 0 && byte_count == 0) {
        printf("用法: %s -m <消息> 或 -s <字节数> [-h 主机] [-p 端口] [-c 次数] [-i 间隔秒数] [-t TLS版本] [-C]\n", argv[0]);
        printf("选项:\n");
        printf("  -m 消息     要发送的消息内容\n");
        printf("  -s 字节数   要发送的字节数（发送01数据）\n");
        printf("  -h 主机     服务器主机地址 (默认: localhost)\n");
        printf("  -p 端口     服务器端口 (默认: 8443)\n");
        printf("  -c 次数     循环测试次数 (默认: 1)\n");
        printf("  -i 间隔     测试间隔秒数 (默认: 1)\n");
        printf("  -t 版本     TLS版本: 1.2, 1.3, auto (默认: auto)\n");
        printf("  -C          TLS版本对比模式：先测试TLS1.2，等待2秒后测试TLS1.3\n");
        printf("\n示例:\n");
        printf("  %s -m \"Hello\" -c 5 -i 2    # 发送Hello消息5次，间隔2秒\n", argv[0]);
        printf("  %s -s 1024                  # 发送1024字节的01数据\n", argv[0]);
        printf("  %s -s 2048 -c 3 -i 1        # 发送2048字节01数据，循环3次\n", argv[0]);
        printf("  %s -m \"Test\" -h 192.168.1.100 -p 9000  # 连接到指定服务器\n", argv[0]);
        printf("  %s -m \"TLS1.2\" -t 1.2      # 强制使用TLS 1.2\n", argv[0]);
        printf("  %s -m \"TLS1.3\" -t 1.3      # 强制使用TLS 1.3\n", argv[0]);
        printf("  %s -m \"Compare\" -C -c 5    # TLS版本对比测试，循环5轮\n", argv[0]);
        exit(EXIT_FAILURE);
    }

    log_with_timestamp("连接到TLS服务器: %s:%d\n", host, port);
    
    // 根据参数类型显示不同的信息
    if (byte_count > 0) {
        // 检查字节数是否超出缓冲区大小
        if (byte_count > BUFFER_SIZE - 1) {
            printf("错误: 字节数 %d 超出最大缓冲区大小 %d\n", byte_count, BUFFER_SIZE - 1);
            exit(EXIT_FAILURE);
        }
        
        log_with_timestamp("发送字节数: %d (01数据)\n", byte_count);
        // 生成01数据
        for (int i = 0; i < byte_count; i++) {
            message[i] = '0' + (i % 2);  // 交替发送0和1
        }
        message[byte_count] = '\0';
    } else {
        log_with_timestamp("发送消息: %s\n", message);
    }
    
    log_with_timestamp("测试次数: %d, 间隔: %d秒\n", test_count, test_interval);
    
    // 显示TLS版本信息
    const char *tls_version_str;
    switch (tls_version) {
        case 1:
            tls_version_str = "TLS 1.2";
            break;
        case 2:
            tls_version_str = "TLS 1.3";
            break;
        case 0:
        default:
            tls_version_str = "自动协商";
            break;
    }
    // 显示测试模式
    if (tls_compare_mode) {
        log_with_timestamp("TLS版本对比模式: 每轮先测试TLS 1.2，等待2秒后测试TLS 1.3\n");
    } else {
        log_with_timestamp("TLS版本: %s\n", tls_version_str);
    }

    // 记录开始时间
    time(&start_time);
    log_with_timestamp("测试开始时间: %s", ctime(&start_time));

    // 初始化OpenSSL
    init_openssl();

    // 执行循环测试
    for (int i = 1; i <= test_count; i++) {
        if (tls_compare_mode) {
            // TLS版本对比模式：先测试TLS 1.2，再测试TLS 1.3
            log_with_timestamp("\n=== 测试轮次 %d/%d ===\n", i, test_count);
            
            // 测试TLS 1.2
            log_with_timestamp("\n--- TLS 1.2 测试 ---\n");
            ssl_ctx = create_context(1);  // 1 = TLS 1.2
            
            memset(&server_addr, 0, sizeof(server_addr));
            server_addr.sin_family = AF_INET;
            server_addr.sin_port = htons(port);
            
            if (inet_pton(AF_INET, host, &server_addr.sin_addr) <= 0) {
                log_with_timestamp("地址解析失败: %s\n", strerror(errno));
                failure_count++;
            } else {
                int retry_count = 0;
                if (perform_test(&server_addr, ssl_ctx, message, byte_count, 3, &retry_count)) {
                    log_with_timestamp("TLS 1.2 测试成功\n");
                    success_count++;
                    total_retry_count += retry_count;
                } else {
                    log_with_timestamp("TLS 1.2 测试失败（已重试 %d 次）\n", retry_count);
                    failure_count++;
                    total_retry_count += retry_count;
                }
            }
            SSL_CTX_free(ssl_ctx);
            
            // 等待2秒
            log_with_timestamp("等待2秒后测试TLS 1.3...\n");
            sleep(1);
            
            // 测试TLS 1.3
            log_with_timestamp("\n--- TLS 1.3 测试 ---\n");
            ssl_ctx = create_context(2);  // 2 = TLS 1.3
            
            memset(&server_addr, 0, sizeof(server_addr));
            server_addr.sin_family = AF_INET;
            server_addr.sin_port = htons(port);
            
            if (inet_pton(AF_INET, host, &server_addr.sin_addr) <= 0) {
                log_with_timestamp("地址解析失败: %s\n", strerror(errno));
                failure_count++;
            } else {
                int retry_count = 0;
                if (perform_test(&server_addr, ssl_ctx, message, byte_count, 3, &retry_count)) {
                    log_with_timestamp("TLS 1.3 测试成功\n");
                    success_count++;
                    total_retry_count += retry_count;
                } else {
                    log_with_timestamp("TLS 1.3 测试失败（已重试 %d 次）\n", retry_count);
                    failure_count++;
                    total_retry_count += retry_count;
                }
            }
            SSL_CTX_free(ssl_ctx);
            
            log_with_timestamp("--- 本轮测试完成 ---\n");
        } else {
            // 普通测试模式
            log_with_timestamp("\n=== 测试 %d/%d ===\n", i, test_count);
            
            // 创建SSL上下文（只创建一次）
            if (i == 1) {
                ssl_ctx = create_context(tls_version);
            }
            
            // 设置服务器地址
            log_with_timestamp("开始设置服务器地址: %s:%d\n", host, port);
            memset(&server_addr, 0, sizeof(server_addr));
            server_addr.sin_family = AF_INET;
            server_addr.sin_port = htons(port);
            
            log_with_timestamp("开始解析IP地址...\n");
            if (inet_pton(AF_INET, host, &server_addr.sin_addr) <= 0) {
                log_with_timestamp("地址解析失败: %s\n", strerror(errno));
                failure_count++;
                continue;
            }
            log_with_timestamp("地址解析成功\n");
            
            // 执行测试（包含重试逻辑）
            log_with_timestamp("开始执行测试...\n");
            int retry_count = 0;
            if (perform_test(&server_addr, ssl_ctx, message, byte_count, 3, &retry_count)) {
                log_with_timestamp("测试成功\n");
                success_count++;
                total_retry_count += retry_count;
            } else {
                log_with_timestamp("测试失败（已重试 %d 次）\n", retry_count);
                failure_count++;
                total_retry_count += retry_count;
            }
        }
        
        // 如果不是最后一次测试，等待指定间隔
        if (i < test_count) {
            log_with_timestamp("等待 %d 秒后进行下一次测试...\n", test_interval);
            sleep(test_interval);
        }
    }
    
    // 清理SSL上下文（非对比模式）
    if (!tls_compare_mode && ssl_ctx) {
        SSL_CTX_free(ssl_ctx);
    }

    // 记录结束时间
    time(&end_time);
    
    // 最终清理
    cleanup_openssl();

    // 计算统计信息
    int total_tests = success_count + failure_count;
    double success_rate = (total_tests > 0) ? ((double)success_count / total_tests) * 100.0 : 0.0;
    double elapsed_time = difftime(end_time, start_time);
    
    // 输出统计信息
    log_with_timestamp("\n");
    log_with_timestamp("========================================\n");
    log_with_timestamp("           测试统计报告\n");
    log_with_timestamp("========================================\n");
    log_with_timestamp("测试开始时间: %s", ctime(&start_time));
    log_with_timestamp("测试结束时间: %s", ctime(&end_time));
    log_with_timestamp("总耗时: %.2f 秒\n", elapsed_time);
    log_with_timestamp("----------------------------------------\n");
    log_with_timestamp("总测试次数: %d\n", total_tests);
    log_with_timestamp("成功次数: %d\n", success_count);
    log_with_timestamp("失败次数: %d\n", failure_count);
    log_with_timestamp("总重试次数: %d\n", total_retry_count);
    log_with_timestamp("成功率: %.2f%%\n", success_rate);
    log_with_timestamp("----------------------------------------\n");
    
    if (success_count > 0) {
        log_with_timestamp("平均每次测试耗时: %.2f 秒\n", elapsed_time / success_count);
    }
    
    if (failure_count > 0) {
        log_with_timestamp("失败率: %.2f%%\n", 100.0 - success_rate);
    }
    
    log_with_timestamp("========================================\n");
    
    // 根据成功率返回适当的退出码
    if (success_rate >= 100.0) {
        log_with_timestamp("✅ 所有测试通过！\n");
        return 0;
    } else if (success_rate >= 80.0) {
        log_with_timestamp("⚠️  大部分测试通过，有少量失败\n");
        return 1;
    } else {
        log_with_timestamp("❌ 测试失败率较高，请检查服务器状态\n");
        return 2;
    }
}