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

#define PROFILE_STORAGE

StgClient *clients;
uint8_t *epoch_tokens;
uint8_t *epoch_msgbundles;

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;

bool storage_generateClientKeys(uint32_t num_clients, uint32_t my_stg_no) {
    uint16_t num_priv_channels = g_teems_config.m_priv_in;
    uint16_t msg_size = g_teems_config.msg_size;
    uint32_t pt_msgbundle_size = num_priv_channels * msg_size;

    clients = new StgClient[num_clients];

    for(uint32_t i =0; i < num_clients; i++) {
        uint32_t mid = storage_state.my_storage_node_id + i;
        clients[i].my_id = mid;
        clients[i].priv_friends = new clientid_t[g_teems_config.m_priv_out];
        // Initialize this client's private channel friends as themself
        for(int j =0; j <g_teems_config.m_priv_out; j++) {
            (clients[i].priv_friends)[j] = mid;
        }
    }

    uint32_t num_stg_nodes = g_teems_config.num_storage_nodes;
    uint32_t c_simid = my_stg_no;

    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 generate_all_tokens()
{
    uint32_t pt_tokens_size = (g_teems_config.m_priv_out * SGX_AESGCM_KEY_SIZE);
    uint32_t enc_tokens_size = (g_teems_config.m_priv_out * SGX_AESGCM_KEY_SIZE) +
        SGX_AESGCM_IV_SIZE + SGX_AESGCM_MAC_SIZE;
    unsigned char token_body[pt_tokens_size];

    const sgx_aes_gcm_128bit_key_t *pTSK = &(g_teems_config.TSK);
    for(uint32_t lcid=0; lcid<storage_state.max_users; lcid++) {

        unsigned char *tkn_iv_ptr = epoch_tokens + enc_tokens_size * lcid;
        unsigned char *tkn_ptr = tkn_iv_ptr + SGX_AESGCM_IV_SIZE;
        unsigned char *tkn_tag = tkn_ptr + pt_tokens_size;

        // We construct the plaintext [S|R|Epoch] underlying the token in
        // the correct location for this client in the epoch_tokens buffer
        // The tokens ( i.e., CMAC over S|R|Epoch) is stored in token_body
        // Then encrypt token_body with the client's storage key and overwrite
        // the correct location in epoch_tokens

        memset(token_body, 0, pt_tokens_size);
        memset(tkn_iv_ptr, 0, SGX_AESGCM_IV_SIZE);
        // IV = epoch_no, for encrypting the token bundle
        memcpy(tkn_iv_ptr, &storage_epoch, sizeof(storage_epoch));

        unsigned char *ptr = tkn_ptr;
        for(int i = 0; i<g_teems_config.m_priv_out; i++)
        {
            memcpy(ptr, (&(clients[lcid].my_id)), sizeof(clientid_t));
            memcpy(ptr + sizeof(clientid_t), (&(clients[lcid].priv_friends[i])), sizeof(clientid_t));
            memcpy(ptr + 2 * sizeof(clientid_t), &storage_epoch, sizeof(storage_epoch));
            ptr+=SGX_AESGCM_KEY_SIZE;
        }

        sgx_status_t ret = SGX_SUCCESS;
        ret = sgx_rijndael128_cmac_msg(pTSK, tkn_ptr, pt_tokens_size,
            (sgx_cmac_128bit_tag_t*) &token_body);
        if(ret!=SGX_SUCCESS) {
            printf("generate_tokens: Creating token FAIL\n");
            return false;
        }

        /*
        if(lcid == 0) {
            printf("Checking generated token_body:");
            for(uint32_t i = 0; i < pt_tokens_size; i++) {
                printf("%x", token_body[i]);
            }
            printf("\n");
        }
        */

        unsigned char *cl_iv = clients[lcid].iv;
        ret = (sgx_rijndael128GCM_encrypt(&(clients[lcid].key), token_body, pt_tokens_size,
            (uint8_t*) tkn_ptr, cl_iv, SGX_AESGCM_IV_SIZE, NULL, 0,
            (sgx_aes_gcm_128bit_tag_t*) tkn_tag));
        if(ret!=SGX_SUCCESS) {
            printf("generate_tokens: Encrypting token FAIL\n");
            return false;
        }
        memcpy(tkn_iv_ptr, cl_iv, SGX_AESGCM_IV_SIZE);
        // Update IV
        uint64_t *iv_ctr = (uint64_t*) cl_iv;
        (*iv_ctr)+=1;

        /*
        if(lcid == 0) {
            printf("Encrypted client token bundle:");
            for(uint32_t i = 0; i < enc_tokens_size; i++) {
                printf("%x", tkn_iv_ptr[i]);
            }
            printf("\n");
        }
        */
    }

    return true;
}


/* processMsgs
   - Take all the messages in storage_state.stg_buf
   - Encrypt them all with their corresponding client key and IV and store into
      epoch_msgbundles
   - ecall_send_msgbundle();
*/
bool processMsgs() {
    unsigned char *epoch_buf_ptr = epoch_msgbundles;
    unsigned char *stg_buf_ptr = storage_state.stg_buf.buf;
    uint32_t msg_bundle_size = g_teems_config.m_priv_in * g_teems_config.msg_size;
    uint32_t enc_msg_bundle_size = msg_bundle_size + SGX_AESGCM_IV_SIZE + SGX_AESGCM_MAC_SIZE;
    sgx_status_t ret = SGX_SUCCESS;
    unsigned char *epoch_buf_ct_ptr = epoch_buf_ptr + SGX_AESGCM_IV_SIZE;
    unsigned char *epoch_buf_tag_ptr = epoch_buf_ct_ptr + msg_bundle_size;

    for(uint32_t lcid = 0; lcid <storage_state.max_users; lcid++) {
        memcpy(epoch_buf_ptr, clients[lcid].iv, SGX_AESGCM_IV_SIZE);
        ret = sgx_rijndael128GCM_encrypt(&(clients[lcid].key), stg_buf_ptr, msg_bundle_size,
            (uint8_t*) epoch_buf_ct_ptr, epoch_buf_ptr, SGX_AESGCM_IV_SIZE, NULL, 0,
            (sgx_aes_gcm_128bit_tag_t*) epoch_buf_tag_ptr);
        if(ret!=SGX_SUCCESS) {
            printf("processMsgs: Encrypting msgs FAIL\n");
            return false;
        }

        // Update IV
        uint64_t *iv_ctr = (uint64_t*) clients[lcid].iv;
        (*iv_ctr)+=1;

        if(lcid==0) {
            printf("\n\nMessage for lcid 0, S, R = %d, %d\n\n\n", *((uint32_t*) stg_buf_ptr),
                *((uint32_t*) (stg_buf_ptr + 4)));
        }

        stg_buf_ptr+=msg_bundle_size;
        epoch_buf_ptr+=enc_msg_bundle_size;
        epoch_buf_ct_ptr+=enc_msg_bundle_size;
        epoch_buf_tag_ptr+=enc_msg_bundle_size;
    }

    return true;
}

bool generateDummies() {
    unsigned char *epoch_buf_ptr = epoch_msgbundles;
    uint32_t msg_bundle_size = g_teems_config.m_priv_in * g_teems_config.msg_size;
    uint32_t enc_msg_bundle_size = msg_bundle_size + SGX_AESGCM_IV_SIZE + SGX_AESGCM_MAC_SIZE;
    sgx_status_t ret = SGX_SUCCESS;
    unsigned char *epoch_buf_ct_ptr = epoch_buf_ptr + SGX_AESGCM_IV_SIZE;
    unsigned char *epoch_buf_tag_ptr = epoch_buf_ct_ptr + msg_bundle_size;
    unsigned char *zeroes = new unsigned char[msg_bundle_size];
    memset(zeroes, 0, msg_bundle_size);

    for(uint32_t lcid = 0; lcid <storage_state.max_users; lcid++) {
        memcpy(epoch_buf_ptr, clients[lcid].iv, SGX_AESGCM_IV_SIZE);
        ret = sgx_rijndael128GCM_encrypt(&(clients[lcid].key), zeroes, msg_bundle_size,
            (uint8_t*) epoch_buf_ct_ptr, epoch_buf_ptr, SGX_AESGCM_IV_SIZE, NULL, 0,
            (sgx_aes_gcm_128bit_tag_t*) epoch_buf_tag_ptr);
        if(ret!=SGX_SUCCESS) {
            printf("processMsgs: Encrypting msgs FAIL\n");
            return false;
        }

        // Update IV
        uint64_t *iv_ctr = (uint64_t*) clients[lcid].iv;
        (*iv_ctr)+=1;

        epoch_buf_ptr+=enc_msg_bundle_size;
        epoch_buf_ct_ptr+=enc_msg_bundle_size;
        epoch_buf_tag_ptr+=enc_msg_bundle_size;
    }

    delete zeroes;
    return true;
}

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

    printf("my_stg_pos = %d\n", my_stg_pos);
    storage_generateClientKeys(max_users, my_stg_pos);
    // sendClientTokens();

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

    generate_all_tokens();

    uint32_t num_expected_msgs = g_teems_config.m_priv_in * storage_state.max_users;
    if(num_msgs!= 0) {
        processMsgs();
    } else {
        generateDummies();
    }

    storage_epoch++;

    storage_state.stg_buf.reset();

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

bool ecall_storage_authenticate(clientid_t cid, unsigned char *auth_message)
{
    bool ret = false;
    uint32_t lcid = cid / g_teems_config.num_storage_nodes;
    const sgx_aes_gcm_128bit_key_t *ckey = &(clients[lcid].key);

    ret = authenticateClient(auth_message, ckey);

    if(!ret) {
        printf("Storage authentication FAIL\n");
    }
    return ret;
}


void ecall_supply_storage_buffers(unsigned char *msgbundles,
    uint32_t msgbundles_size, unsigned char *tokens, uint32_t tokens_size)
{
    epoch_msgbundles = msgbundles;
    epoch_tokens = tokens;
}