#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <strings.h> /* for strncasecmp */
#include <unistd.h>
#include <pthread.h>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <errno.h>
#include <openssl/evp.h>
#include <openssl/err.h>
#include <ctype.h>
#include <time.h>

#define MAX_MESSAGE 1024
#define MAX_USERNAME 32
#define MAX_PASSWORD 64
#define SERVER_IP "127.0.0.1"
#define PORT 8888

/* Color palette for structured display */
#define COLOR_RESET    "\033[0m"
#define COLOR_PROMPT   "\033[36m" /* cyan */
#define COLOR_MSG      "\033[37m" /* white for normal messages */
#define COLOR_USER     "\033[1;33m" /* bold yellow for usernames */
#define COLOR_JOIN     "\033[32m" /* green for joins */
#define COLOR_LEAVE    "\033[31m" /* red for leaves */
#define COLOR_SYSTEM   "\033[35m" /* magenta for system/info */
#define COLOR_COMMAND  "\033[33m" /* yellow for local commands */
#define COLOR_ERROR    "\033[1;31m" /* bright red for errors */

int client_socket;
char username[MAX_USERNAME];
char current_room[MAX_USERNAME];

#define SUPPRESS_WINDOW 2 /* seconds to suppress identical consecutive messages */

pthread_mutex_t room_mutex = PTHREAD_MUTEX_INITIALIZER;
pthread_mutex_t socket_mutex = PTHREAD_MUTEX_INITIALIZER;

/* prevent duplicate consecutive prints (thread-safe) */
static pthread_mutex_t print_mutex = PTHREAD_MUTEX_INITIALIZER;
static char last_printed[MAX_MESSAGE] = {0};
static time_t last_print_ts = 0;

/* helper to rtrim newline characters */
static void rtrim_newlines(char *s) {
    if (!s) return;
    size_t len = strlen(s);
    while (len > 0 && (s[len-1] == '\n' || s[len-1] == '\r')) {
        s[len-1] = '\0';
        len--;
    }
}

/* classify and print a single server message with color */
static int dedupe_and_store(const char *msg) {
    int should_print = 0;
    time_t now = time(NULL);

    pthread_mutex_lock(&print_mutex);
    if (last_printed[0] == '\0' || strcmp(msg, last_printed) != 0) {
        /* new message -> store and allow print */
        strncpy(last_printed, msg, sizeof(last_printed) - 1);
        last_printed[sizeof(last_printed) - 1] = '\0';
        last_print_ts = now;
        should_print = 1;
    } else {
        /* identical to last message: only allow if outside suppression window */
        if (difftime(now, last_print_ts) > SUPPRESS_WINDOW) {
            last_print_ts = now;
            should_print = 1;
        } else {
            should_print = 0;
        }
    }
    pthread_mutex_unlock(&print_mutex);
    return should_print;
}

static void print_colored_message(const char *msg_in) {
    if (!msg_in) return;
    char msg[MAX_MESSAGE];
    strncpy(msg, msg_in, sizeof(msg)-1);
    msg[sizeof(msg)-1] = '\0';
    rtrim_newlines(msg);

    /* dedupe consecutive identical messages */
    if (!dedupe_and_store(msg)) return;

    /* system responses */
    if (strstr(msg, "You joined room:") || strstr(msg, "Available Rooms:") ||
        strstr(msg, "Authentication successful") || strstr(msg, "Welcome,") ||
        strstr(msg, "Registration successful") || strstr(msg, "Server is full") ||
        strncasecmp(msg, "Error", 5) == 0) {
        printf(COLOR_SYSTEM "%s" COLOR_RESET "\n", msg);
        return;
    }

    /* join / leave notifications */
    if (strstr(msg, " joined the room") || strstr(msg, " joined the chat")) {
        printf(COLOR_JOIN "%s" COLOR_RESET "\n", msg);
        return;
    }
    if (strstr(msg, " left the room") || strstr(msg, " left the chat")) {
        printf(COLOR_LEAVE "%s" COLOR_RESET "\n", msg);
        return;
    }

    /* Private message indicator */
    if (strstr(msg, "[PM") || strstr(msg, "PM from") || strstr(msg, "PM to")) {
        printf(COLOR_USER "%s" COLOR_RESET "\n", msg);
        return;
    }

    /* General chat message: try to color username differently when possible */
    /* Expected format: [room] username: message */
    const char *p = strchr(msg, ']');
    const char *colon = strchr(msg, ':');
    if (msg[0] == '[' && p && colon && colon > p) {
        /* print prefix ([room]) + username in USER color + rest in MSG color */
        int prefix_len = (int)(p - msg + 1);
        char prefix[64] = {0};
        strncpy(prefix, msg, prefix_len);
        const char *user_start = p + 2; /* skip "] " */
        int user_len = (int)(colon - user_start);
        if (user_len <= 0 || user_len >= 64) {
            /* fallback */
            printf(COLOR_MSG "%s" COLOR_RESET "\n", msg);
            return;
        }
        char user[64] = {0};
        strncpy(user, user_start, user_len);
        const char *rest = colon + 1;
        printf(COLOR_MSG "%s " COLOR_USER "%s" COLOR_MSG ":%s" COLOR_RESET "\n",
               prefix, user, rest);
        return;
    }

    /* default: regular message color */
    printf(COLOR_MSG "%s" COLOR_RESET "\n", msg);
}

void print_help() {
    printf(COLOR_SYSTEM
           "=== Available Commands ===\n"
           "  /rooms           - List all available chat rooms\n"
           "  /join <room>     - Join or create a chat room\n"
           "  /help            - Show this help message\n"
           "  /quit            - Exit the chat\n\n"
           "Tips:\n"
           "- Messages will be sent to your current room\n"
           "- Your current room is shown in the prompt [room]>\n"
           "- Use /join to switch between rooms\n"
           "- Private messages don't work yet\n"
           COLOR_RESET);
}

void clear_input_line() {
    printf("\033[2K\r"); // Clear the current line
}

void show_prompt() {
    pthread_mutex_lock(&room_mutex);
    printf("\033[2K\r" COLOR_PROMPT "[%s]> " COLOR_RESET, current_room[0] ? current_room : "lobby");
    fflush(stdout);
    pthread_mutex_unlock(&room_mutex);
}

void *receive_messages(void *arg) {
    (void)arg; /* silence unused parameter warning */
    char message[MAX_MESSAGE];
    int read_size;

    while ((read_size = recv(client_socket, message, MAX_MESSAGE - 1, 0)) > 0) {
        message[read_size] = '\0';
        rtrim_newlines(message);     /* remove any trailing newlines sent by server */
        clear_input_line();          /* Clear current input line */
        print_colored_message(message);
        show_prompt();               /* Show the prompt again */
    }

    /* server disconnected or error */
    clear_input_line();
    printf(COLOR_ERROR "Disconnected from server." COLOR_RESET "\n");
    show_prompt();
    return NULL;
}

void hash_password(const char *password, char *hashed, size_t hashed_size) {
    (void)hashed_size; /* silence unused parameter warning */
    EVP_MD_CTX *mdctx;
    const EVP_MD *md;
    unsigned char hash[EVP_MAX_MD_SIZE];
    unsigned int hash_len;

    // Initialize the hashing context
    mdctx = EVP_MD_CTX_new();
    md = EVP_sha256();
    
    if (mdctx == NULL) {
        fprintf(stderr, "Error creating hash context\n");
        return;
    }

    EVP_DigestInit_ex(mdctx, md, NULL);
    EVP_DigestUpdate(mdctx, password, strlen(password));
    EVP_DigestFinal_ex(mdctx, hash, &hash_len);
    EVP_MD_CTX_free(mdctx);
    
    // Convert to hex string
    for(unsigned int i = 0; i < hash_len; i++) {
        snprintf(hashed + (i * 2), 3, "%02x", hash[i]);
    }
    hashed[hash_len * 2] = '\0';
}

int validate_username(const char *username) {
    if (!username || strlen(username) < 3 || strlen(username) >= MAX_USERNAME) {
        return 0;
    }
    
    for (int i = 0; username[i]; i++) {
        if (!isalnum(username[i]) && username[i] != '_') {
            return 0;
        }
    }
    return 1;
}

int validate_password(const char *password) {
    if (!password || strlen(password) < 8 || strlen(password) >= MAX_PASSWORD) {
        return 0;
    }
    
    int has_upper = 0, has_lower = 0, has_digit = 0;
    for (int i = 0; password[i]; i++) {
        if (isupper(password[i])) has_upper = 1;
        if (islower(password[i])) has_lower = 1;
        if (isdigit(password[i])) has_digit = 1;
    }
    
    return has_upper && has_lower && has_digit;
}

// Add error handling macro
#define HANDLE_ERROR(condition, message) \
    do { \
        if (condition) { \
            fprintf(stderr, "\033[31mError: %s (%s)\033[0m\n", message, strerror(errno)); \
            return 1; \
        } \
    } while(0)

int main() {
    struct sockaddr_in server_addr;
    char password[MAX_PASSWORD];

    char choice[10];
    printf(COLOR_SYSTEM "1. Login\n2. Register\nChoice: " COLOR_RESET);
    fgets(choice, sizeof(choice), stdin);

    printf(COLOR_SYSTEM "Username: " COLOR_RESET);
    fgets(username, MAX_USERNAME, stdin);
    username[strcspn(username, "\n")] = 0;

    printf(COLOR_SYSTEM "Password: " COLOR_RESET);
    fgets(password, MAX_PASSWORD, stdin);
    password[strcspn(password, "\n")] = 0;

    // Validate username and password
    if (!validate_username(username)) {
        printf(COLOR_ERROR "Invalid username. Use 3-31 characters, alphanumeric and underscore only." COLOR_RESET "\n");
        return 1;
    }

    if (!validate_password(password)) {
        printf(COLOR_ERROR "Invalid password. Must be 8+ characters with uppercase, lowercase and numbers." COLOR_RESET "\n");
        return 1;
    }

    printf(COLOR_SYSTEM "Connecting to %s:%d...\n" COLOR_RESET, SERVER_IP, PORT);
    
    client_socket = socket(AF_INET, SOCK_STREAM, 0);
    HANDLE_ERROR(client_socket == -1, "Failed to create socket");

    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(PORT);

    HANDLE_ERROR(inet_pton(AF_INET, SERVER_IP, &server_addr.sin_addr) <= 0, 
                "Invalid address");

    HANDLE_ERROR(connect(client_socket, (struct sockaddr *)&server_addr, 
                        sizeof(server_addr)) < 0, "Connection failed");

    // Send credentials with command
    char credentials[MAX_USERNAME + MAX_PASSWORD + 10];
    char hashed_password[65];
    hash_password(password, hashed_password, sizeof(hashed_password));

    snprintf(credentials, sizeof(credentials), "%s %s %s",
             choice[0] == '2' ? "REGISTER" : "LOGIN",
             username, hashed_password);

    send(client_socket, credentials, strlen(credentials), 0);

    // Handle response
    char response[MAX_MESSAGE];
    ssize_t response_len = recv(client_socket, response, sizeof(response) - 1, 0);
    if (response_len <= 0) {
        printf(COLOR_ERROR "Server disconnected" COLOR_RESET "\n");
        close(client_socket);
        return 1;
    }
    response[response_len] = '\0';
    rtrim_newlines(response);
    print_colored_message(response);

    // If registration, exit after showing response
    if (choice[0] == '2') {
        close(client_socket);
        return 0;
    }

    // Continue only if authentication was successful
    if (strncmp(response, "Authentication successful", 23) != 0) {
        close(client_socket);
        return 1;
    }

    printf(COLOR_SYSTEM "Successfully connected to server!" COLOR_RESET "\n\n");
    printf(COLOR_SYSTEM "Type /help to see available commands\n\n" COLOR_RESET);
    print_help();
    printf("\n");

    pthread_t receive_thread;
    pthread_create(&receive_thread, NULL, receive_messages, NULL);

    char message[MAX_MESSAGE];
    while (1) {
        show_prompt();
        if (!fgets(message, MAX_MESSAGE, stdin)) break;
        message[strcspn(message, "\n")] = 0;

        if (strcmp(message, "/help") == 0) {
            print_help();
            continue;
        }

        if (strcmp(message, "/quit") == 0) break;

        if (strncmp(message, "/join ", 6) == 0) {
            char *new_room = message + 6;
            pthread_mutex_lock(&room_mutex);
            strncpy(current_room, new_room, MAX_USERNAME - 1);
            current_room[MAX_USERNAME - 1] = '\0';
            pthread_mutex_unlock(&room_mutex);
        }

        // Send message
        if (strlen(message) > 0) {
            pthread_mutex_lock(&socket_mutex);
            
            // Only print local messages for commands starting with /
            if (message[0] == '/') {
                /* Local command feedback (colored) */
                if (strcmp(message, "/help") != 0) { /* /help already handled above */
                    printf(COLOR_COMMAND "Executing command: %s" COLOR_RESET "\n", message);
                }
            }
            // Send message to server
            ssize_t bytes_sent = send(client_socket, message, strlen(message), 0);
            if (bytes_sent <= 0) {
                printf(COLOR_ERROR "Failed to send message" COLOR_RESET "\n");
            }
            
            pthread_mutex_unlock(&socket_mutex);
        }
    }

    close(client_socket);
    return 0;
}