#include <unistd.h>
#include "common.h"

// include/ path is here for a reason - to make sure we compile against boringssl (temporary solution)
#include <include/openssl/ssl.h>
#include <pthread.h>

// Buffer used for read/write tests
unsigned char rw_buf[BUFFER_SIZE];
static const char* DefaultCurves = "CECPQ2b:CECPQ2:X25519";
static const uint16_t TLS_PROT_VERSION = TLS1_3_VERSION;

SSL_CTX *setup_client_ctx(const char* curves)
{
    SSL_CTX* ctx = NULL;

    ctx = SSL_CTX_new(TLS_method());
    assert(ctx != NULL);
    SSL_CTX_set_verify(ctx, SSL_VERIFY_PEER, NULL);
    if( SSL_CTX_load_verify_locations(ctx, CACERT, NULL) != 1) {
        ERR("Error loading CA DIR");
    }

    if (SSL_CTX_set1_curves_list(ctx, curves)!=1) {
        ERR("Can't set SIDH group");
    }

    if (SSL_CTX_set_min_proto_version(ctx, TLS_PROT_VERSION) != 1 ||
        SSL_CTX_set_max_proto_version(ctx, TLS_PROT_VERSION) != 1) {
        ERR("Enforcing protocol to TLSv1.2");
    }

    SSL_CTX_enable_pq_experiment_signal(ctx);
    return ctx;
}

int do_client_loop(SSL* ssl)
{
    int err, nwritten;
    char buf[80];

    for(;;)
    {
        if( !fgets(buf, sizeof(buf), stdin) )
            break;
        for(nwritten = 0; nwritten < strlen(buf); nwritten += err)
        {
            err = SSL_write(ssl, buf + nwritten, strlen(buf) - nwritten);
            if(err<=0)
                return 0;
        }
    }
    return 1;
}

void test_Handshake(const char* IP, const char* curves, size_t handshake_nb) {
    SSL* ssl;
    SSL_CTX* ctx;
    int err;

    init();
    ctx = setup_client_ctx(curves);

    DBG("Trying to connect");
    int fd = connect_once(IP);

    for (size_t i=0; i<handshake_nb; i++) {

        DBG("SSL ctx setup");
        if (!(ssl = SSL_new(ctx)))  {
            ERR("Error creating an SSL context");
        }
        SSL_set_fd(ssl, fd);

        DBG("SSL handshake");
        err = SSL_connect(ssl);
        if (err<=0) {
            ERR("Error connecting SSL err=%d", err);
        }
        //printf("Signal seen: %X\n", SSL_pq_experiment_signal_seen(ssl));
        assert(!SSL_session_reused(ssl));
        assert(SSL_shutdown(ssl) == 0);
        assert(SSL_shutdown(ssl) == 1);
        SSL_free(ssl);
    }

    SSL_CTX_free(ctx);
    close(fd);
}

void test_Read(const char* IP, const char* curves) {
    int err, nread=0;
    SSL* ssl;
    SSL_CTX* ctx;

    init();
    ctx = setup_client_ctx(curves);

    DBG("Trying to connect");
    int fd = connect_once(IP);

    DBG("SSL ctx setup");
    if (!(ssl = SSL_new(ctx)))  {
        ERR("Error creating an SSL context");
    }
    SSL_set_fd(ssl, fd);

    DBG("SSL handshake");
    err = SSL_connect(ssl);
    if (err<=0) {
        ERR("Error connecting SSL err=%d", err);
    }

    DBG("SSL writing");
    if (SSL_write(ssl, (unsigned char*)READ_REQUEST, sizeof(READ_REQUEST) ) <= 0) {
        ERR("SSL_write");
    }

    DBG("Continue reading");
    do {
#ifdef DEBUG
        fwrite(rw_buf,1, nread, stdout);
#endif
        for(nread=0; nread<sizeof(rw_buf); nread+=err) {
            err = SSL_read(ssl, rw_buf+nread, sizeof(rw_buf) - nread);
            if(err <= 0) {
                break;
            }
        }
    } while(err>0);

    if (0==(SSL_get_shutdown(ssl) & SSL_RECEIVED_SHUTDOWN)) {
        ERR("SSL_Read");
    }

    DBG("Shutdown");
    assert(SSL_shutdown(ssl) == 1);

    SSL_free(ssl);
    SSL_CTX_free(ctx);
    close(fd);
}

void test_Write(const char* IP, const char* curves) {
    SSL* ssl;
    SSL_CTX* ctx;
    int err,nread=0;

    init();
    ctx = setup_client_ctx(curves);

    DBG("Trying to connect");
    int fd = connect_once(IP);

    DBG("SSL ctx setup");
    if (!(ssl = SSL_new(ctx)))  {
        ERR("Error creating an SSL context");
    }
    SSL_set_fd(ssl, fd);

    DBG("SSL handshake");
    err = SSL_connect(ssl);
    if (err<=0) {
        ERR("Error connecting SSL err=%d", err);
    }

    DBG("Writing");

    fill_buffer_from_file(rw_buf, sizeof(rw_buf));
    for (size_t i=0; i<RW_ITERATIONS; i++) {
        for(nread=0; nread<sizeof(rw_buf); nread+=err) {
            err = SSL_write(ssl, rw_buf+nread, sizeof(rw_buf) - nread);
            if(err <= 0) {
                ERR("SSL_write");
            }
        }
    }

    DBG("Shutdown");
    assert(SSL_shutdown(ssl) == 0);
    assert(SSL_shutdown(ssl) == 1);

    SSL_free(ssl);
    SSL_CTX_free(ctx);
    close(fd);
}

int main(int argc, char* argv[]) {

    if (argc < 3) {
        goto usage;
    }

    const char *curves = NULL;
    if (argc==4) {
        curves = argv[3];
    } else {
        curves = DefaultCurves;
    }

    if (!strncmp("test_Handshake", argv[2],  strlen("test_Handshake"))) {
        test_Handshake(argv[1], curves, HANDHAKE_REPS);
    } else if (!strncmp("test_Read", argv[2], strlen("test_Read"))) {
        test_Read(argv[1], curves);
    } else if (!strncmp("test_Write", argv[2], strlen("test_Write"))) {
        test_Write(argv[1], curves);
    } else {
        printf("Unknown test");
        goto usage;
    }

exit:
    cleanup();
    return 0;

usage:
    ERR("\n\nUsage: %s host_ip test_name [groups]\nOptions for 'test_name':\n\t"
        "test_Handshake\n\ttest_Write\n\ttest_Read\n", argv[0]);
    goto exit;
}