#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 "appconfig.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)
{
    sgx_ec256_public_t pubkey;
    sgx_sealed_data_t *sealedprivkey =
        (sgx_sealed_data_t *)malloc(SEALED_PRIVKEY_SIZE);

    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, SEALED_PRIVKEY_SIZE, 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)
{
    sgx_ec256_public_t pubkey;
    sgx_sealed_data_t *sealedprivkey =
        (sgx_sealed_data_t *)malloc(SEALED_PRIVKEY_SIZE);

    if (fread(sealedprivkey, SEALED_PRIVKEY_SIZE, 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 -k sealedprivkeyfile -n myname "
        "[-t nthreads] [-d epoch wait time (in s)] [-e number of epochs] "
        "[-w number of Waksman Networks to precompute before starting "
        "any epochs] [command] < config.json\n",
        argv0);
    exit(1);
}

int main(int argc, char **argv)
{
    // Unbuffer stdout
    setbuf(stdout, NULL);

    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]);
        }

        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);
    }

    const char *sealedprivkeyfile = NULL;
    const char *myname = NULL;
    uint16_t nthreads = 1;

    const char *progname = argv[0];
    ++argv;
    // Parse options
    while (*argv && (*argv)[0] == '-') {
        if (!strcmp(*argv, "--")) {
            argv += 1;
            break;
        } else if (!strcmp(*argv, "-k")) {
            if (argv[1] == NULL) {
                usage(progname);
            }
            sealedprivkeyfile = argv[1];
            argv += 2;
        } else if (!strcmp(*argv, "-n")) {
            if (argv[1] == NULL) {
                usage(progname);
            }
            myname = argv[1];
            argv += 2;
        } else if (!strcmp(*argv, "-t")) {
            if (argv[1] == NULL) {
                usage(progname);
            }
            nthreads = uint16_t(atoi(argv[1]));
            argv += 2;
        } else if (!strcmp(*argv, "-d")) {
            if (argv[1] == NULL) {
                usage(progname);
            }
            epoch_wait_time = int(atoi(argv[1]));
            argv += 2;
        } else if (!strcmp(*argv, "-e")) {
            if (argv[1] == NULL) {
                usage(progname);
            }
            num_epochs = int(atoi(argv[1]));
            argv += 2;
        } else if (!strcmp(*argv, "-w")) {
            if (argv[1] == NULL) {
                usage(progname);
            }
            num_WN_to_precompute = int(atoi(argv[1]));
            argv += 2;
        } else {
            usage(progname);
        }
    }

    if (sealedprivkeyfile == NULL || myname == NULL) {
        usage(progname);
    }

    // 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, nthreads)) {
        exit(1);
    }

    boost::asio::io_context io_context;

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

    // Queue up the actual work
    boost::asio::post(io_context, [&]{

        // Start enclave-to-enclave communications
        ecall_comms_start([&]{
            boost::asio::post(io_context, [&]{
                // This runs when we have completed our handshakes with
                // all other nodes
                printf("Starting\n");
                start(netio, argv);
            });
        });
        printf("Reading\n");
        for (nodenum_t node_num = 0; node_num < netio.num_nodes; ++node_num) {
            if (node_num == netio.me) continue;
            NodeIO &node = netio.node(node_num);
            node.recv_commands(
                // error_cb
                [node_num, &config](boost::system::error_code ec) {
                    printf("Error %s from %d %s\n",
                        ec.message().c_str(), node_num,
                        config.nodes[node_num].name.c_str());
                },
                // epoch_cb
                [](uint32_t epoch) {
                    printf("Epoch %u\n", epoch);
                });
        }
    });

    // Start threads to handle the async io handlers
    boost::thread t0([&]{io_context.run();});
    boost::thread t1([&]{io_context.run();});
    io_context.run();
    t0.join();
    t1.join();

    // All done
    ecall_close();
    g_netio = NULL;
    sgx_destroy_enclave(global_eid);

    return 0;
}