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

#define PORT 7777
#define BUFFER_SIZE 8192
#define MAX_CLIENTS 100
#define HISTORY_SIZE 1000
#define USERNAME_MAX 50
#define AES_IV_SIZE 12
#define AES_TAG_SIZE 16

typedef struct {
    int socket;
    char username[USERNAME_MAX];
} client_t;

client_t *clients[MAX_CLIENTS];
pthread_mutex_t clients_mutex = PTHREAD_MUTEX_INITIALIZER;

typedef struct {
    int data_len;
    unsigned char *data;
} encrypted_message_t;

encrypted_message_t message_history[HISTORY_SIZE];
int history_count = 0;
int history_start = 0;
pthread_mutex_t history_mutex = PTHREAD_MUTEX_INITIALIZER;

void add_to_history(const unsigned char *data, int len) {
    pthread_mutex_lock(&history_mutex);
    int idx = (history_start + history_count) % HISTORY_SIZE;

    free(message_history[idx].data);
    message_history[idx].data = malloc(len);
    memcpy(message_history[idx].data, data, len);
    message_history[idx].data_len = len;

    if (history_count < HISTORY_SIZE) {
        history_count++;
    } else {
        history_start = (history_start + 1) % HISTORY_SIZE;
    }
    pthread_mutex_unlock(&history_mutex);
}

void send_history(int socket) {
    pthread_mutex_lock(&history_mutex);

    int net_count = htonl(history_count);
    write(socket, &net_count, sizeof(int));

    for (int i = 0; i < history_count; i++) {
        int idx = (history_start + i) % HISTORY_SIZE;
        int net_len = htonl(message_history[idx].data_len);
        write(socket, &net_len, sizeof(int));
        write(socket, message_history[idx].data, message_history[idx].data_len);
    }

    pthread_mutex_unlock(&history_mutex);
}

// Рассылка ВСЕМ клиентам включая отправителя (-1 = всем)
void broadcast_to_all(const unsigned char *data, int len) {
    pthread_mutex_lock(&clients_mutex);
    for (int i = 0; i < MAX_CLIENTS; i++) {
        if (clients[i]) {
            int net_len = htonl(len);
            send(clients[i]->socket, &net_len, sizeof(int), 0);
            send(clients[i]->socket, data, len, 0);
        }
    }
    pthread_mutex_unlock(&clients_mutex);
}

void *handle_client(void *arg) {
    client_t *cli = (client_t *)arg;

    int valread = read(cli->socket, cli->username, USERNAME_MAX);
    if (valread <= 0) goto cleanup;
    cli->username[valread] = '\0';
    cli->username[strcspn(cli->username, "\n")] = '\0';

    printf("User '%s' connected\n", cli->username);

    send_history(cli->socket);

    while (1) {
        int net_total_len;
        valread = read(cli->socket, &net_total_len, sizeof(int));
        if (valread <= 0) break;

        int total_len = ntohl(net_total_len);
        if (total_len <= 0 || total_len > BUFFER_SIZE) {
            printf("Invalid length from %s: %d\n", cli->username, total_len);
            break;
        }

        unsigned char *data = malloc(total_len);
        int total_read = 0;
        while (total_read < total_len) {
            valread = read(cli->socket, data + total_read, total_len - total_read);
            if (valread <= 0) {
                free(data);
                goto cleanup;
            }
            total_read += valread;
        }

        // Сохраняем
        add_to_history(data, total_len);

        // Рассылаем ВСЕМ включая отправителя
        broadcast_to_all(data, total_len);

        free(data);
    }

cleanup:
    printf("User '%s' disconnected\n", cli->username);
    close(cli->socket);

    pthread_mutex_lock(&clients_mutex);
    for (int i = 0; i < MAX_CLIENTS; i++) {
        if (clients[i] == cli) {
            clients[i] = NULL;
            break;
        }
    }
    pthread_mutex_unlock(&clients_mutex);

    free(cli);
    pthread_exit(NULL);
}

int main() {
    int server_fd, new_socket;
    struct sockaddr_in address;
    int opt = 1;
    int addrlen = sizeof(address);

    for (int i = 0; i < HISTORY_SIZE; i++) {
        message_history[i].data = NULL;
        message_history[i].data_len = 0;
    }

    if ((server_fd = socket(AF_INET, SOCK_STREAM, 0)) == 0) {
        perror("socket failed");
        exit(EXIT_FAILURE);
    }

    if (setsockopt(server_fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt))) {
        perror("setsockopt");
        exit(EXIT_FAILURE);
    }

    address.sin_family = AF_INET;
    address.sin_addr.s_addr = INADDR_ANY;
    address.sin_port = htons(PORT);

    if (bind(server_fd, (struct sockaddr *)&address, sizeof(address)) < 0) {
        perror("bind failed");
        exit(EXIT_FAILURE);
    }

    if (listen(server_fd, 3) < 0) {
        perror("listen");
        exit(EXIT_FAILURE);
    }

    printf("E2E Encrypted Server listening on port %d\n", PORT);

    while (1) {
        if ((new_socket = accept(server_fd, (struct sockaddr *)&address, (socklen_t*)&addrlen)) < 0) {
            perror("accept");
            continue;
        }

        client_t *cli = malloc(sizeof(client_t));
        cli->socket = new_socket;

        pthread_mutex_lock(&clients_mutex);
        int added = 0;
        for (int i = 0; i < MAX_CLIENTS; i++) {
            if (!clients[i]) {
                clients[i] = cli;
                added = 1;
                break;
            }
        }
        pthread_mutex_unlock(&clients_mutex);

        if (!added) {
            printf("Max clients reached\n");
            close(new_socket);
            free(cli);
            continue;
        }

        pthread_t tid;
        pthread_create(&tid, NULL, handle_client, cli);
        pthread_detach(tid);
    }

    return 0;
}