#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <pthread.h>
#include <signal.h>

#define MAX_CLIENTS 10
#define BUFFER_SIZE 1024
#define USERNAME_LENGTH 50
#define PASSWORD_LENGTH 50
#define FILENAME "./users.txt"

typedef struct {
    int socket;
    char username[USERNAME_LENGTH];
    pthread_t thread;
    int active;
} Client;

Client clients[MAX_CLIENTS];
pthread_mutex_t clients_mutex = PTHREAD_MUTEX_INITIALIZER;
int server_socket;

// 函数声明
void handle_signal(int sig);
void *handle_client(void *arg);
int find_client_index(int socket);
int find_client_index_by_username(const char *username);
void send_private_message(const char *sender, const char *recipient, const char *message);
void list_users(int client_socket);
int authenticate_user(const char *username, const char *password);
int register_user(const char *username, const char *password);
void save_user(const char *username, const char *password);

int main() {
    int client_socket;
    struct sockaddr_in server_addr, client_addr;
    socklen_t client_addr_len = sizeof(client_addr);

    // 创建socket
    server_socket = socket(AF_INET, SOCK_STREAM, 0);
    if (server_socket == -1) {
        perror("Socket creation failed");
        exit(EXIT_FAILURE);
    }

    // 设置socket选项
    int opt = 1;
    if (setsockopt(server_socket, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)) == -1) {
        perror("Setsockopt failed");
        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(8888);

    // 绑定socket
    if (bind(server_socket, (struct sockaddr *)&server_addr, sizeof(server_addr)) == -1) {
        perror("Bind failed");
        exit(EXIT_FAILURE);
    }

    // 监听连接
    if (listen(server_socket, 5) == -1) {
        perror("Listen failed");
        exit(EXIT_FAILURE);
    }

    printf("Server started. Waiting for connections...\n");

    // 初始化客户端数组
    for (int i = 0; i < MAX_CLIENTS; i++) {
        clients[i].socket = -1;
        clients[i].active = 0;
    }

    // 设置信号处理
    signal(SIGINT, handle_signal);

    // 接受客户端连接
    while (1) {
        client_socket = accept(server_socket, (struct sockaddr *)&client_addr, &client_addr_len);
        if (client_socket == -1) {
            perror("Accept failed");
            continue;
        }

        // 查找空闲位置
        int i;
        for (i = 0; i < MAX_CLIENTS; i++) {
            if (clients[i].socket == -1) {
                pthread_mutex_lock(&clients_mutex);
                clients[i].socket = client_socket;
                clients[i].active = 1;
                pthread_mutex_unlock(&clients_mutex);

                // 创建线程处理客户端
                if (pthread_create(&clients[i].thread, NULL, handle_client, &clients[i]) != 0) {
                    perror("Thread creation failed");
                    pthread_mutex_lock(&clients_mutex);
                    clients[i].socket = -1;
                    clients[i].active = 0;
                    pthread_mutex_unlock(&clients_mutex);
                    close(client_socket);
                }
                pthread_detach(clients[i].thread);
                break;
            }
        }

        if (i == MAX_CLIENTS) {
            // 服务器已满
            char response[] = "Server is full. Try again later.\n";
            send(client_socket, response, strlen(response), 0);
            close(client_socket);
        }
    }

    // 关闭服务器socket
    close(server_socket);
    return 0;
}

void handle_signal(int sig) {
    printf("\nShutting down server...\n");
    
    // 关闭所有客户端socket
    pthread_mutex_lock(&clients_mutex);
    for (int i = 0; i < MAX_CLIENTS; i++) {
        if (clients[i].socket != -1) {
            close(clients[i].socket);
            clients[i].socket = -1;
            clients[i].active = 0;
        }
    }
    pthread_mutex_unlock(&clients_mutex);
    
    // 关闭服务器socket
    close(server_socket);
    exit(EXIT_SUCCESS);
}

void *handle_client(void *arg) {
    Client *client = (Client *)arg;
    char buffer[BUFFER_SIZE];
    int bytes_received;

    // 登录/注册流程
    while (1) {
        // 发送菜单
        char menu[] = "1. Login\n2. Register\nEnter choice: ";
        send(client->socket, menu, strlen(menu), 0);

        // 接收选择
        memset(buffer, 0, BUFFER_SIZE);
        bytes_received = recv(client->socket, buffer, BUFFER_SIZE - 1, 0);
        if (bytes_received <= 0) {
            pthread_mutex_lock(&clients_mutex);
            client->active = 0;
            pthread_mutex_unlock(&clients_mutex);
            close(client->socket);
            pthread_exit(NULL);
        }

        int choice = atoi(buffer);

        // 处理登录
        if (choice == 1) {
            // 请求用户名
            char username_prompt[] = "Username: ";
            send(client->socket, username_prompt, strlen(username_prompt), 0);

            // 接收用户名
            memset(buffer, 0, BUFFER_SIZE);
            bytes_received = recv(client->socket, buffer, BUFFER_SIZE - 1, 0);
            if (bytes_received <= 0) {
                pthread_mutex_lock(&clients_mutex);
                client->active = 0;
                pthread_mutex_unlock(&clients_mutex);
                close(client->socket);
                pthread_exit(NULL);
            }
            char username[USERNAME_LENGTH];
            strncpy(username, buffer, USERNAME_LENGTH - 1);

            // 请求密码
            char password_prompt[] = "Password: ";
            send(client->socket, password_prompt, strlen(password_prompt), 0);

            // 接收密码
            memset(buffer, 0, BUFFER_SIZE);
            bytes_received = recv(client->socket, buffer, BUFFER_SIZE - 1, 0);
            if (bytes_received <= 0) {
                pthread_mutex_lock(&clients_mutex);
                client->active = 0;
                pthread_mutex_unlock(&clients_mutex);
                close(client->socket);
                pthread_exit(NULL);
            }
            char password[PASSWORD_LENGTH];
            strncpy(password, buffer, PASSWORD_LENGTH - 1);

            // 验证用户
            if (authenticate_user(username, password)) {
                // 检查用户名是否已被使用
                int username_in_use = 0;
                pthread_mutex_lock(&clients_mutex);
                for (int i = 0; i < MAX_CLIENTS; i++) {
                    if (i != find_client_index(client->socket) && clients[i].active && 
                        strcmp(clients[i].username, username) == 0) {
                        username_in_use = 1;
                        break;
                    }
                }
                pthread_mutex_unlock(&clients_mutex);

                if (username_in_use) {
                    char response[] = "Username already in use. Please try again.\n";
                    send(client->socket, response, strlen(response), 0);
                } else {
                    // 登录成功
                    strcpy(client->username, username);
                    char success[] = "Login successful. Welcome to the chat!\n";
                    send(client->socket, success, strlen(success), 0);
                    
                    // 通知用户可用命令
                    char help[] = "Available commands:\n/private <username> <message> - Send private message\n/list - List online users\n/quit - Exit chat\n";
                    send(client->socket, help, strlen(help), 0);
                    
                    break;
                }
            } else {
                char response[] = "Authentication failed. Please try again.\n";
                send(client->socket, response, strlen(response), 0);
            }
        }
        // 处理注册
        else if (choice == 2) {
            // 请求用户名
            char username_prompt[] = "Enter username: ";
            send(client->socket, username_prompt, strlen(username_prompt), 0);

            // 接收用户名
            memset(buffer, 0, BUFFER_SIZE);
            bytes_received = recv(client->socket, buffer, BUFFER_SIZE - 1, 0);
            if (bytes_received <= 0) {
                pthread_mutex_lock(&clients_mutex);
                client->active = 0;
                pthread_mutex_unlock(&clients_mutex);
                close(client->socket);
                pthread_exit(NULL);
            }
            char username[USERNAME_LENGTH];
            strncpy(username, buffer, USERNAME_LENGTH - 1);

            // 请求密码
            char password_prompt[] = "Enter password: ";
            send(client->socket, password_prompt, strlen(password_prompt), 0);

            // 接收密码
            memset(buffer, 0, BUFFER_SIZE);
            bytes_received = recv(client->socket, buffer, BUFFER_SIZE - 1, 0);
            if (bytes_received <= 0) {
                pthread_mutex_lock(&clients_mutex);
                client->active = 0;
                pthread_mutex_unlock(&clients_mutex);
                close(client->socket);
                pthread_exit(NULL);
            }
            char password[PASSWORD_LENGTH];
            strncpy(password, buffer, PASSWORD_LENGTH - 1);

            // 注册用户
            if (register_user(username, password)) {
                char success[] = "Registration successful. You can now login.\n";
                send(client->socket, success, strlen(success), 0);
            } else {
                char response[] = "Registration failed. Username already exists.\n";
                send(client->socket, response, strlen(response), 0);
            }
        }
        else {
            char response[] = "Invalid choice. Please try again.\n";
            send(client->socket, response, strlen(response), 0);
        }
    }

    // 聊天循环
    while (1) {
        memset(buffer, 0, BUFFER_SIZE);
        bytes_received = recv(client->socket, buffer, BUFFER_SIZE - 1, 0);
        
        if (bytes_received <= 0) {
            // 客户端断开连接
            break;
        }

        // 检查是否是退出命令
        if (strncmp(buffer, "/quit", 5) == 0) {
            break;
        }
        
        // 检查是否是列出用户命令
        if (strncmp(buffer, "/list", 5) == 0) {
            list_users(client->socket);
            continue;
        }
        
        // 检查是否是私聊命令
        if (strncmp(buffer, "/private", 7) == 0) {
            char recipient[USERNAME_LENGTH];
            char message[BUFFER_SIZE - USERNAME_LENGTH - 8];
            
            // 解析命令
            sscanf(buffer, "/private %s %[^\n]", recipient, message);
            
            // 发送私聊消息
            send_private_message(client->username, recipient, message);
            continue;
        }
        
        // 发送帮助信息
        char help[] = "Invalid command. Available commands:\n/private <username> <message> - Send private message\n/list - List online users\n/quit - Exit chat\n";
        send(client->socket, help, strlen(help), 0);
    }

    // 客户端断开连接
    char leave_message[BUFFER_SIZE];
    snprintf(leave_message, BUFFER_SIZE, "%s has left the chat.\n", client->username);
    pthread_mutex_lock(&clients_mutex);
    for (int i = 0; i < MAX_CLIENTS; i++) {
        if (i != find_client_index(client->socket) && clients[i].active) {
            send(clients[i].socket, leave_message, strlen(leave_message), 0);
        }
    }
    pthread_mutex_unlock(&clients_mutex);

    // 关闭客户端socket
    pthread_mutex_lock(&clients_mutex);
    close(client->socket);
    client->socket = -1;
    client->active = 0;
    pthread_mutex_unlock(&clients_mutex);

    pthread_exit(NULL);
}

int find_client_index(int socket) {
    for (int i = 0; i < MAX_CLIENTS; i++) {
        if (clients[i].socket == socket) {
            return i;
        }
    }
    return -1;
}

int find_client_index_by_username(const char *username) {
    pthread_mutex_lock(&clients_mutex);
    for (int i = 0; i < MAX_CLIENTS; i++) {
        if (clients[i].active && strcmp(clients[i].username, username) == 0) {
            pthread_mutex_unlock(&clients_mutex);
            return i;
        }
    }
    pthread_mutex_unlock(&clients_mutex);
    return -1;
}

void send_private_message(const char *sender, const char *recipient, const char *message) {
    int recipient_index = find_client_index_by_username(recipient);
    
    if (recipient_index == -1) {
        char error[BUFFER_SIZE];
        snprintf(error, BUFFER_SIZE, "Error: User '%s' not found or offline.\n", recipient);
        
        int sender_index = find_client_index_by_username(sender);
        if (sender_index != -1) {
            send(clients[sender_index].socket, error, strlen(error), 0);
        }
        return;
    }
    
    // 发送消息给接收者
    char private_message[BUFFER_SIZE];
    snprintf(private_message, BUFFER_SIZE, "[Private from %s]: %s\n", sender, message);
    send(clients[recipient_index].socket, private_message, strlen(private_message), 0);
    
    // 发送确认给发送者
    char confirmation[BUFFER_SIZE];
    snprintf(confirmation, BUFFER_SIZE, "[Private to %s]: %s\n", recipient, message);
    int sender_index = find_client_index_by_username(sender);
    if (sender_index != -1) {
        send(clients[sender_index].socket, confirmation, strlen(confirmation), 0);
    }
}

void list_users(int client_socket) {
    char response[BUFFER_SIZE] = "Online users:\n";
    
    pthread_mutex_lock(&clients_mutex);
    for (int i = 0; i < MAX_CLIENTS; i++) {
        if (clients[i].active) {
            strncat(response, clients[i].username, USERNAME_LENGTH - strlen(response) - 2);
            strcat(response, "\n");
        }
    }
    pthread_mutex_unlock(&clients_mutex);
    
    send(client_socket, response, strlen(response), 0);
}

int authenticate_user(const char *username, const char *password) {
    FILE *file = fopen(FILENAME, "r");
    if (file == NULL) {
        return 0;
    }

    char stored_username[USERNAME_LENGTH];
    char stored_password[PASSWORD_LENGTH];
    int found = 0;

    while (fscanf(file, "%s %s", stored_username, stored_password) != EOF) {
        if (strcmp(username, stored_username) == 0 && strcmp(password, stored_password) == 0) {
            found = 1;
            break;
        }
    }

    fclose(file);
    return found;
}

int register_user(const char *username, const char *password) {
    FILE *file = fopen(FILENAME, "r");
    if (file == NULL) {
        file = fopen(FILENAME, "w");
        if (file == NULL) {
            return 0;
        }
        fclose(file);
        file = fopen(FILENAME, "r");
    }

    char stored_username[USERNAME_LENGTH];
    while (fscanf(file, "%s", stored_username) != EOF) {
        if (strcmp(username, stored_username) == 0) {
            fclose(file);
            return 0; // 用户名已存在
        }
        // 跳过密码
        fscanf(file, "%*s");
    }
    fclose(file);

    // 添加新用户
    save_user(username, password);
    return 1;
}

void save_user(const char *username, const char *password) {
    FILE *file = fopen(FILENAME, "a");
    if (file != NULL) {
        fprintf(file, "%s %s\n", username, password);
        fclose(file);
    }
}