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

#define BUFFER_SIZE 1024
#define SERVER_IP "127.0.0.1"
#define PORT 8888

int client_socket;
pthread_t receive_thread;
int running = 1;

// 函数声明
void *receive_messages(void *arg);
void handle_signal(int sig);
void login_register_menu();

int main() {
    struct sockaddr_in server_addr;

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

    // 准备服务器地址
    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(PORT);

    // 将IPv4地址从点分十进制转换为二进制形式
    if (inet_pton(AF_INET, SERVER_IP, &server_addr.sin_addr) <= 0) {
        perror("Invalid address/ Address not supported");
        exit(EXIT_FAILURE);
    }

    // 连接到服务器
    if (connect(client_socket, (struct sockaddr *)&server_addr, sizeof(server_addr)) == -1) {
        perror("Connection failed");
        exit(EXIT_FAILURE);
    }

    printf("Connected to server.\n");

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

    // 创建接收消息的线程
    if (pthread_create(&receive_thread, NULL, receive_messages, NULL) != 0) {
        perror("Thread creation failed");
        close(client_socket);
        exit(EXIT_FAILURE);
    }

    // 登录/注册菜单
    login_register_menu();

    // 聊天循环
    printf("\nWelcome to the chat! Type /help for commands.\n");
    char message[BUFFER_SIZE];
    while (running) {
        memset(message, 0, BUFFER_SIZE);
        fgets(message, BUFFER_SIZE - 1, stdin);
        
        // 移除换行符
        size_t len = strlen(message);
        if (len > 0 && message[len - 1] == '\n') {
            message[len - 1] = '\0';
        }

        // 检查是否是退出命令
        if (strcmp(message, "/quit") == 0) {
            running = 0;
            break;
        }

        // 发送消息到服务器
        if (send(client_socket, message, strlen(message), 0) == -1) {
            perror("Send failed");
            running = 0;
        }
    }

    // 等待接收线程结束
    pthread_join(receive_thread, NULL);

    // 关闭socket
    close(client_socket);
    printf("Disconnected from server.\n");

    return 0;
}

void *receive_messages(void *arg) {
    char buffer[BUFFER_SIZE];
    int bytes_received;

    while (running) {
        memset(buffer, 0, BUFFER_SIZE);
        bytes_received = recv(client_socket, buffer, BUFFER_SIZE - 1, 0);
        
        if (bytes_received <= 0) {
            // 服务器断开连接
            printf("Server disconnected.\n");
            running = 0;
            break;
        }

        // 打印接收到的消息
        printf("%s", buffer);
    }

    pthread_exit(NULL);
}

void handle_signal(int sig) {
    printf("\nExiting...\n");
    running = 0;
    close(client_socket);
    exit(EXIT_SUCCESS);
}

void login_register_menu() {
    char buffer[BUFFER_SIZE];
    int choice;
    char username[BUFFER_SIZE];
    char password[BUFFER_SIZE];

    while (running) {
        // 接收菜单
        memset(buffer, 0, BUFFER_SIZE);
        if (recv(client_socket, buffer, BUFFER_SIZE - 1, 0) <= 0) {
            printf("Server disconnected.\n");
            running = 0;
            return;
        }
        printf("%s", buffer);

        // 发送选择
        scanf("%d", &choice);
        getchar(); // 消耗换行符
        sprintf(buffer, "%d", choice);
        if (send(client_socket, buffer, strlen(buffer), 0) == -1) {
            perror("Send failed");
            running = 0;
            return;
        }

        // 处理登录
        if (choice == 1) {
            // 接收用户名提示
            memset(buffer, 0, BUFFER_SIZE);
            if (recv(client_socket, buffer, BUFFER_SIZE - 1, 0) <= 0) {
                printf("Server disconnected.\n");
                running = 0;
                return;
            }
            printf("%s", buffer);

            // 发送用户名
            fgets(username, BUFFER_SIZE - 1, stdin);
            username[strcspn(username, "\n")] = 0; // 移除换行符
            if (send(client_socket, username, strlen(username), 0) == -1) {
                perror("Send failed");
                running = 0;
                return;
            }

            // 接收密码提示
            memset(buffer, 0, BUFFER_SIZE);
            if (recv(client_socket, buffer, BUFFER_SIZE - 1, 0) <= 0) {
                printf("Server disconnected.\n");
                running = 0;
                return;
            }
            printf("%s", buffer);

            // 发送密码
            fgets(password, BUFFER_SIZE - 1, stdin);
            password[strcspn(password, "\n")] = 0; // 移除换行符
            if (send(client_socket, password, strlen(password), 0) == -1) {
                perror("Send failed");
                running = 0;
                return;
            }

            // 接收登录结果
            memset(buffer, 0, BUFFER_SIZE);
            if (recv(client_socket, buffer, BUFFER_SIZE - 1, 0) <= 0) {
                printf("Server disconnected.\n");
                running = 0;
                return;
            }
            printf("%s", buffer);

            // 检查是否登录成功
            if (strstr(buffer, "Login successful") != NULL) {
                // 接收帮助信息
                memset(buffer, 0, BUFFER_SIZE);
                if (recv(client_socket, buffer, BUFFER_SIZE - 1, 0) > 0) {
                    printf("%s", buffer);
                }
                break;
            }
        }
        // 处理注册
        else if (choice == 2) {
            // 接收用户名提示
            memset(buffer, 0, BUFFER_SIZE);
            if (recv(client_socket, buffer, BUFFER_SIZE - 1, 0) <= 0) {
                printf("Server disconnected.\n");
                running = 0;
                return;
            }
            printf("%s", buffer);

            // 发送用户名
            fgets(username, BUFFER_SIZE - 1, stdin);
            username[strcspn(username, "\n")] = 0; // 移除换行符
            if (send(client_socket, username, strlen(username), 0) == -1) {
                perror("Send failed");
                running = 0;
                return;
            }

            // 接收密码提示
            memset(buffer, 0, BUFFER_SIZE);
            if (recv(client_socket, buffer, BUFFER_SIZE - 1, 0) <= 0) {
                printf("Server disconnected.\n");
                running = 0;
                return;
            }
            printf("%s", buffer);

            // 发送密码
            fgets(password, BUFFER_SIZE - 1, stdin);
            password[strcspn(password, "\n")] = 0; // 移除换行符
            if (send(client_socket, password, strlen(password), 0) == -1) {
                perror("Send failed");
                running = 0;
                return;
            }

            // 接收注册结果
            memset(buffer, 0, BUFFER_SIZE);
            if (recv(client_socket, buffer, BUFFER_SIZE - 1, 0) <= 0) {
                printf("Server disconnected.\n");
                running = 0;
                return;
            }
            printf("%s", buffer);
        }
        else {
            // 接收无效选择提示
            memset(buffer, 0, BUFFER_SIZE);
            if (recv(client_socket, buffer, BUFFER_SIZE - 1, 0) <= 0) {
                printf("Server disconnected.\n");
                running = 0;
                return;
            }
            printf("%s", buffer);
        }
    }
}