#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_mailboxes;

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

struct UserRange {
    uint32_t start, num;
    bool ret;
};

static void* generate_all_tokens_launch(void *voidargs)
{
    UserRange *args = (UserRange *)voidargs;

    uint32_t pt_tokens_size = (g_teems_config.m_priv_out * SGX_CMAC_MAC_SIZE);
    uint32_t enc_tokens_size = pt_tokens_size +
        SGX_AESGCM_IV_SIZE + SGX_AESGCM_MAC_SIZE;
    unsigned char token_body[pt_tokens_size];
    uint32_t user_start = args->start;
    uint32_t user_end = args->start + args->num;

    const sgx_aes_gcm_128bit_key_t *pTSK = &(g_teems_config.TSK);
    for(uint32_t lcid=user_start; lcid < user_end; 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
        // epoch_no used is for the next epoch
        unsigned long epoch_val = storage_epoch + 1;
        memcpy(tkn_iv_ptr, &epoch_val, sizeof(epoch_val));

        sgx_status_t ret = SGX_SUCCESS;
        unsigned char *ptr = tkn_ptr;
        unsigned char *tkn_body_ptr = token_body;
        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), &epoch_val, sizeof(epoch_val));

            ret = sgx_rijndael128_cmac_msg(pTSK, ptr, pt_tokens_size,
                (sgx_cmac_128bit_tag_t*) tkn_body_ptr);
            if(ret!=SGX_SUCCESS) {
                printf("generate_tokens: Creating token FAIL\n");
                args->ret = false;
                return NULL;
            }

            ptr+=SGX_CMAC_MAC_SIZE;
            tkn_body_ptr+=SGX_CMAC_MAC_SIZE;
        }


        /*
        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");
            args->ret = false;
            return NULL;
        }
        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");
        }
        */
    }

    args->ret = true;
    return NULL;
}

static bool launch_all_users(void *(*launch)(void *)) {
    threadid_t nthreads = g_teems_config.nthreads;

    // Special-case nthread=1 for efficiency
    if (nthreads <= 1) {
        UserRange args = { 0, storage_state.max_users };
        return launch(&args);
    }
    UserRange args[nthreads];
    uint32_t inc = storage_state.max_users / nthreads;
    uint32_t extra = storage_state.max_users % nthreads;
    uint32_t last = 0;
    for (threadid_t i=0; i<nthreads; ++i) {
        uint32_t num = inc + (i < extra);
        args[i] = { last, num };
        last += num;
    }

    // Launch all but the first section into other threads
    for (threadid_t i=1; i<nthreads; ++i) {
        threadpool_dispatch(g_thread_id+i, launch, args+i);
    }

    // Do the first section ourselves
    launch(args);

    // Join the threads
    for (threadid_t i=1; i<nthreads; ++i) {
        threadpool_join(g_thread_id+i, NULL);
    }

    bool ret = true;
    for (threadid_t i=0; i<nthreads; ++i) {
        ret &= args[i].ret;
    }
    return ret;
}

bool generate_all_tokens() {
    return launch_all_users(generate_all_tokens_launch);
}

/* processMsgs
   - Take all the messages in storage_state.stg_buf
   - Encrypt them all with their corresponding client key and IV and store into
      epoch_mailboxes
*/
static void *processMsgs_launch(void *voidargs) {
    UserRange *args = (UserRange *)voidargs;
    uint32_t user_start = args->start;
    uint32_t user_end = args->start + args->num;

    uint32_t mailbox_size = g_teems_config.m_priv_in * g_teems_config.msg_size;
    uint32_t enc_mailbox_size = mailbox_size + SGX_AESGCM_IV_SIZE + SGX_AESGCM_MAC_SIZE;
    unsigned char *epoch_buf_ptr = epoch_mailboxes +
        enc_mailbox_size * user_start;
    unsigned char *stg_buf_ptr = storage_state.stg_buf.buf +
        mailbox_size * user_start;
    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 + mailbox_size;

    for(uint32_t lcid = user_start; lcid < user_end; lcid++) {
        memcpy(epoch_buf_ptr, clients[lcid].iv, SGX_AESGCM_IV_SIZE);
        ret = sgx_rijndael128GCM_encrypt(&(clients[lcid].key), stg_buf_ptr, mailbox_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");
            args->ret = false;
            return NULL;
        }

        // 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+=mailbox_size;
        epoch_buf_ptr+=enc_mailbox_size;
        epoch_buf_ct_ptr+=enc_mailbox_size;
        epoch_buf_tag_ptr+=enc_mailbox_size;
    }

    args->ret = true;
    return NULL;
}

bool processMsgs() {
    return launch_all_users(processMsgs_launch);
}

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

    bool ret = generate_all_tokens();

    uint32_t num_expected_msgs = g_teems_config.m_priv_in * storage_state.max_users;
    processMsgs();

    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 *mailboxes,
    uint32_t mailboxes_size, unsigned char *tokens, uint32_t tokens_size)
{
    epoch_mailboxes = mailboxes;
    epoch_tokens = tokens;
}