#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 <pthread.h>
#include <netinet/tcp.h>

#define MAX_BUFFER 1024
#define MAX_USERS 100
#define MAX_FRIENDS 50
#define PORT 8888

// 用户结构体
typedef struct {
    char username[50];
    char password[50];
    int friends[MAX_FRIENDS];
    int friend_count;
    int online;
    int socket;  // 添加socket字段
} User;

// 全局变量
User users[MAX_USERS];
int user_count = 0;
pthread_mutex_t users_mutex = PTHREAD_MUTEX_INITIALIZER;

// 函数声明
void load_users();
void save_users();
int find_user(const char* username);
void* handle_client(void* arg);
void process_login(int client_socket, char* buffer);
void process_register(int client_socket, char* buffer);
void process_add_friend(int client_socket, char* buffer);
void process_remove_friend(int client_socket, char* buffer);
void process_chat(int client_socket, char* buffer);
void display_online_users();  // 新增函数声明
void process_get_friends(int client_socket, char* buffer);

int main() {
    int server_socket, client_socket;
    struct sockaddr_in server_addr, client_addr;
    socklen_t client_len = sizeof(client_addr);
    pthread_t thread_id;

    // 加载用户数据
    load_users();

    // 创建socket
    server_socket = socket(AF_INET, SOCK_STREAM, 0);
    if (server_socket < 0) {
        perror("创建Socket失败");
        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);

    // 绑定socket
    if (bind(server_socket, (struct sockaddr*)&server_addr, sizeof(server_addr)) < 0) {
        perror("绑定失败");
        exit(EXIT_FAILURE);
    }

    // 监听连接
    if (listen(server_socket, 5) < 0) {
        perror("监听失败");
        exit(EXIT_FAILURE);
    }

    printf("服务器已启动，监听端口 %d\n", PORT);

    // 接受客户端连接
    while (1) {
        client_socket = accept(server_socket, (struct sockaddr*)&client_addr, &client_len);
        if (client_socket < 0) {
            perror("接受连接失败");
            continue;
        }

        printf("新客户端连接：%s:%d\n",inet_ntoa(client_addr.sin_addr), ntohs(client_addr.sin_port));

        // 为每个客户端创建新线程
        if (pthread_create(&thread_id, NULL, handle_client, (void*)&client_socket) != 0) {
            perror("创建线程失败");
            close(client_socket);
        }
    }

    close(server_socket);
    return 0;
}

void load_users() {
    FILE* file = fopen("Users.txt", "r");
    if (file == NULL) {
        printf("未找到用户文件，将创建新文件。\n");
        return;
    }

    while (fscanf(file, "%s %s %d", users[user_count].username, 
                  users[user_count].password, &users[user_count].friend_count) == 3) {
        for (int i = 0; i < users[user_count].friend_count; i++) {
            fscanf(file, "%d", &users[user_count].friends[i]);
        }
        users[user_count].online = 0;
        users[user_count].socket = -1;
        user_count++;
    }

    fclose(file);
}

void save_users() {
    FILE* file = fopen("Users.txt", "w");
    if (file == NULL) {
        perror("无法打开用户文件");
        return;
    }

    for (int i = 0; i < user_count; i++) {
        fprintf(file, "%s %s %d", users[i].username, 
                users[i].password, users[i].friend_count);
        for (int j = 0; j < users[i].friend_count; j++) {
            fprintf(file, " %d", users[i].friends[j]);
        }
        fprintf(file, "\n");
    }

    fclose(file);
}

int find_user(const char* username) {
    for (int i = 0; i < user_count; i++) {
        if (strcmp(users[i].username, username) == 0) {
            return i;
        }
    }
    return -1;
}

void* handle_client(void* arg) {
    int client_socket = *((int*)arg);
    char buffer[MAX_BUFFER];
    int user_index = -1;

    // 设置socket选项，启用保活机制
    int keepalive = 1;
    int keepidle = 5;  // 5秒没有数据交互就开始探测
    int keepinterval = 2;  // 探测间隔2秒
    int keepcount = 3;  // 探测3次无响应就断开

    setsockopt(client_socket, SOL_SOCKET, SO_KEEPALIVE, &keepalive, sizeof(keepalive));
    setsockopt(client_socket, IPPROTO_TCP, TCP_KEEPIDLE, &keepidle, sizeof(keepidle));
    setsockopt(client_socket, IPPROTO_TCP, TCP_KEEPINTVL, &keepinterval, sizeof(keepinterval));
    setsockopt(client_socket, IPPROTO_TCP, TCP_KEEPCNT, &keepcount, sizeof(keepcount));

    while (1) {
        memset(buffer, 0, MAX_BUFFER);
        int bytes_received = recv(client_socket, buffer, MAX_BUFFER, 0);
        
        // 检查连接状态
        if (bytes_received <= 0) {
            pthread_mutex_lock(&users_mutex);
            // 查找并更新断开连接的用户状态
            for (int i = 0; i < user_count; i++) {
                if (users[i].socket == client_socket) {
                    user_index = i;
                    users[i].online = 0;
                    users[i].socket = -1;
                    printf("\n用户 %s 断开连接\n", users[i].username);
                    
                    // 通知所有在线好友
                    for (int j = 0; j < users[i].friend_count; j++) {
                        int friend_index = users[i].friends[j];
                        if (users[friend_index].online && users[friend_index].socket != -1) {
                            char notify[MAX_BUFFER];
                            sprintf(notify, "系统消息: 您的好友 %s 已离线", users[i].username);
                            send(users[friend_index].socket, notify, MAX_BUFFER, 0);
                        }
                    }
                    break;
                }
            }
            save_users();
            display_online_users();
            pthread_mutex_unlock(&users_mutex);
            break;
        }

        char command[20];
        sscanf(buffer, "%s", command);

        if (strcmp(command, "LOGIN") == 0) {
            process_login(client_socket, buffer);
            // 更新user_index
            for (int i = 0; i < user_count; i++) {
                if (users[i].socket == client_socket) {
                    user_index = i;
                    break;
                }
            }
        }
        else if (strcmp(command, "REGISTER") == 0) {
            process_register(client_socket, buffer);
        }
        else if (strcmp(command, "ADDFRIEND") == 0) {
            process_add_friend(client_socket, buffer);
        }
        else if (strcmp(command, "REMOVEFRIEND") == 0) {
            process_remove_friend(client_socket, buffer);
        }
        else if (strcmp(command, "CHAT") == 0) {
            process_chat(client_socket, buffer);
        }
        else if (strcmp(command, "LOGOUT") == 0) {
            pthread_mutex_lock(&users_mutex);
            // 查找当前用户
            for (int i = 0; i < user_count; i++) {
                if (users[i].socket == client_socket) {
                    user_index = i;
                    break;
                }
            }
            
            if (user_index != -1) {
                users[user_index].online = 0;
                users[user_index].socket = -1;
                // 通知所有在线好友
                for (int i = 0; i < users[user_index].friend_count; i++) {
                    int friend_index = users[user_index].friends[i];
                    if (users[friend_index].online && users[friend_index].socket != -1) {
                        char notify[MAX_BUFFER];
                        sprintf(notify, "系统消息: 您的好友 %s 已离线", users[user_index].username);
                        send(users[friend_index].socket, notify, MAX_BUFFER, 0);
                    }
                }
                save_users();
                printf("\n用户 %s 退出登录\n", users[user_index].username);
                display_online_users();
            }
            pthread_mutex_unlock(&users_mutex);
            break;
        }
        else if (strcmp(command, "GETFRIENDS") == 0) {
            process_get_friends(client_socket, buffer);
        }
    }

    // 确保在断开连接时清理用户状态
    if (user_index != -1) {
        pthread_mutex_lock(&users_mutex);
        users[user_index].online = 0;
        users[user_index].socket = -1;
        save_users();
        display_online_users();
        pthread_mutex_unlock(&users_mutex);
    }

    close(client_socket);
    return NULL;
}

void process_login(int client_socket, char* buffer) {
    char username[50], password[50];
    sscanf(buffer, "LOGIN %s %s", username, password);

    pthread_mutex_lock(&users_mutex);
    int user_index = find_user(username);
    
    if (user_index == -1) {
        send(client_socket, "FAIL 用户不存在", MAX_BUFFER, 0);
    }
    else if (strcmp(users[user_index].password, password) != 0) {
        send(client_socket, "FAIL 密码错误", MAX_BUFFER, 0);
    }
    else {
        // 如果用户已经在线，先将其标记为离线
        if (users[user_index].online) {
            users[user_index].online = 0;
            users[user_index].socket = -1;
        }
        
        // 更新用户状态
        users[user_index].online = 1;
        users[user_index].socket = client_socket;
        save_users();
        
        char response[MAX_BUFFER];
        snprintf(response, MAX_BUFFER, "SUCCESS %d", user_index);
        send(client_socket, response, MAX_BUFFER, 0);
        printf("\n用户 %s 登录成功\n", username);
        
        // 通知所有在线好友
        for (int i = 0; i < users[user_index].friend_count; i++) {
            int friend_index = users[user_index].friends[i];
            if (users[friend_index].online && users[friend_index].socket != -1) {
                char notify[MAX_BUFFER];
                snprintf(notify, MAX_BUFFER, "系统消息: 您的好友 %s 已上线", username);
                send(users[friend_index].socket, notify, MAX_BUFFER, 0);
            }
        }
        display_online_users();
    }
    pthread_mutex_unlock(&users_mutex);
}

void process_register(int client_socket, char* buffer) {
    char username[50], password[50];
    sscanf(buffer, "REGISTER %s %s", username, password);

    pthread_mutex_lock(&users_mutex);
    if (find_user(username) != -1) {
        send(client_socket, "FAIL 用户名已存在", MAX_BUFFER, 0);
    }
    else if (user_count >= MAX_USERS) {
        send(client_socket, "FAIL 用户数量已达上限", MAX_BUFFER, 0);
    }
    else {
        strcpy(users[user_count].username, username);
        strcpy(users[user_count].password, password);
        users[user_count].friend_count = 0;
        users[user_count].online = 0;
        users[user_count].socket = -1;
        user_count++;
        save_users();
        send(client_socket, "SUCCESS", MAX_BUFFER, 0);
    }
    pthread_mutex_unlock(&users_mutex);
}

void process_add_friend(int client_socket, char* buffer) {
    int user_index;
    char friend_name[50];
    sscanf(buffer, "ADDFRIEND %d %s", &user_index, friend_name);

    pthread_mutex_lock(&users_mutex);
    int friend_index = find_user(friend_name);
    
    if (friend_index == -1) {
        send(client_socket, "FAIL 好友不存在", MAX_BUFFER, 0);
    }
    else if (user_index == friend_index) {
        send(client_socket, "FAIL 不能添加自己为好友", MAX_BUFFER, 0);
    }
    else {
        // 检查是否已经是好友
        for (int i = 0; i < users[user_index].friend_count; i++) {
            if (users[user_index].friends[i] == friend_index) {
                send(client_socket, "FAIL 已经是好友关系", MAX_BUFFER, 0);
                pthread_mutex_unlock(&users_mutex);
                return;
            }
        }

        // 添加好友
        users[user_index].friends[users[user_index].friend_count++] = friend_index;
        users[friend_index].friends[users[friend_index].friend_count++] = user_index;
        save_users();
        send(client_socket, "SUCCESS", MAX_BUFFER, 0);
        
        // 通知被添加的用户
        if (users[friend_index].online && users[friend_index].socket != -1) {
            char notify[MAX_BUFFER];
            snprintf(notify, MAX_BUFFER, "系统消息: %s 添加您为好友", users[user_index].username);
            send(users[friend_index].socket, notify, MAX_BUFFER, 0);
        }
    }
    pthread_mutex_unlock(&users_mutex);
}

void process_remove_friend(int client_socket, char* buffer) {
    int user_index;
    char friend_name[50];
    sscanf(buffer, "REMOVEFRIEND %d %s", &user_index, friend_name);

    pthread_mutex_lock(&users_mutex);
    int friend_index = find_user(friend_name);
    
    if (friend_index == -1) {
        send(client_socket, "FAIL 好友不存在", MAX_BUFFER, 0);
        pthread_mutex_unlock(&users_mutex);
        return;
    }

    int found = 0;
    // 从用户的好友列表中移除
    for (int i = 0; i < users[user_index].friend_count; i++) {
        if (users[user_index].friends[i] == friend_index) {
            for (int j = i; j < users[user_index].friend_count - 1; j++) {
                users[user_index].friends[j] = users[user_index].friends[j + 1];
            }
            users[user_index].friend_count--;
            found = 1;
            break;
        }
    }

    if (found) {
        // 从好友的好友列表中移除
        for (int i = 0; i < users[friend_index].friend_count; i++) {
            if (users[friend_index].friends[i] == user_index) {
                for (int j = i; j < users[friend_index].friend_count - 1; j++) {
                    users[friend_index].friends[j] = users[friend_index].friends[j + 1];
                }
                users[friend_index].friend_count--;
                break;
            }
        }

        save_users();
        
        // 发送成功响应
        send(client_socket, "FRIEND_REMOVE_SUCCESS", MAX_BUFFER, 0);

        // 发送通知给被删除的好友
        if (users[friend_index].online && users[friend_index].socket != -1) {
            char notify[MAX_BUFFER];
            snprintf(notify, MAX_BUFFER, "FRIEND_REMOVE %s 将您从好友列表中移除", 
                    users[user_index].username);
            send(users[friend_index].socket, notify, MAX_BUFFER, 0);
        }
    } else {
        send(client_socket, "FAIL 不是好友关系", MAX_BUFFER, 0);
    }
    
    pthread_mutex_unlock(&users_mutex);
}

void process_chat(int client_socket, char* buffer) {
    int user_index, friend_index;
    char message[MAX_BUFFER];
    sscanf(buffer, "CHAT %d %d %[^\n]", &user_index, &friend_index, message);
    pthread_mutex_lock(&users_mutex);
    // 验证是否是好友关系
    int is_friend = 0;
    for (int i = 0; i < users[user_index].friend_count; i++) {
        if (users[user_index].friends[i] == friend_index) {
            is_friend = 1;
            break;
        }
    }

    if (!is_friend) {
        send(client_socket, "FAIL 不是好友关系", MAX_BUFFER, 0);
    }
    else {
        char response[MAX_BUFFER];
        // 限制 message 长度，防止溢出
        char safe_message[MAX_BUFFER/2];
        strncpy(safe_message, message, sizeof(safe_message) - 1);
        safe_message[sizeof(safe_message) - 1] = '\0';
        int n = snprintf(response, MAX_BUFFER, "CHAT %s: %s", users[user_index].username, safe_message);
        if (n < 0 || n >= MAX_BUFFER) {
            // 被截断，补上 ...
            strcpy(&response[MAX_BUFFER-5], "...");
        }
        // 发送给发送者
        send(client_socket, response, MAX_BUFFER, 0);
        // 发送给接收者
        if (users[friend_index].online && users[friend_index].socket != -1) {
            send(users[friend_index].socket, response, MAX_BUFFER, 0);
        } else {
            // 好友不在线，给发送者提示
            send(client_socket, "系统消息: 对方不在线，消息未送达", MAX_BUFFER, 0);
        }
    }
    pthread_mutex_unlock(&users_mutex);
}

// 新增函数：显示在线用户
void display_online_users() {
    printf("\n当前在线用户：\n");
    int online_count = 0;
    for (int i = 0; i < user_count; i++) {
        if (users[i].online) {
            printf("%d. %s\n", ++online_count, users[i].username);
        }
    }
    if (online_count == 0) {
        printf("暂无在线用户\n");
    }
    printf("------------------------\n");
}

// 新增函数：处理获取好友列表请求
void process_get_friends(int client_socket, char* buffer) {
    int user_index;
    sscanf(buffer, "GETFRIENDS %d", &user_index);

    pthread_mutex_lock(&users_mutex);
    char response[MAX_BUFFER] = "FRIENDLIST ";
    char temp[64];
    
    for (int i = 0; i < users[user_index].friend_count; i++) {
        int friend_index = users[user_index].friends[i];
        int n = snprintf(temp, sizeof(temp), "%s,%d ", users[friend_index].username, friend_index);
        if (n < 0 || n >= sizeof(temp)) {
            strcpy(&temp[sizeof(temp)-5], "...");
        }
        if (strlen(response) + strlen(temp) < MAX_BUFFER) {
            strcat(response, temp);
        } else {
            strcat(response, "...");
            break;
        }
    }
    send(client_socket, response, MAX_BUFFER, 0);
    pthread_mutex_unlock(&users_mutex);
}