#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 <stdbool.h>

#define MAX_DATAGRAM_SIZE 1350
#define TOTAL_MB 100
#define NUM_STREAMS 5
#define MB_PER_STREAM (TOTAL_MB / NUM_STREAMS)

typedef struct {
    uint64_t stream_id;
    long long bytes_sent;
    long long bytes_total;
    bool finished;
} StreamState;

int main(int argc, char *argv[]) {
    quiche_config *config = quiche_config_new(QUICHE_PROTOCOL_VERSION);
    if (config == NULL) return -1;
    
    quiche_config_verify_peer(config, false); 
    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);
    quiche_config_set_initial_max_data(config, 1024 * 1024 * 500);
    quiche_config_set_initial_max_stream_data_bidi_local(config, 1024 * 1024 * 100);
    quiche_config_set_initial_max_stream_data_bidi_remote(config, 1024 * 1024 * 100);
    quiche_config_set_initial_max_streams_bidi(config, 100);
    quiche_config_set_cc_algorithm(config, QUICHE_CC_RENO);

    int sock = socket(AF_INET, SOCK_DGRAM, 0);
    if (sock < 0) return -1;
    
    struct sockaddr_in peer_addr;
    memset(&peer_addr, 0, sizeof(peer_addr));
    peer_addr.sin_family = AF_INET;
    peer_addr.sin_port = htons(8889);
    inet_pton(AF_INET, "100.115.45.1", &peer_addr.sin_addr);
    
    if (connect(sock, (struct sockaddr *)&peer_addr, sizeof(peer_addr)) < 0) return -1;
    
    struct sockaddr_in local_addr;
    socklen_t local_addr_len = sizeof(local_addr);
    if (getsockname(sock, (struct sockaddr *)&local_addr, &local_addr_len) < 0) return -1;

    int flags = fcntl(sock, F_GETFL, 0);
    fcntl(sock, F_SETFL, flags | O_NONBLOCK);
    
    uint8_t scid[QUICHE_MAX_CONN_ID_LEN];
    int rng = open("/dev/urandom", O_RDONLY);
    read(rng, scid, sizeof(scid));
    close(rng);
    
    quiche_conn *conn = quiche_connect("100.115.45.1", (const uint8_t *)scid, sizeof(scid), (struct sockaddr *)&local_addr, local_addr_len, (struct sockaddr *)&peer_addr, sizeof(peer_addr), config);
    if (conn == NULL) return -1;
    
    printf("Connecting to QUIC Multi-Stream Server...\n");
    printf("Sending %d streams, %d MB each (Total %d MB)...\n", NUM_STREAMS, MB_PER_STREAM, TOTAL_MB);
    
    uint8_t buf[65535];
    uint8_t out[MAX_DATAGRAM_SIZE];
    uint8_t payload[4096];
    memset(payload, 'C', sizeof(payload));

    // Initialize stream states
    StreamState streams[NUM_STREAMS];
    for (int i = 0; i < NUM_STREAMS; i++) {
        streams[i].stream_id = 4 * i + 4; // 4, 8, 12, 16, 20... (Client Bidi) or simple increment if library handles it? 
        // Note: Client initiated bidi streams usually start at 0, then 4, 8... 
        // but let's stick to explicit IDs or check quiche docs.
        // Quiche expects us to use IDs. 0, 4, 8, 12, 16 are valid client bidi.
        streams[i].stream_id = i * 4; 
        streams[i].bytes_sent = 0;
        streams[i].bytes_total = (long long)MB_PER_STREAM * 1024 * 1024;
        streams[i].finished = false;
    }

    bool all_finished = false;

    while (1) {
        ssize_t read_len = recv(sock, buf, sizeof(buf), MSG_DONTWAIT);
        if (read_len > 0) {
             quiche_conn_recv(conn, buf, read_len, &(quiche_recv_info){
                 .to = (struct sockaddr *)&local_addr,
                 .to_len = local_addr_len,
                 .from = (struct sockaddr *)&peer_addr,
                 .from_len = sizeof(peer_addr),
             });
        }

        if (quiche_conn_is_closed(conn)) {
            printf("Connection closed.\n");
            break;
        }

        if (quiche_conn_is_established(conn)) {
             all_finished = true;
             for (int i = 0; i < NUM_STREAMS; i++) {
                 if (!streams[i].finished) {
                     all_finished = false;
                     // Try to send on this stream
                     while (streams[i].bytes_sent < streams[i].bytes_total) {
                         uint64_t err_code = 0;
                         // Determine payload size
                         ssize_t sent = quiche_conn_stream_send(conn, streams[i].stream_id, payload, sizeof(payload), false, &err_code);
                         if (sent > 0) {
                             streams[i].bytes_sent += sent;
                             if (streams[i].bytes_sent >= streams[i].bytes_total) {
                                 // Send FIN
                                 quiche_conn_stream_send(conn, streams[i].stream_id, NULL, 0, true, &err_code);
                                 streams[i].finished = true;
                                 printf("Stream %ld finished.\n", streams[i].stream_id);
                             }
                         } else {
                             // E.g. Done (congestion) or Stream Limit
                             break; 
                         }
                     }
                 }
             }

             if (all_finished) {
                 // Wait a bit to ensure ACKs or just exit?
                 // Ideally wait for close, but let's just loop a bit or wait for idle.
                 // Actually the server will likely not close connection, we can just idle.
             }
        }

        bool has_outgoing = false;
        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;
            send(sock, out, written, 0);
            has_outgoing = true;
        }

        quiche_conn_on_timeout(conn);
        if (!has_outgoing && !all_finished) usleep(100);
        if (all_finished && !has_outgoing) {
             // Maybe wait for connection close or timeout
             usleep(100000); // Wait 100ms
             // quiche_conn_close(conn, true, 0, "Done", 4);
             // break;
             // Let's keep it alive for a moment for server to ack then exit
             static int linger = 0;
             if (linger++ > 20) {
                 printf("All streams finished. Closing.\n");
                 uint8_t reason[] = "done";
                 quiche_conn_close(conn, true, 0, reason, sizeof(reason));
                 break; 
             }
        }
    }

    quiche_conn_free(conn);
    quiche_config_free(config);
    return 0;
}