#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <time.h>
#include <sys/epoll.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <arpa/inet.h>
#include <fcntl.h>
#include <errno.h>
#include <math.h>
#include <stdint.h>

// --- Test Configuration ---
#define HOST "127.0.0.1"
#define PORT 8080
#define NUM_SUBSCRIBERS 1000
#define NUM_PUBLISHERS 10
#define TOTAL_CLIENTS (NUM_SUBSCRIBERS + NUM_PUBLISHERS)
#define TEST_DURATION_S 15
#define MESSAGES_PER_SECOND_PER_PUBLISHER 100

// --- Internal Configuration ---
#define MAX_EVENTS TOTAL_CLIENTS
#define RW_BUFFER_SIZE 8192
#define MAX_LATENCIES 20000000 // Pre-allocate for ~1.3M messages/sec

// --- WebSocket Constants ---
#define WEBSOCKET_KEY_MAGIC "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"

// --- Helper Enums ---
typedef enum {
    CLIENT_SUBSCRIBER,
    CLIENT_PUBLISHER
} ClientType;

typedef enum {
    STATE_CONNECTING,
    STATE_HANDSHAKE_SEND,
    STATE_HANDSHAKE_RECV,
    STATE_SUBSCRIBING,
    STATE_RUNNING,
    STATE_CLOSED
} ClientState;

// --- Client State Structure ---
typedef struct {
    int fd;
    ClientType type;
    ClientState state;
    char read_buf[RW_BUFFER_SIZE];
    size_t read_len;
    char write_buf[RW_BUFFER_SIZE];
    size_t write_len;
    size_t write_pos;
    double next_send_time;
} Client;

// --- Global State & Metrics ---
double* latencies;
size_t latencies_count = 0;
uint64_t messages_sent = 0;
uint64_t messages_received = 0;
int subscriber_setup_count = 0;
int all_subscribed = 0;
int epoll_fd;

const char* CHANNELS[] = {"news", "sports", "tech", "finance", "weather"};
const int NUM_CHANNELS = sizeof(CHANNELS) / sizeof(CHANNELS[0]);


// =============================================================================
// START: Clean SHA-1 and Base64 Implementations
// =============================================================================

// --- SHA-1 Implementation ---
typedef struct {
    uint32_t state[5];
    uint32_t count[2];
    unsigned char buffer[64];
} SHA1_CTX;

#define SHA1_ROTLEFT(n,c) (((n) << (c)) | ((n) >> (32 - (c))))

void SHA1_Transform(uint32_t state[5], const unsigned char buffer[64]) {
    uint32_t a, b, c, d, e;
    uint32_t block[16];
    memcpy(block, buffer, 64);
    for(int i = 0; i < 16; i++) {
        uint8_t *p = (uint8_t*)&block[i];
        block[i] = (p[0] << 24) | (p[1] << 16) | (p[2] << 8) | p[3];
    }
    
    a = state[0]; b = state[1]; c = state[2]; d = state[3]; e = state[4];
    
    uint32_t W[80];
    for (int t = 0; t < 80; t++) {
        if (t < 16) {
            W[t] = block[t];
        } else {
            W[t] = SHA1_ROTLEFT(W[t - 3] ^ W[t - 8] ^ W[t - 14] ^ W[t - 16], 1);
        }
        uint32_t temp = SHA1_ROTLEFT(a, 5) + e + W[t];
        if (t < 20) temp += ((b & c) | (~b & d)) + 0x5A827999;
        else if (t < 40) temp += (b ^ c ^ d) + 0x6ED9EBA1;
        else if (t < 60) temp += ((b & c) | (b & d) | (c & d)) + 0x8F1BBCDC;
        else temp += (b ^ c ^ d) + 0xCA62C1D6;
        e = d; d = c; c = SHA1_ROTLEFT(b, 30); b = a; a = temp;
    }
    state[0] += a; state[1] += b; state[2] += c; state[3] += d; state[4] += e;
}

void SHA1_Init(SHA1_CTX* context) {
    context->state[0] = 0x67452301;
    context->state[1] = 0xEFCDAB89;
    context->state[2] = 0x98BADCFE;
    context->state[3] = 0x10325476;
    context->state[4] = 0xC3D2E1F0;
    context->count[0] = context->count[1] = 0;
}

void SHA1_Update(SHA1_CTX* context, const unsigned char* data, uint32_t len) {
    uint32_t i, j;
    j = context->count[0];
    if ((context->count[0] += len << 3) < j) context->count[1]++;
    context->count[1] += (len >> 29);
    j = (j >> 3) & 63;
    if ((j + len) > 63) {
        memcpy(&context->buffer[j], data, (i = 64 - j));
        SHA1_Transform(context->state, context->buffer);
        for (; i + 63 < len; i += 64) {
            SHA1_Transform(context->state, &data[i]);
        }
        j = 0;
    } else {
        i = 0;
    }
    memcpy(&context->buffer[j], &data[i], len - i);
}

void SHA1_Final(unsigned char digest[20], SHA1_CTX* context) {
    uint32_t i;
    unsigned char finalcount[8];
    for (i = 0; i < 8; i++) {
        finalcount[i] = (unsigned char)((context->count[(i >= 4 ? 0 : 1)] >> ((3 - (i & 3)) * 8)) & 255);
    }
    SHA1_Update(context, (unsigned char*)"\x80", 1);
    while ((context->count[0] & 504) != 448) {
        SHA1_Update(context, (unsigned char*)"\0", 1);
    }
    SHA1_Update(context, finalcount, 8);
    for (i = 0; i < 20; i++) {
        digest[i] = (unsigned char)((context->state[i >> 2] >> ((3 - (i & 3)) * 8)) & 255);
    }
}

// --- Base64 Implementation ---
char* base64_encode(const unsigned char *data, size_t input_length) {
    const char b64_table[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
    size_t output_length = 4 * ((input_length + 2) / 3);
    char *encoded_data = malloc(output_length + 1);
    if (encoded_data == NULL) return NULL;

    for (size_t i = 0, j = 0; i < input_length;) {
        uint32_t octet_a = i < input_length ? data[i++] : 0;
        uint32_t octet_b = i < input_length ? data[i++] : 0;
        uint32_t octet_c = i < input_length ? data[i++] : 0;
        uint32_t triple = (octet_a << 16) + (octet_b << 8) + octet_c;

        encoded_data[j++] = b64_table[(triple >> 18) & 0x3F];
        encoded_data[j++] = b64_table[(triple >> 12) & 0x3F];
        encoded_data[j++] = b64_table[(triple >> 6) & 0x3F];
        encoded_data[j++] = b64_table[triple & 0x3F];
    }

    for (size_t i = 0; i < (3 - input_length % 3) % 3; i++) {
        encoded_data[output_length - 1 - i] = '=';
    }
    encoded_data[output_length] = '\0';
    return encoded_data;
}

// =============================================================================
// END: Clean SHA-1 and Base64 Implementations
// =============================================================================


// --- Utility Functions ---
double get_time_double() {
    struct timespec ts;
    clock_gettime(CLOCK_REALTIME, &ts);
    return ts.tv_sec + ts.tv_nsec / 1e9;
}

void epoll_ctl_mod(int fd, uint32_t events, void* ptr) {
    struct epoll_event ev;
    ev.events = events;
    ev.data.ptr = ptr;
    epoll_ctl(epoll_fd, EPOLL_CTL_MOD, fd, &ev);
}

void close_client(Client* client) {
    if (client->state != STATE_CLOSED) {
        epoll_ctl(epoll_fd, EPOLL_CTL_DEL, client->fd, NULL);
        close(client->fd);
        client->state = STATE_CLOSED;
    }
}

// --- WebSocket Core Functions ---
size_t create_ws_frame(const char* payload, size_t payload_len, char* out_buffer) {
    size_t frame_len = 2 + payload_len + 4; // Header + Mask + Payload
    if (payload_len > 125) frame_len += 2; // For 16-bit length
    
    out_buffer[0] = 0x81; // FIN + Text Frame
    if (payload_len <= 125) {
        out_buffer[1] = 0x80 | payload_len;
    } else {
        out_buffer[1] = 0x80 | 126;
        *(uint16_t*)(out_buffer + 2) = htons(payload_len);
    }

    size_t header_len = (payload_len <= 125) ? 2 : 4;
    uint32_t mask = rand();
    *(uint32_t*)(out_buffer + header_len) = mask;
    
    uint8_t* mask_bytes = (uint8_t*)&mask;
    for(size_t i = 0; i < payload_len; ++i) {
        out_buffer[header_len + 4 + i] = payload[i] ^ mask_bytes[i % 4];
    }
    return header_len + 4 + payload_len;
}

void send_handshake(Client* client) {
    unsigned char key_bytes[16];
    for (int i = 0; i < 16; i++) key_bytes[i] = rand() % 256;
    char* b64_key = base64_encode(key_bytes, 16);

    client->write_len = snprintf(client->write_buf, RW_BUFFER_SIZE,
        "GET / HTTP/1.1\r\n"
        "Host: %s:%d\r\n"
        "Upgrade: websocket\r\n"
        "Connection: Upgrade\r\n"
        "Sec-WebSocket-Key: %s\r\n"
        "Sec-WebSocket-Version: 13\r\n\r\n",
        HOST, PORT, b64_key);
    
    client->state = STATE_HANDSHAKE_SEND;
    epoll_ctl_mod(client->fd, EPOLLIN | EPOLLOUT | EPOLLET, client);
    free(b64_key);
}

// --- Event Handlers ---
void handle_write(Client* client) {
    ssize_t sent = send(client->fd, client->write_buf + client->write_pos, client->write_len - client->write_pos, 0);
    if (sent < 0) {
        if (errno != EAGAIN && errno != EWOULDBLOCK) close_client(client);
        return;
    }
    client->write_pos += sent;

    if (client->write_pos >= client->write_len) {
        client->write_pos = 0;
        client->write_len = 0;
        epoll_ctl_mod(client->fd, EPOLLIN | EPOLLET, client); // Done writing, wait for reads

        if (client->state == STATE_HANDSHAKE_SEND) {
            client->state = STATE_HANDSHAKE_RECV;
        } else if (client->state == STATE_SUBSCRIBING) {
            client->state = STATE_RUNNING;
            subscriber_setup_count++;
            if (subscriber_setup_count == NUM_SUBSCRIBERS) {
                 printf("✅ All subscribers are connected and subscribed. Starting publishers...\n");
                 all_subscribed = 1;
            }
        }
    }
}

void handle_read(Client* client) {
    ssize_t n = read(client->fd, client->read_buf + client->read_len, RW_BUFFER_SIZE - client->read_len);
    if (n <= 0) {
        if (n < 0 && errno != EAGAIN && errno != EWOULDBLOCK) perror("read error");
        close_client(client);
        return;
    }
    client->read_len += n;

    if (client->state == STATE_HANDSHAKE_RECV) {
        if (strstr(client->read_buf, "\r\n\r\n")) {
            if (strstr(client->read_buf, " 101 ") == NULL) {
                fprintf(stderr, "Handshake failed for fd %d\n", client->fd);
                close_client(client);
                return;
            }
            client->read_len = 0; // Handshake complete, clear buffer

            if (client->type == CLIENT_SUBSCRIBER) {
                const char* channel = CHANNELS[rand() % NUM_CHANNELS];
                char sub_msg[128];
                int msg_len = snprintf(sub_msg, sizeof(sub_msg), "sub %s", channel);
                client->write_len = create_ws_frame(sub_msg, msg_len, client->write_buf);
                client->state = STATE_SUBSCRIBING;
                epoll_ctl_mod(client->fd, EPOLLIN | EPOLLOUT | EPOLLET, client);
            } else { // Publisher
                client->state = STATE_RUNNING;
            }
        }
    } else if (client->state == STATE_RUNNING && client->type == CLIENT_SUBSCRIBER) {
        // Simple WebSocket frame parsing for this specific test case
        while (client->read_len >= 2) {
            uint64_t payload_len = client->read_buf[1] & 0x7F;
            size_t header_len = 2;
            if (payload_len == 126) {
                if (client->read_len < 4) break;
                payload_len = ntohs(*(uint16_t*)(client->read_buf + 2));
                header_len = 4;
            } else if (payload_len == 127) {
                // Not expected for this test, would require 64-bit length handling
                close_client(client);
                break;
            }

            if (client->read_len >= header_len + payload_len) {
                char* payload = client->read_buf + header_len;
                double sent_time = atof(payload);
                if (sent_time > 0) {
                    double latency = get_time_double() - sent_time;
                    if (latencies_count < MAX_LATENCIES) {
                        latencies[latencies_count++] = latency;
                    }
                    messages_received++;
                }
                
                size_t frame_size = header_len + payload_len;
                memmove(client->read_buf, client->read_buf + frame_size, client->read_len - frame_size);
                client->read_len -= frame_size;
            } else {
                break; // Incomplete frame
            }
        }
    }
}

// --- Statistics ---
int compare_doubles(const void* a, const void* b) {
    double da = *(const double*)a;
    double db = *(const double*)b;
    if (da < db) return -1;
    if (da > db) return 1;
    return 0;
}

void print_report() {
    printf("\n"
           "================================================================================\n");
    printf("%s\n", "                              PERFORMANCE REPORT                              ");
    printf("================================================================================\n");

    if (latencies_count == 0) {
        printf("No messages were received. Cannot generate a report. Is the server running?\n");
        return;
    }

    uint64_t message_loss = (messages_sent > messages_received) ? (messages_sent - messages_received) : 0;
    double loss_rate = (messages_sent > 0) ? ((double)message_loss / messages_sent * 100.0) : 0;
    double throughput = (double)messages_received / TEST_DURATION_S;

    printf("Test Duration:       %d seconds\n", TEST_DURATION_S);
    printf("Total Messages Sent: %lu\n", messages_sent);
    printf("Total Messages Rcvd: %lu\n", messages_received);
    printf("Message Loss:        %lu (%.2f%%)\n", message_loss, loss_rate);
    printf("Actual Throughput:   %.2f msg/sec\n", throughput);
    printf("--------------------------------------------------------------------------------\n");

    qsort(latencies, latencies_count, sizeof(double), compare_doubles);
    
    double sum = 0;
    for (size_t i = 0; i < latencies_count; ++i) sum += latencies[i];
    
    printf("Latency Statistics (ms):\n");
    printf("  Average:           %.4f ms\n", (sum / latencies_count) * 1000.0);
    printf("  Min:               %.4f ms\n", latencies[0] * 1000.0);
    printf("  Max:               %.4f ms\n", latencies[latencies_count - 1] * 1000.0);
    printf("  Median (p50):      %.4f ms\n", latencies[(size_t)(latencies_count * 0.50)] * 1000.0);
    printf("  95th Percentile:   %.4f ms\n", latencies[(size_t)(latencies_count * 0.95)] * 1000.0);
    printf("  99th Percentile:   %.4f ms\n", latencies[(size_t)(latencies_count * 0.99)] * 1000.0);
    printf("================================================================================\n");
}

// --- Main Function ---
int main() {
    srand(time(NULL));
    latencies = malloc(sizeof(double) * MAX_LATENCIES);
    if (!latencies) {
        perror("malloc latencies");
        return 1;
    }

    printf("Starting WebSocket Pub/Sub Load Test...\n");
    printf("Simulating %d subscribers and %d publishers.\n", NUM_SUBSCRIBERS, NUM_PUBLISHERS);
    printf("Publishing at ~%d msg/sec for %d seconds.\n", NUM_PUBLISHERS * MESSAGES_PER_SECOND_PER_PUBLISHER, TEST_DURATION_S);
    printf("--------------------------------------------------------------------------------\n");
    
    epoll_fd = epoll_create1(0);
    if (epoll_fd == -1) {
        perror("epoll_create1");
        free(latencies);
        return 1;
    }

    // *** FIX: Allocate clients on the heap, not the stack ***
    Client* clients = malloc(sizeof(Client) * TOTAL_CLIENTS);
    if (!clients) {
        perror("malloc clients");
        free(latencies);
        close(epoll_fd);
        return 1;
    }

    struct sockaddr_in server_addr;
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(PORT);
    inet_pton(AF_INET, HOST, &server_addr.sin_addr);

    for (int i = 0; i < TOTAL_CLIENTS; ++i) {
        clients[i].fd = socket(AF_INET, SOCK_STREAM, 0);
        fcntl(clients[i].fd, F_SETFL, O_NONBLOCK);
        
        clients[i].type = (i < NUM_SUBSCRIBERS) ? CLIENT_SUBSCRIBER : CLIENT_PUBLISHER;
        clients[i].state = STATE_CONNECTING;
        clients[i].read_len = 0;
        clients[i].write_len = 0;
        clients[i].write_pos = 0;
        clients[i].next_send_time = 0;

        connect(clients[i].fd, (struct sockaddr*)&server_addr, sizeof(server_addr));

        struct epoll_event ev;
        ev.events = EPOLLIN | EPOLLOUT | EPOLLET;
        ev.data.ptr = &clients[i];
        if (epoll_ctl(epoll_fd, EPOLL_CTL_ADD, clients[i].fd, &ev) == -1) {
            perror("epoll_ctl: add");
            free(latencies);
            free(clients);
            close(epoll_fd);
            return 1;
        }
    }

    struct epoll_event events[MAX_EVENTS];
    double start_time = get_time_double();
    double end_time = start_time + TEST_DURATION_S;

    while (get_time_double() < end_time) {
        int n = epoll_wait(epoll_fd, events, MAX_EVENTS, 10); // 10ms timeout
        double now = get_time_double();

        for (int i = 0; i < n; ++i) {
            Client* client = (Client*)events[i].data.ptr;
            if (events[i].events & (EPOLLERR | EPOLLHUP)) {
                close_client(client);
                continue;
            }

            if (client->state == STATE_CONNECTING && (events[i].events & EPOLLOUT)) {
                int result;
                socklen_t result_len = sizeof(result);
                getsockopt(client->fd, SOL_SOCKET, SO_ERROR, &result, &result_len);
                if (result == 0) {
                    send_handshake(client);
                } else {
                    close_client(client);
                }
            } else {
                if (events[i].events & EPOLLIN) handle_read(client);
                if (events[i].events & EPOLLOUT) handle_write(client);
            }
        }
        
        // Publisher logic
        if (all_subscribed) {
            for (int i = NUM_SUBSCRIBERS; i < TOTAL_CLIENTS; ++i) {
                Client* client = &clients[i];
                if (client->state == STATE_RUNNING && client->write_len == 0 && now >= client->next_send_time) {
                    const char* channel = CHANNELS[rand() % NUM_CHANNELS];
                    char message[256];
                    int msg_len = snprintf(message, sizeof(message), "%.6f:Hello from publisher %d on channel %s", now, i - NUM_SUBSCRIBERS, channel);
                    
                    char pub_msg[384];
                    int pub_msg_len = snprintf(pub_msg, sizeof(pub_msg), "pub %s %s", channel, message);
                    
                    client->write_len = create_ws_frame(pub_msg, pub_msg_len, client->write_buf);
                    messages_sent++;

                    client->next_send_time = now + (1.0 / MESSAGES_PER_SECOND_PER_PUBLISHER);
                    epoll_ctl_mod(client->fd, EPOLLOUT | EPOLLIN | EPOLLET, client);
                }
            }
        }
    }

    printf("\nTest duration finished. Shutting down clients...\n");
    for (int i = 0; i < TOTAL_CLIENTS; ++i) {
        close_client(&clients[i]);
    }

    print_report();
    
    free(clients);
    free(latencies);
    close(epoll_fd);

    return 0;
}