#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>

#define LOCAL_CONN_ID_LEN 16
#define MAX_DATAGRAM_SIZE 1350

typedef struct {
    int sock;
    struct sockaddr_storage peer_addr;
    socklen_t peer_addr_len;
    quiche_conn *conn;
} Client;

int main(int argc, char *argv[]) {
    quiche_config *config = quiche_config_new(QUICHE_PROTOCOL_VERSION);
    if (config == NULL) {
        fprintf(stderr, "failed to create config\n");
        return -1;
    }

    if (quiche_config_load_cert_chain_from_pem_file(config, "cert.crt") < 0) {
        fprintf(stderr, "failed to load certificate chain\n");
        return -1;
    }
    
    if (quiche_config_load_priv_key_from_pem_file(config, "cert.key") < 0) {
        fprintf(stderr, "failed to load private key\n");
        return -1;
    }
    
    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, 5000);
    quiche_config_set_max_recv_udp_payload_size(config, MAX_DATAGRAM_SIZE);
    quiche_config_set_max_send_udp_payload_size(config, MAX_DATAGRAM_SIZE);
    quiche_config_set_initial_max_data(config, 10000000);
    quiche_config_set_initial_max_stream_data_bidi_local(config, 1000000);
    quiche_config_set_initial_max_stream_data_bidi_remote(config, 1000000);
    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(8888);
    sa.sin_addr.s_addr = INADDR_ANY;

    int sock = socket(AF_INET, SOCK_DGRAM, 0);
    if (sock < 0) {
        perror("socket");
        return -1;
    }

    if (bind(sock, (struct sockaddr *)&sa, sizeof(sa)) < 0) {
        perror("bind");
        return -1;
    }

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

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

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

    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) {
            if (errno != EWOULDBLOCK && errno != EAGAIN) {
                perror("recvfrom");
                break;
            }
        } else {
            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 {
                         // Skip retry and accept connection directly
                         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);
                         printf("New 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;
                break;  // Exit the main loop and terminate
            }

            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)) {
                    uint8_t recv_buf[1024];
                    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) {
                         printf("Received %zd bytes on stream %lu: %.*s\n", recv_bytes, s, (int)recv_bytes, recv_buf);
                         char resp[1200];
                         snprintf(resp, sizeof(resp), "Server received: %.*s", (int)recv_bytes, recv_buf);
                         quiche_conn_stream_send(conn, s, (uint8_t*)resp, strlen(resp), true, &err_code);
                    }
                }
                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(1000); 
    }
    
    quiche_config_free(config);
    return 0;
}