#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdbool.h>
#include <sys/types.h>
#include <usrsctp.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <pthread.h>
#include <errno.h>

int done = 0;

const size_t send_buffer_size = 1024;

pthread_mutex_t lock;

int receive_cb(struct socket *sock, union sctp_sockstore addr, void *data, size_t datalen, struct sctp_rcvinfo info, int flags, void *ulp_info) {
    if (data == NULL) {
        // Closed connection
        done = 1;
        usrsctp_close(sock);

        // The mainthread might be stuck, unlock it so it can exit
        pthread_mutex_unlock(&lock);
    } else {
        if (write(fileno(stdout), data, datalen) < 0) {
            perror("write");
        }
    }

    return 1;
}

int send_cb(struct socket *sock, uint32_t sb_free, void *ulp_info) {
    // Callback called when the socket has space in its send cue. When this function is called is
    // rather bizarre:
    //
    // 1. Once called, this function is not called again until something has been added to the send
    //    queue;
    // 2. On the client side, this will be called BEFORE a connection has been established;
    // 3. If it blocks, at least on the client we will get stuck in the establishment process until
    //    it unblocks; and
    // 4. On the server this is never called after connection establishment until something has
    //    been added to the send queue.
    //
    // The easiest solution here is to do all the sending in the main thread and have this unlock
    // a mutex that the main thread will try to lock whenever it gets `EAGAIN` from usrsctp.
    
    pthread_mutex_unlock(&lock);

    return 1;
}

int main(int argc, char **argv) {
    if (argc < 2 || !(strncmp(argv[1], "-c", 3) == 0 || strncmp(argv[1], "-s", 3) == 0)) {
        fprintf(stderr, "Usage: %s [-c|-s] ...\n", argv[0]);
        return 1;
    }

    const bool is_server = strncmp(argv[1], "-s", 3) == 0;
    const bool is_client = !is_server;

    pthread_mutex_init(&lock, NULL);

    int local_udp_port;
    int remote_sctp_port;
    int remote_ip;

    int local_ip = 0; 
    int local_sctp_port = 0;
    int remote_udp_port = 0; // Only client

    struct socket *sock;

    // Server-only variables
    struct socket *listen_sock = NULL;

    if (is_server) {
        if (argc != 7) {
            fprintf(stderr, "Usage: %s -s <local IP> <local UDP port> <local SCTP port> <remote IP> <remote SCTP port>\n", argv[0]);
            return 1;
        }

        local_ip = inet_addr(argv[2]);
        local_udp_port = atoi(argv[3]);
        local_sctp_port = atoi(argv[4]);
        remote_ip = inet_addr(argv[5]);
        remote_sctp_port = atoi(argv[6]);
    } else {
        if (argc != 8) {
            fprintf(stderr, "Usage: %s -c <local IP> <local UDP port> <local SCTP port> <remote IP> <remote UDP port> <remote SCTP port>\n", argv[0]);
            return 1;
        }

        local_ip = inet_addr(argv[2]);
        local_udp_port = atoi(argv[3]);
        local_sctp_port = atoi(argv[4]);

        remote_ip = inet_addr(argv[5]);
        remote_udp_port = atoi(argv[6]);
        remote_sctp_port = atoi(argv[7]);
    }

    // This is the local UDP encapsulation port
    usrsctp_init(local_udp_port, NULL, NULL);

    // We use the callback API since usrsctp doesn't seem to support using poll/select
    sock = usrsctp_socket(
        PF_INET,
        SOCK_STREAM,
        IPPROTO_SCTP, 
        receive_cb,
        send_cb, 
        send_buffer_size,
        NULL
    );

    if (is_server) {
        // Server needs two sockets; this is the listening one
        listen_sock = sock;
    }

    // This is my address
    struct sockaddr_in my_addr;
    memset(&my_addr, 0, sizeof(my_addr));
    my_addr.sin_family = AF_INET;
    my_addr.sin_port = htons(local_sctp_port);
    my_addr.sin_addr.s_addr = local_ip;

    // Bind to the socket
    if (usrsctp_bind(sock, (struct sockaddr *) &my_addr, sizeof(my_addr)) < 0) {
        perror("bind");
        return 1;
    }

    if (is_client) {
        // Set the remote UDP port
        struct sctp_udpencaps encaps;
        memset(&encaps, 0, sizeof(encaps));
        encaps.sue_address.ss_family = AF_INET;
        encaps.sue_port = htons(remote_udp_port);

        usrsctp_setsockopt(sock, IPPROTO_SCTP, SCTP_REMOTE_UDP_ENCAPS_PORT, &encaps, sizeof(encaps));
    }

    // This is the remote address
    struct sockaddr_in addr;
    socklen_t addr_len = sizeof(addr);

    if (is_client) {
        // Set the remote address
        memset(&addr, 0, sizeof(addr));
        addr.sin_family = AF_INET;
        addr.sin_port = htons(remote_sctp_port);
        addr.sin_addr.s_addr = remote_ip;

        // Connect the socket
        fprintf(stderr, "Trying to connect...\n");
        if (usrsctp_connect(sock, (struct sockaddr *)&addr, sizeof(addr)) < 0) {
            perror("connect");
            return 1;
        }
        fprintf(stderr, "Connected!\n");
    } else {
        // Wait until we have a connection from the right IP/port
        fprintf(stderr, "Waiting for a connection...\n");
        while (1) {
            if (usrsctp_listen(listen_sock, 2) < 0) {
                perror("listen");
                return 1;
            }

            if ((sock = usrsctp_accept(listen_sock, (struct sockaddr *) &addr, &addr_len)) == NULL) {
                perror("accept");
                return 1;
            }

            if (remote_ip != 0 && (addr.sin_addr.s_addr != remote_ip || addr.sin_port != htons(remote_sctp_port))) {
                fprintf(stderr, "Bad IP/port tried: %d %s:%d\n", addr_len, inet_ntoa(addr.sin_addr), ntohs(addr.sin_port));
                usrsctp_close(sock);
            } else {
                // Valid connection
                break;
            }
        }

        fprintf(stderr, "Got a connection!\n");

        // Stop listening for new connections
        usrsctp_listen(listen_sock, 0);
    }

    // Mainloop: read from stdin and write to the SCTP socket
    char buff[send_buffer_size];
    int data_read;

    while (done == 0) {
        if ((data_read = read(fileno(stdin), &buff, sizeof(buff))) != -1) {
            // Send this to the remote
            if (data_read == 0) {
                // EOF
                done = 1;
            } else {
                while (usrsctp_sendv(sock, buff, data_read, NULL, 0, NULL, 0, 0, 0) < 0) {
                    // usrsctp has a rather small send buffer, so on e.g. a file transfer we will
                    // run into EAGAIN. Using usrsctp_set_non_blocking(sock, 0) doesn't seem to
                    // have any effect on this.
                    if (errno == EAGAIN) {
                        pthread_mutex_lock(&lock);
                        if (done == 1) {
                            break;
                        }
                    } else {
                        perror("sendv");
                        done = 1;
                        break;
                    }
                }
            }
        }
    }

    usrsctp_close(sock);

    if (is_server) {
        usrsctp_close(listen_sock);
    }

    usrsctp_finish();
}