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

#define MAX_CLIENTS 100
#define BUFFER_SIZE 4096
#define MAX_MESSAGE_SIZE 4096
#define CLIENT_TIMEOUT 30  // 客户端超时时间（秒）

// 全局变量
SSL_CTX *ssl_ctx;
int server_socket;
int client_count = 0;
pthread_mutex_t client_mutex = PTHREAD_MUTEX_INITIALIZER;

// IP统计结构体
typedef struct {
    char ip[INET_ADDRSTRLEN];
    int total_connections;
    int successful_connections;
    time_t first_connection_time;
    time_t last_connection_time;
} ip_stats_t;

ip_stats_t ip_statistics[MAX_CLIENTS];
pthread_mutex_t ip_stats_mutex = PTHREAD_MUTEX_INITIALIZER;

// 时间戳日志函数
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);
}

// 更新IP统计信息
void update_ip_stats(const char *ip, int success) {
    pthread_mutex_lock(&ip_stats_mutex);
    
    // 查找或创建IP统计记录
    int index = -1;
    for (int i = 0; i < MAX_CLIENTS; i++) {
        if (strlen(ip_statistics[i].ip) == 0) {
            index = i;
            strcpy(ip_statistics[i].ip, ip);
            ip_statistics[i].total_connections = 0;
            ip_statistics[i].successful_connections = 0;
            ip_statistics[i].first_connection_time = time(NULL);
            break;
        } else if (strcmp(ip_statistics[i].ip, ip) == 0) {
            index = i;
            break;
        }
    }
    
    if (index >= 0) {
        ip_statistics[index].total_connections++;
        if (success) {
            ip_statistics[index].successful_connections++;
        }
        ip_statistics[index].last_connection_time = time(NULL);
        
        // 每10次连接统计一次
        if (ip_statistics[index].total_connections % 10 == 0) {
            double success_rate = (double)ip_statistics[index].successful_connections / 
                                 ip_statistics[index].total_connections * 100.0;
            time_t duration = ip_statistics[index].last_connection_time - 
                             ip_statistics[index].first_connection_time;
            
            log_with_timestamp("IP统计 [%s]: 总次数=%d, 成功次数=%d, 成功率=%.2f%%, 耗时=%ld秒\n",
                             ip, ip_statistics[index].total_connections,
                             ip_statistics[index].successful_connections,
                             success_rate, duration);
        }
    }
    
    pthread_mutex_unlock(&ip_stats_mutex);
}

// 客户端结构体
typedef struct {
    int socket;
    SSL *ssl;
    pthread_t thread;
    int active;
} client_t;

client_t clients[MAX_CLIENTS];

// 初始化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() {
    const SSL_METHOD *method;
    SSL_CTX *ctx;

    // 使用TLS方法，支持1.2和1.3
    method = TLS_server_method();
    ctx = SSL_CTX_new(method);
    if (!ctx) {
        perror("无法创建SSL上下文");
        ERR_print_errors_fp(stderr);
        exit(EXIT_FAILURE);
    }

    // 设置最小TLS版本为1.2
    SSL_CTX_set_min_proto_version(ctx, TLS1_2_VERSION);
    
    // 加载证书和私钥（从cert目录加载）
    if (SSL_CTX_use_certificate_file(ctx, "cert/server.crt", SSL_FILETYPE_PEM) <= 0) {
        ERR_print_errors_fp(stderr);
        exit(EXIT_FAILURE);
    }

    if (SSL_CTX_use_PrivateKey_file(ctx, "cert/server.key", SSL_FILETYPE_PEM) <= 0) {
        ERR_print_errors_fp(stderr);
        exit(EXIT_FAILURE);
    }

    return ctx;
}

// 计算MD5哈希
void calculate_md5(const char *input, char *output) {
    unsigned char digest[MD5_DIGEST_LENGTH];
    MD5_CTX ctx;
    
    MD5_Init(&ctx);
    MD5_Update(&ctx, input, strlen(input));
    MD5_Final(digest, &ctx);
    
    for (int i = 0; i < MD5_DIGEST_LENGTH; i++) {
        sprintf(output + (i * 2), "%02x", digest[i]);
    }
    output[32] = '\0';
}

// 处理客户端连接
void *handle_client(void *arg) {
    client_t *client = (client_t *)arg;
    char buffer[BUFFER_SIZE];
    char md5_result[33];
    int bytes_received;
    char client_ip[INET_ADDRSTRLEN];
    int connection_success = 0;
    
    // 获取客户端IP地址
    struct sockaddr_in client_addr;
    socklen_t addr_len = sizeof(client_addr);
    if (getpeername(client->socket, (struct sockaddr*)&client_addr, &addr_len) == 0) {
        inet_ntop(AF_INET, &client_addr.sin_addr, client_ip, INET_ADDRSTRLEN);
    } else {
        strcpy(client_ip, "unknown");
    }
    
    log_with_timestamp("新客户端连接，IP: %s，线程ID: %lu，套接字: %d\n", 
                     client_ip, client->thread, client->socket);
    
    while (client->active) {
        // 检查SSL连接状态
        if (SSL_get_shutdown(client->ssl) & SSL_RECEIVED_SHUTDOWN) {
            printf("检测到SSL关闭信号\n");
            break;
        }
        
        // 设置套接字超时
        struct timeval timeout;
        timeout.tv_sec = 30;  // 30秒超时
        timeout.tv_usec = 0;
        setsockopt(client->socket, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout));
        
        // 检查连接是否仍然有效
        if (!client->active || client->socket < 0) {
            log_with_timestamp("客户端连接已失效，退出处理，IP: %s\n", client_ip);
            goto cleanup;
        }
        
        // 接收客户端消息
        bytes_received = SSL_read(client->ssl, buffer, MAX_MESSAGE_SIZE);
        
        if (bytes_received <= 0) {
            int ssl_error = SSL_get_error(client->ssl, bytes_received);
            switch (ssl_error) {
                case SSL_ERROR_ZERO_RETURN:
                    log_with_timestamp("客户端正常断开连接，IP: %s\n", client_ip);
                    goto cleanup;
                case SSL_ERROR_WANT_READ:
                    printf("SSL需要更多数据，继续等待...\n");
                    continue;
                case SSL_ERROR_WANT_WRITE:
                    printf("SSL需要写入数据，继续等待...\n");
                    continue;
                case SSL_ERROR_SSL:
                    // 检查是否是客户端关闭连接导致的错误
                    if (bytes_received == 0) {
                        log_with_timestamp("客户端关闭连接，IP: %s\n", client_ip);
                    } else {
                        log_with_timestamp("SSL协议错误，可能是数据格式问题或连接异常，IP: %s\n", client_ip);
                        ERR_print_errors_fp(stderr);
                    }
                    // 尝试清理SSL错误队列
                    ERR_clear_error();
                    goto cleanup;
                case SSL_ERROR_SYSCALL:
                    if (bytes_received == 0) {
                        log_with_timestamp("客户端关闭连接（系统调用），IP: %s\n", client_ip);
                    } else {
                        log_with_timestamp("SSL系统调用错误，可能是网络问题，IP: %s\n", client_ip);
                    }
                    // 不要直接break，让连接正常清理
                    goto cleanup;
                default:
                    log_with_timestamp("SSL读取错误: %d，IP: %s\n", ssl_error, client_ip);
                    goto cleanup;
            }
            break;
        }
        
        buffer[bytes_received] = '\0';
        
        // 检查消息长度
        if (bytes_received >= MAX_MESSAGE_SIZE - 1) {
            printf("警告: 消息可能被截断，长度: %d\n", bytes_received);
        }
        
        log_with_timestamp("收到客户端消息，IP: %s，长度: %d，内容: %s\n", 
                         client_ip, bytes_received, buffer);
        
        // 计算MD5
        calculate_md5(buffer, md5_result);
        log_with_timestamp("MD5计算完成，IP: %s，结果: %s\n", client_ip, md5_result);
        
        // 发送MD5结果给客户端
        if (SSL_write(client->ssl, md5_result, strlen(md5_result)) <= 0) {
            int write_error = SSL_get_error(client->ssl, -1);
            log_with_timestamp("发送MD5结果失败，SSL错误: %d，IP: %s\n", write_error, client_ip);
            goto cleanup;
        }
        
        log_with_timestamp("已发送MD5结果给客户端，IP: %s\n", client_ip);
        connection_success = 1;  // 标记连接成功
    }
    
cleanup:
    // 更新IP统计
    update_ip_stats(client_ip, connection_success);
    
    // 清理客户端连接
    log_with_timestamp("开始清理客户端连接，IP: %s，套接字: %d\n", client_ip, client->socket);
    
    if (client->ssl) {
        SSL_shutdown(client->ssl);
        SSL_free(client->ssl);
        client->ssl = NULL;
    }
    
    if (client->socket >= 0) {
        close(client->socket);
        client->socket = -1;
    }
    
    pthread_mutex_lock(&client_mutex);
    client->active = 0;
    client_count--;
    pthread_mutex_unlock(&client_mutex);
    
    log_with_timestamp("客户端连接已关闭，IP: %s，当前连接数: %d\n", client_ip, client_count);
    return NULL;
}

// 信号处理函数
void signal_handler(int sig) {
    printf("\n收到信号 %d，正在关闭服务器...\n", sig);
    
    // 关闭所有客户端连接
    for (int i = 0; i < MAX_CLIENTS; i++) {
        if (clients[i].active) {
            clients[i].active = 0;
            if (clients[i].ssl) {
                SSL_shutdown(clients[i].ssl);
                SSL_free(clients[i].ssl);
            }
            if (clients[i].socket > 0) {
                close(clients[i].socket);
            }
        }
    }
    
    if (server_socket > 0) {
        close(server_socket);
    }
    
    if (ssl_ctx) {
        SSL_CTX_free(ssl_ctx);
    }
    
    cleanup_openssl();
    exit(0);
}

int main(int argc, char *argv[]) {
    int port = 8443; // 默认端口
    struct sockaddr_in server_addr, client_addr;
    socklen_t client_len = sizeof(client_addr);
    int opt = 1;
    int c;
    
    // 解析命令行参数
    while ((c = getopt(argc, argv, "p:h")) != -1) {
        switch (c) {
            case 'p':
                port = atoi(optarg);
                if (port <= 0 || port > 65535) {
                    printf("错误: 端口号必须在1-65535之间\n");
                    exit(EXIT_FAILURE);
                }
                break;
            case 'h':
                printf("用法: %s [-p 端口] [-h]\n", argv[0]);
                printf("选项:\n");
                printf("  -p 端口    指定服务器端口 (默认: 8443)\n");
                printf("  -h         显示此帮助信息\n");
                printf("\n示例:\n");
                printf("  %s -p 9000    # 在端口9000运行服务器\n", argv[0]);
                printf("  %s             # 在默认端口8443运行服务器\n", argv[0]);
                exit(EXIT_SUCCESS);
            case '?':
                printf("使用 -h 查看帮助信息\n");
                exit(EXIT_FAILURE);
            default:
                abort();
        }
    }
    
    log_with_timestamp("启动TLS服务器，端口: %d\n", port);
    
    // 设置信号处理
    signal(SIGINT, signal_handler);
    signal(SIGTERM, signal_handler);
    
    // 初始化OpenSSL
    init_openssl();
    
    // 创建SSL上下文
    ssl_ctx = create_context();
    
    // 创建服务器套接字
    server_socket = socket(AF_INET, SOCK_STREAM, 0);
    if (server_socket < 0) {
        perror("创建套接字失败");
        exit(EXIT_FAILURE);
    }
    
    // 设置套接字选项
    if (setsockopt(server_socket, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)) < 0) {
        perror("设置套接字选项失败");
        exit(EXIT_FAILURE);
    }
    
    // 绑定地址
    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_addr.s_addr = INADDR_ANY;
    server_addr.sin_port = htons(port);
    
    if (bind(server_socket, (struct sockaddr *)&server_addr, sizeof(server_addr)) < 0) {
        perror("绑定地址失败");
        exit(EXIT_FAILURE);
    }
    
    // 开始监听
    if (listen(server_socket, 10) < 0) {
        perror("监听失败");
        exit(EXIT_FAILURE);
    }
    
    log_with_timestamp("服务器已启动，等待客户端连接...\n");
    log_with_timestamp("支持TLS 1.2和1.3\n");
    log_with_timestamp("按Ctrl+C停止服务器\n");
    
    // 初始化客户端数组
    for (int i = 0; i < MAX_CLIENTS; i++) {
        clients[i].socket = -1;
        clients[i].ssl = NULL;
        clients[i].active = 0;
    }
    
    // 主循环：接受客户端连接
    while (1) {
        int client_socket = accept(server_socket, (struct sockaddr *)&client_addr, &client_len);
        if (client_socket < 0) {
            perror("接受连接失败");
            continue;
        }
        
        log_with_timestamp("新连接来自: %s:%d\n", inet_ntoa(client_addr.sin_addr), ntohs(client_addr.sin_port));
        log_with_timestamp("当前活跃连接数: %d/%d\n", client_count, MAX_CLIENTS);
        
        // 检查是否达到最大客户端数
        if (client_count >= MAX_CLIENTS) {
            log_with_timestamp("已达到最大客户端连接数，拒绝新连接\n");
            close(client_socket);
            continue;
        }
        
        // 创建SSL连接
        SSL *ssl = SSL_new(ssl_ctx);
        SSL_set_fd(ssl, client_socket);
        
        if (SSL_accept(ssl) <= 0) {
            log_with_timestamp("SSL握手失败，IP: %s\n", inet_ntoa(client_addr.sin_addr));
            ERR_print_errors_fp(stderr);
            SSL_free(ssl);
            close(client_socket);
            continue;
        }
        
        log_with_timestamp("SSL握手成功，IP: %s，TLS版本: %s\n", 
                         inet_ntoa(client_addr.sin_addr), SSL_get_version(ssl));
        
        // 查找空闲的客户端槽位
        int client_index = -1;
        pthread_mutex_lock(&client_mutex);
        for (int i = 0; i < MAX_CLIENTS; i++) {
            if (!clients[i].active) {
                client_index = i;
                break;
            }
        }
        
        if (client_index == -1) {
            log_with_timestamp("没有可用的客户端槽位\n");
            SSL_free(ssl);
            close(client_socket);
            pthread_mutex_unlock(&client_mutex);
            continue;
        }
        
        // 设置客户端信息
        clients[client_index].socket = client_socket;
        clients[client_index].ssl = ssl;
        clients[client_index].active = 1;
        client_count++;
        
        // 创建处理线程
        if (pthread_create(&clients[client_index].thread, NULL, handle_client, &clients[client_index]) != 0) {
            log_with_timestamp("创建线程失败\n");
            clients[client_index].active = 0;
            client_count--;
            SSL_free(ssl);
            close(client_socket);
            pthread_mutex_unlock(&client_mutex);
            continue;
        } else {
            // 分离线程，让系统自动回收线程资源
            pthread_detach(clients[client_index].thread);
            log_with_timestamp("客户端处理线程已创建，当前连接数: %d\n", client_count);
        }
        
        pthread_mutex_unlock(&client_mutex);
    }
    
    return 0;
}