#include <iostream>
#include <cstdio>
#include <cstring>

#include <boost/asio.hpp>
#include <boost/thread.hpp>

#include "sgx_urts.h"
#include "sgx_tcrypto.h"
#include "sgx_tseal.h"
#include "Untrusted.hpp"
#include "config.hpp"
#include "net.hpp"
#include "start.hpp"

static bool hexdump(FILE *outf, const char *label, void *p, size_t len)
{
    unsigned char *pc = (unsigned char *)p;
    if (label) {
        if (fprintf(outf, "%s: ", label) < 0) {
            return false;
        }
    }
    for (size_t i=0; i<len; ++i) {
        if (fprintf(outf, "%02x", pc[i]) < 0) {
            return false;
        }
    }
    if (fprintf(outf, "\n") < 0) {
        return false;
    }
    return true;
}

// Generate a new public/private keypair, and save the sealed private
// key and the public key to the given filenames (overwriting if they're
// already there).  Return true on success, false on failure.
static bool genkey(const char *sealedprivkeyfile, const char *pubkeyfile)
{
    size_t sealedprivsize =
        sizeof(sgx_sealed_data_t) + sizeof(sgx_ec256_private_t) + 19;

    sgx_ec256_public_t pubkey;
    sgx_sealed_data_t *sealedprivkey =
        (sgx_sealed_data_t *)malloc(sealedprivsize);

    ecall_identity_key_new(&pubkey, sealedprivkey);

    printf("Generating and saving public key and sealed private key\n");

    FILE *sprivf = fopen(sealedprivkeyfile, "wb");
    if (!sprivf) {
        free(sealedprivkey);
        perror("Sealed privkey write failed");
        return false;
    }
    if (fwrite(sealedprivkey, sealedprivsize, 1, sprivf) != 1) {
        fclose(sprivf);
        free(sealedprivkey);
        perror("Sealed privkey write failed");
        return false;
    }
    free(sealedprivkey);
    if (fclose(sprivf)) {
        perror("Sealed privkey write failed");
        return false;
    }
    FILE *pubf = fopen(pubkeyfile, "wb");
    if (!pubf) {
        perror("Pubkey write failed");
        return false;
    }
    if (!hexdump(pubf, NULL, &pubkey, sizeof(pubkey))) {
        fclose(pubf);
        perror("Pubkey write failed");
        return false;
    }
    if (fclose(pubf)) {
        perror("Pubkey write failed");
        return false;
    }
    hexdump(stdout, "Pubkey", &pubkey, sizeof(pubkey));

    return true;
}

// Try to load a sealed private key from sprivf.  If successful, save
// the public key to pubkeyfile (if non-NULL), overwriting if already
// present.  Return true on success, false on failure.
static bool loadkey(FILE *sprivf, const char *pubkeyfile)
{
    size_t sealedprivsize =
        sizeof(sgx_sealed_data_t) + sizeof(sgx_ec256_private_t) + 19;

    sgx_ec256_public_t pubkey;
    sgx_sealed_data_t *sealedprivkey =
        (sgx_sealed_data_t *)malloc(sealedprivsize);

    if (fread(sealedprivkey, sealedprivsize, 1, sprivf) != 1) {
        fprintf(stderr, "Could not read sealed privkey file\n");
        return false;
    }

    bool res = ecall_identity_key_load(&pubkey, sealedprivkey);
    free(sealedprivkey);

    if (!res) {
        fprintf(stderr, "Key load failed\n");
        return false;
    }

    printf("Loaded sealed private key\n");
    if (pubkeyfile) {
        FILE *pubf = fopen(pubkeyfile, "wb");
        if (!pubf) {
            perror("Pubkey write failed");
            return false;
        }
        if (!hexdump(pubf, NULL, &pubkey, sizeof(pubkey))) {
            fclose(pubf);
            perror("Pubkey write failed");
            return false;
        }
        if (fclose(pubf)) {
            perror("Pubkey write failed");
            return false;
        }
    }
    hexdump(stdout, "Pubkey", &pubkey, sizeof(pubkey));

    return true;
}

static void usage(const char *argv0)
{
    fprintf(stderr, "Usage: %s --gen sealedprivkeyfile pubkeyfile\n",
        argv0);
    fprintf(stderr, "or     %s sealedprivkeyfile myname [args] < config.json\n",
        argv0);
}

int main(int argc, char **argv)
{
    if (initialize_enclave() < 0) {
        return -1;
    }

    // --gen sealedprivkeyfile pubkeyfile
    // If sealedprivkeyfile exists and is valid, regenerate the
    // pubkeyfile (overwriting if there was already one there).
    // Otherwise, generate and write a new sealedprivkeyfile and
    // pubkeyfile (again overwriting if needed).
    if (argc > 1 && !strcmp(argv[1], "--gen")) {
        if (argc != 4) {
            usage(argv[0]);
            exit(1);
        }

        const char *sealedprivkeyfile = argv[2];
        const char *pubkeyfile = argv[3];

        FILE *sprivf = fopen(sealedprivkeyfile, "rb");
        if (sprivf && loadkey(sprivf, pubkeyfile)) {
            // We successfully used an existing sealedprivkeyfile
            fclose(sprivf);
            return 0;
        }
        if (sprivf) {
            fclose(sprivf);
        }
        // Generate a new keypair
        if (genkey(sealedprivkeyfile, pubkeyfile)) {
            return 0;
        }
        // Something went wrong
        exit(1);
    }

    if (argc < 3) {
        usage(argv[0]);
        exit(1);
    }

    const char *sealedprivkeyfile = argv[1];
    std::string myname(argv[2]);

    // Read the config.json from the first line of stdin.  We have to do
    // this before outputting anything to avoid potential deadlock with
    // the launch program.
    std::string configstr;
    std::getline(std::cin, configstr);

    // Load the sealed private key
    FILE *sprivf = fopen(sealedprivkeyfile, "rb");
    if (!sprivf) {
        perror("Cannot read sealed private key file");
        exit(1);
    }
    if (!loadkey(sprivf, NULL)) {
        fprintf(stderr, "Could not load sealed private key\n");
        exit(1);
    }
    fclose(sprivf);

    Config config;

    if (!config_parse(config, configstr, myname)) {
        exit(1);
    }

    boost::asio::io_context io_context;

    // The NetIO will keep a (const) reference to the config
    NetIO netio(io_context, config);

    // Queue up the actual work
    boost::asio::post(io_context, [&]{
        start(netio, argc, argv);
    });

    // Start another thread; one will perform the work and the other
    // will execute the async_write handlers
    boost::thread t([&]{io_context.run();});
    io_context.run();
    t.join();

    // All done
    sgx_destroy_enclave(global_eid);

    return 0;
}