#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <quiche.h>
#include <time.h>

#define MAX_DATAGRAM_SIZE 1350
#define LOCAL_CONN_ID_LEN 16

// Simple structure to hold connection state
typedef struct {
    int sock;
    struct sockaddr_storage peer_addr;
    socklen_t peer_addr_len;
    quiche_conn *conn;
    long long total_bytes;
    struct timespec start_time;
    int timer_started;
} Client;

int main(int argc, char *argv[]) {
    // 1. Configuration
    quiche_config *config = quiche_config_new(QUICHE_PROTOCOL_VERSION);
    if (config == NULL) return -1;

    quiche_config_load_cert_chain_from_pem_file(config, "cert.crt");
    quiche_config_load_priv_key_from_pem_file(config, "cert.key");
    
    quiche_config_set_application_protos(config, (uint8_t *) "\x0ahq-interop\x05hq-29\x05hq-28\x05hq-27\x08http/0.9", 38);
    quiche_config_set_max_idle_timeout(config, 10000);
    quiche_config_set_max_recv_udp_payload_size(config, MAX_DATAGRAM_SIZE);
    quiche_config_set_max_send_udp_payload_size(config, MAX_DATAGRAM_SIZE);
    // Increase limits for performance
    quiche_config_set_initial_max_data(config, 1024 * 1024 * 200);
    quiche_config_set_initial_max_stream_data_bidi_local(config, 1024 * 1024 * 200);
    quiche_config_set_initial_max_stream_data_bidi_remote(config, 1024 * 1024 * 200);
    quiche_config_set_initial_max_streams_bidi(config, 100);
    quiche_config_set_cc_algorithm(config, QUICHE_CC_RENO);

    struct sockaddr_in sa;
    memset(&sa, 0, sizeof(sa));
    sa.sin_family = AF_INET;
    sa.sin_port = htons(8889);
    sa.sin_addr.s_addr = INADDR_ANY;

    int sock = socket(AF_INET, SOCK_DGRAM, 0);
    if (sock < 0) return -1;
    if (bind(sock, (struct sockaddr *)&sa, sizeof(sa)) < 0) return -1;

    int flags = fcntl(sock, F_GETFL, 0);
    fcntl(sock, F_SETFL, flags | O_NONBLOCK);

    printf("QUIC Performance Server listening on port 8889\n");

    Client *client = NULL; 
    uint8_t buf[65535];
    uint8_t out[MAX_DATAGRAM_SIZE];
    bool done_printing = false;

    while (1) {
        struct sockaddr_storage peer_addr;
        socklen_t peer_addr_len = sizeof(peer_addr);
        ssize_t read_len = recvfrom(sock, buf, sizeof(buf), 0, (struct sockaddr *)&peer_addr, &peer_addr_len);

        if (read_len > 0) {
            uint8_t type;
            uint32_t version;
            uint8_t scid[QUICHE_MAX_CONN_ID_LEN];
            size_t scid_len = sizeof(scid);
            uint8_t dcid[QUICHE_MAX_CONN_ID_LEN];
            size_t dcid_len = sizeof(dcid);
            uint8_t token[256];
            size_t token_len = sizeof(token);

            int rc = quiche_header_info(buf, read_len, LOCAL_CONN_ID_LEN, &version, &type, scid, &scid_len, dcid, &dcid_len, token, &token_len);

            if (rc >= 0) {
                if (client == NULL) {
                    if (!quiche_version_is_supported(version)) {
                         ssize_t written = quiche_negotiate_version(scid, scid_len, dcid, dcid_len, out, sizeof(out));
                         if (written > 0) sendto(sock, out, written, 0, (struct sockaddr *)&peer_addr, peer_addr_len);
                    } else if (token_len == 0) {
                        uint8_t new_scid[QUICHE_MAX_CONN_ID_LEN];
                        int rng = open("/dev/urandom", O_RDONLY);
                        if (rng >= 0) {
                            read(rng, new_scid, sizeof(new_scid));
                            close(rng);
                        }
                        ssize_t written = quiche_retry(scid, scid_len, dcid, dcid_len, new_scid, sizeof(new_scid), token, token_len, version, out, sizeof(out));
                        if (written > 0) sendto(sock, out, written, 0, (struct sockaddr *)&peer_addr, peer_addr_len);
                    } else {
                         client = malloc(sizeof(Client));
                         client->sock = sock;
                         client->peer_addr = peer_addr;
                         client->peer_addr_len = peer_addr_len;
                         
                         uint8_t server_scid[QUICHE_MAX_CONN_ID_LEN];
                         int rng = open("/dev/urandom", O_RDONLY);
                         if (rng >= 0) {
                            read(rng, server_scid, sizeof(server_scid));
                            close(rng);
                         }

                         client->conn = quiche_accept(server_scid, sizeof(server_scid), dcid, dcid_len, (struct sockaddr *)&sa, sizeof(sa), (struct sockaddr *)&peer_addr, peer_addr_len, config);
                         client->total_bytes = 0;
                         client->timer_started = 0;
                         printf("New performance connection accepted.\n");
                    }
                }
                
                if (client != NULL) {
                    quiche_conn_recv(client->conn, buf, read_len, &(quiche_recv_info){
                        .to = (struct sockaddr *)&sa,
                        .to_len = sizeof(sa),
                        .from = (struct sockaddr *)&peer_addr,
                        .from_len = peer_addr_len,
                    });
                }
            }
        }

        if (client != NULL) {
            quiche_conn *conn = client->conn;
            
            if (quiche_conn_is_closed(conn)) {
                printf("Connection closed.\n");
                quiche_conn_free(conn);
                free(client);
                client = NULL;
                continue;
            }

            if (quiche_conn_is_established(conn)) {
                uint64_t s = 0;
                quiche_stream_iter *readable = quiche_conn_readable(conn);
                while (quiche_stream_iter_next(readable, &s)) {
                    if (!client->timer_started) {
                        clock_gettime(CLOCK_MONOTONIC, &client->start_time);
                        client->timer_started = 1;
                    }
                    
                    uint8_t recv_buf[65535];
                    bool fin = false;
                    uint64_t err_code = 0;
                    ssize_t recv_bytes = quiche_conn_stream_recv(conn, s, recv_buf, sizeof(recv_buf), &fin, &err_code);
                    if (recv_bytes > 0) {
                        client->total_bytes += recv_bytes;
                    }
                    if (fin) {
                        struct timespec end;
                        clock_gettime(CLOCK_MONOTONIC, &end);
                        double time_taken = (end.tv_sec - client->start_time.tv_sec) + (end.tv_nsec - client->start_time.tv_nsec) / 1e9;
                        double mb = client->total_bytes / (1024.0 * 1024.0);
                        double throughput = mb / time_taken;
                        
                        if (!done_printing) {
                             printf("Received %.2f MB in %.2f seconds.\n", mb, time_taken);
                             printf("Throughput: %.2f MB/s\n", throughput);
                             done_printing = true;
                        }
                    }
                }
                quiche_stream_iter_free(readable);
            }
            
            while (1) {
                quiche_send_info send_info;
                ssize_t written = quiche_conn_send(conn, out, sizeof(out), &send_info);
                if (written == QUICHE_ERR_DONE) break;
                if (written < 0) break;
                sendto(sock, out, written, 0, (struct sockaddr *)&send_info.to, send_info.to_len);
            }
            
            quiche_conn_on_timeout(conn);
        }
        usleep(100); 
    }
    
    quiche_config_free(config);
    return 0;
}