#include "utils.hpp"
#include "config.hpp"
#include "ORExpand.hpp"
#include "sort.hpp"
#include "storage.hpp"
#include "client.hpp"

#define PROFILE_STORAGE

StgClient *clients;

static struct {
    uint32_t max_users;
    uint32_t my_storage_node_id;
    // A local storage buffer, used when we need to do non-in-place
    // sorts of the messages that have arrived
    MsgBuffer stg_buf;
    // The destination vector for ORExpand
    std::vector<uint32_t> dest;
} storage_state;

// route_init will call this function; no one else should call it
// explicitly.  The parameter is the number of messages that can fit in
// the storage-side MsgBuffer.  Returns true on success, false on
// failure.
bool storage_init(uint32_t max_users, uint32_t msg_buf_size)
{
    storage_state.max_users = max_users;
    storage_state.stg_buf.alloc(msg_buf_size);
    storage_state.dest.resize(msg_buf_size);
    uint32_t my_storage_node_id = 0;
    uint32_t my_stg_pos = 0;
    for (nodenum_t i=0; i<g_teems_config.num_nodes; ++i) {
        if (g_teems_config.roles[i] & ROLE_STORAGE) {
            if (i == g_teems_config.my_node_num) {
                storage_state.my_storage_node_id = my_storage_node_id << DEST_UID_BITS;
                my_stg_pos = my_storage_node_id;
            } else {
                ++my_storage_node_id;
            }
        }
    }

    storage_generateClientKeys(max_users, my_stg_pos);

    return true;
}

// Handle the messages received by a storage node.  Pass a _locked_
// MsgBuffer.  This function will itself reset and unlock it when it's
// done with it.
void storage_received(MsgBuffer &storage_buf)
{
    uint16_t msg_size = g_teems_config.msg_size;
    nodenum_t my_node_num = g_teems_config.my_node_num;
    const uint8_t *msgs = storage_buf.buf;
    uint32_t num_msgs = storage_buf.inserted;
    uint32_t real = 0, padding = 0;
    uint32_t uid_mask = (1 << DEST_UID_BITS) - 1;
    uint32_t nid_mask = ~uid_mask;

#ifdef PROFILE_STORAGE
    unsigned long start_received = printf_with_rtclock("begin storage_received (%u)\n", storage_buf.inserted);
#endif

    // It's OK to test for errors in a way that's non-oblivous if
    // there's an error (but it should be oblivious if there are no
    // errors)
    for (uint32_t i=0; i<num_msgs; ++i) {
        uint32_t uid = *(const uint32_t*)(storage_buf.buf+(i*msg_size));
        bool ok = ((((uid & nid_mask) == storage_state.my_storage_node_id)
            & ((uid & uid_mask) < storage_state.max_users))
            | ((uid & uid_mask) == uid_mask));
        if (!ok) {
            printf("Received bad uid: %08x\n", uid);
            assert(ok);
        }
    }

    // Testing: report how many real and dummy messages arrived
    printf("Storage server received %u messages:\n", num_msgs);
    for (uint32_t i=0; i<num_msgs; ++i) {
        uint32_t dest_addr = *(const uint32_t*)msgs;
        nodenum_t dest_node =
            g_teems_config.storage_map[dest_addr >> DEST_UID_BITS];
        if (dest_node != my_node_num) {
            char hexbuf[2*msg_size + 1];
            for (uint32_t j=0;j<msg_size;++j) {
                snprintf(hexbuf+2*j, 3, "%02x", msgs[j]);
            }
            printf("Misrouted message: %s\n", hexbuf);
        } else if ((dest_addr & uid_mask) == uid_mask) {
            ++padding;
        } else {
            ++real;
        }
        msgs += msg_size;
    }
    printf("%u real, %u padding\n", real, padding);

    /*
    for (uint32_t i=0;i<num_msgs; ++i) {
        printf("%3d: %08x %08x\n", i,
        *(uint32_t*)(storage_buf.buf+(i*msg_size)),
        *(uint32_t*)(storage_buf.buf+(i*msg_size+4)));
    }
    */
    // Sort the received messages by userid into the
    // storage_state.stg_buf MsgBuffer.
#ifdef PROFILE_STORAGE
    unsigned long start_sort = printf_with_rtclock("begin oblivious sort (%u)\n", storage_buf.inserted);
#endif
    sort_mtobliv<UidKey>(g_teems_config.nthreads, storage_buf.buf,
        msg_size, storage_buf.inserted, storage_buf.bufsize,
        storage_state.stg_buf.buf);
#ifdef PROFILE_STORAGE
    printf_with_rtclock_diff(start_sort, "end oblivious sort (%u)\n", storage_buf.inserted);
#endif

    /*
    for (uint32_t i=0;i<num_msgs; ++i) {
        printf("%3d: %08x %08x\n", i,
        *(uint32_t*)(storage_state.stg_buf.buf+(i*msg_size)),
        *(uint32_t*)(storage_state.stg_buf.buf+(i*msg_size+4)));
    }
    */

#ifdef PROFILE_STORAGE
    unsigned long start_dest = printf_with_rtclock("begin setting dests (%u)\n", storage_state.stg_buf.bufsize);
#endif
    // Obliviously set the dest array
    uint32_t *dests = storage_state.dest.data();
    uint32_t stg_size = storage_state.stg_buf.bufsize;
    const uint8_t *buf = storage_state.stg_buf.buf;
    uint32_t m_priv_in = g_teems_config.m_priv_in;

    uint32_t uid = *(uint32_t*)(buf);
    uid &= uid_mask;
    // num_msgs is not a private value
    if (num_msgs > 0) {
        dests[0] = oselect_uint32_t(uid * m_priv_in, 0xffffffff,
            uid == uid_mask);
    }
    uint32_t prev_uid = uid;
    for (uint32_t i=1; i<num_msgs; ++i) {
        uid = *(uint32_t*)(buf + i*msg_size);
        uid &= uid_mask;
        uint32_t next = oselect_uint32_t(uid * m_priv_in, dests[i-1]+1,
            uid == prev_uid);
        dests[i] = oselect_uint32_t(next, 0xffffffff, uid == uid_mask);
        prev_uid = uid;
    }
    for (uint32_t i=num_msgs; i<stg_size; ++i) {
        dests[i] = 0xffffffff;
        *(uint32_t*)(buf + i*msg_size) = 0xffffffff;
    }
#ifdef PROFILE_STORAGE
    printf_with_rtclock_diff(start_dest, "end setting dests (%u)\n", stg_size);
#endif
    /*
    for (uint32_t i=0;i<stg_size; ++i) {
        printf("%3d: %08x %08x %u\n", i,
        *(uint32_t*)(storage_state.stg_buf.buf+(i*msg_size)),
        *(uint32_t*)(storage_state.stg_buf.buf+(i*msg_size+4)),
        dests[i]);
    }
    */

#ifdef PROFILE_STORAGE
    unsigned long start_expand = printf_with_rtclock("begin ORExpand (%u)\n", stg_size);
#endif
    ORExpand_parallel<OSWAP_16X>(storage_state.stg_buf.buf, dests,
        msg_size, stg_size, g_teems_config.nthreads);
#ifdef PROFILE_STORAGE
    printf_with_rtclock_diff(start_expand, "end ORExpand (%u)\n", stg_size);
#endif
    /*
    for (uint32_t i=0;i<stg_size; ++i) {
        printf("%3d: %08x %08x %u\n", i,
        *(uint32_t*)(storage_state.stg_buf.buf+(i*msg_size)),
        *(uint32_t*)(storage_state.stg_buf.buf+(i*msg_size+4)),
        dests[i]);
    }
    */

    // You can do more processing after these lines, as long as they
    // don't touch storage_buf.  They _can_ touch the backing buffer
    // storage_state.stg_buf.
    storage_buf.reset();
    pthread_mutex_unlock(&storage_buf.mutex);

    storage_state.stg_buf.reset();

#ifdef PROFILE_STORAGE
    printf_with_rtclock_diff(start_received, "end storage_received (%u)\n", storage_buf.inserted);
#endif
}

bool storage_generateClientKeys(uint32_t num_clients, uint32_t my_stg_no) {

    clients = new StgClient[num_clients];
    uint32_t num_stg_nodes = g_teems_config.num_storage_nodes;
    uint32_t c_simid = my_stg_no;

    //printf("In Ingestion::genCK, num_clients = %d\n", num_clients);
    for (uint32_t i=0; i<num_clients; i++) {
        const sgx_aes_gcm_128bit_key_t *pESK = &(g_teems_config.ESK);
        unsigned char zeroes[SGX_AESGCM_KEY_SIZE];
        unsigned char iv[SGX_AESGCM_IV_SIZE];
        sgx_aes_gcm_128bit_tag_t tag;
        memset(zeroes, 0, SGX_AESGCM_KEY_SIZE);
        memset(iv, 0, SGX_AESGCM_IV_SIZE);

        memcpy(iv, (uint8_t*) (&c_simid), sizeof(c_simid));
        memcpy(iv + sizeof(c_simid), "STG", sizeof("STG"));

        sgx_status_t ret = SGX_SUCCESS;
        ret = sgx_rijndael128GCM_encrypt(pESK, zeroes, SGX_AESGCM_KEY_SIZE,
            (uint8_t*) (clients[i].key), iv, SGX_AESGCM_IV_SIZE, NULL, 0, &tag);
        if(ret!=SGX_SUCCESS) {
            printf("stg_generateClientKeys FAIL\n");
            return false;
        }

        /*
        if(c_simid % 10 == 0) {
            printf("Storage: c_simid = %d, Key:", c_simid);
            for (int k = 0; k<SGX_AESGCM_KEY_SIZE; k++) {
                printf("%x", (clients[i].key)[k]);
            }
            printf("\n");
        }
        */
        c_simid+=num_stg_nodes;

    }

    return true;
}

bool authenticateClient()
{

}