#include <assert.h>
#include <openssl/ssl.h>
#include <openssl/rsa.h>
#include <openssl/x509.h>
#include <openssl/evp.h>
#include <openssl/err.h>
#include <string.h>

#ifdef _WIN32
#include <ws2tcpip.h>
#include <winsock2.h>
#include <synchapi.h>
#include <sysinfoapi.h>
#include <wincrypt.h>
#include <cryptuiapi.h>
#endif

#ifdef __linux__
#include <sys/time.h>
#include <fcntl.h>
#include <unistd.h>
#include <sys/select.h>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <sys/types.h>
#include <netdb.h>
typedef int SOCKET;
#endif

#ifdef _WIN32
void set_nonblocking(SOCKET sock) {
    u_long mode = 1;  // 1 to enable non-blocking socket
    int status = ioctlsocket(sock, FIONBIO, &mode);
    assert(status == 0);
}

void load_root_certs(SSL_CTX *ctx) {

    HCERTSTORE hStore = CertOpenSystemStoreW(0, L"ROOT");
    PCCERT_CONTEXT pContext = NULL;
    X509 *x509;
    X509_STORE *store = SSL_CTX_get_cert_store(ctx);

    if (!hStore) {
        fprintf(stderr, "Cannot open system certificate store.\n");
        exit(1);
    }

    while (pContext = CertEnumCertificatesInStore(hStore, pContext)) {
        const unsigned char *encoded_cert = pContext->pbCertEncoded; 
        x509 = d2i_X509(NULL, &encoded_cert, pContext->cbCertEncoded);
        if (x509) {
            int i = X509_STORE_add_cert(store, x509);
            if (i == 1) {
                fprintf(stderr, "Certificate added.\n");
            }
            X509_free(x509);
        }
    }

    CertCloseStore(hStore, 0);
}

void init_libraries() {
    WSADATA whocares;
    assert(WSAStartup(MAKEWORD(2,2), &whocares) == 0);
}

void print_error_and_exit() {
    int err = WSAGetLastError();
    fprintf(stderr, "err #%d\n", err);
    exit(1);
}

#endif

#ifdef __linux__
void set_nonblocking(int fd) {
    int flags = fcntl(fd, F_GETFL, 0);
    assert(flags != -1);
    int status = fcntl(fd, F_SETFL, flags | O_NONBLOCK);
    assert(status != -1);
}

void load_root_certs(SSL_CTX *ssl_ctx) {
    SSL_CTX_set_default_verify_paths(ssl_ctx);
}

void init_libraries() {
}

void print_error_and_exit() {
    fprintf(stderr, "error=%d\n", errno);
    exit(1);
}
#endif


void main(int argc, char **argv) {

    init_libraries();
    
    SSL_CTX *ssl_ctx = SSL_CTX_new(TLS_method());
    SSL_CTX_set_mode(ssl_ctx, SSL_MODE_ACCEPT_MOVING_WRITE_BUFFER);
    SSL_CTX_set_mode(ssl_ctx, SSL_MODE_ENABLE_PARTIAL_WRITE);
    SSL_CTX_set_verify(ssl_ctx, SSL_VERIFY_PEER, NULL);
    load_root_certs(ssl_ctx);

    SOCKET sock_fd = socket(AF_INET, SOCK_STREAM, 0);

    /* this is mit.edu */
    unsigned int ip0 = 104;
    unsigned int ip1 = 90;
    unsigned int ip2 = 66;
    unsigned int ip3 = 18;
    unsigned int ip = (ip0<<24)|(ip1<<16)|(ip2<<8)|ip3;
    
    struct sockaddr_in addr;
    addr.sin_family = AF_INET;
    addr.sin_port = htons(443);
    addr.sin_addr.s_addr = htonl(ip);

    int status = connect(sock_fd, (struct sockaddr *) &addr, sizeof(struct sockaddr_in));
    if (status != 0) print_error_and_exit();
    fprintf(stderr, "Connect successful.\n");

    set_nonblocking(sock_fd);

    SSL *ssl = SSL_new(ssl_ctx);
    SSL_set_fd(ssl, sock_fd);
    SSL_set1_host(ssl, "web.mit.edu");
    
    while (1) {
        int ret = SSL_connect(ssl);
        if (ret == 1) break;
        int err = SSL_get_error(ssl, ret);
        if (SSL_ERROR_WANT_READ == err || SSL_ERROR_WANT_WRITE == err) {
            /* sleep(1); */
        } else {
            ERR_print_errors_fp(stderr);
            exit(1);
        }
    }
    fprintf(stderr, "SSL connect done.\n");

    
    STACK_OF(X509)* certCollection = SSL_get_peer_cert_chain(ssl);
    for (size_t i = 0; i < sk_X509_num(certCollection); i++) {
        X509* cert = sk_X509_value(certCollection, i);
        X509_print_fp(stderr, cert);
        fprintf(stderr, "\n----\n");
    }

    const char *req =
	"HEAD / HTTP/1.1\r\nHost: mit.edu:https\r\nConnection: close\r\n\r\n";
    int reqlen = strlen(req);

    while (reqlen > 0) {
        size_t nwrote;
        int ret = SSL_write_ex(ssl, req, reqlen, &nwrote);
        if (ret == 1) {
            req += nwrote;
            reqlen -= nwrote;
        } else {
            int err = SSL_get_error(ssl, ret);
            if (SSL_ERROR_WANT_READ == err || SSL_ERROR_WANT_WRITE == err) {
                /* sleep(1); */
            } else {
                ERR_print_errors_fp(stderr);
                exit(1);
            }
        }
    }
    fprintf(stderr, "Header transmitted.\n");

    while (1) {
        char buf[256];
        size_t nread;
        int ret = SSL_read_ex(ssl, buf, 256, &nread);
        if (ret == 1) {
            fwrite(buf, 1, nread, stdout);
            fflush(stdout);
        } else {
            int err = SSL_get_error(ssl, ret);
            if (SSL_ERROR_WANT_READ == err || SSL_ERROR_WANT_WRITE == err) {
                /* sleep(1); */
            } else {
                ERR_print_errors_fp(stderr);
                exit(1);
            }
        }
    }
}