#include <pthread.h>
#include "Enclave_t.h"
#include "utils.hpp"
#include "config.hpp"
#include "route.hpp"
#include "ingest.hpp"

Ingestion g_ing;

bool ecall_ingest_msgbundle(clientid_t cid, unsigned char *msgbundle,
    uint32_t num_msgs)
{
    bool ret;
    ret = g_ing.processMsgBundle(cid, msgbundle, num_msgs);
    return ret;
}

bool ecall_ingestion_authenticate(clientid_t cid, unsigned char *auth_message)
{
    bool ret;
    ret = g_ing.authenticate(cid, auth_message);
    return ret;
}

void Ingestion::initialize(uint32_t num, uint32_t start,
        sgx_aes_gcm_128bit_key_t &ESK) {
    cnum = num;
    cstart = start;
    clients = new IngClient[cnum];

    generateClientKeys(ESK);

    if(g_teems_config.token_channel) {
        max_buffer_size = g_teems_config.m_token_out * cnum;
    } else {
        max_buffer_size = g_teems_config.m_id_out * cnum;
    }
    buffer = &(route_state.ingbuf);
}

bool Ingestion::authenticate(clientid_t cid, unsigned char *auth_message)
{
    uint32_t num_ing_nodes = g_teems_config.num_ingestion_nodes;
    uint32_t lcid = cid / num_ing_nodes;
    const sgx_aes_gcm_128bit_key_t *ckey = &(clients[lcid].key);
    return authenticateClient(auth_message, ckey);
}

bool Ingestion::processMsgBundle(clientid_t cid, unsigned char *msgbundle,
    uint32_t num_msgs) {

    // Fetch corresponding client key
    uint32_t num_ing_nodes = g_teems_config.num_ingestion_nodes;
    uint32_t num_stg_nodes = g_teems_config.num_storage_nodes;
    clientid_t lcid = cid / num_ing_nodes;
    clientid_t global_client_id =
        ((cid % num_stg_nodes) << DEST_UID_BITS) + (cid / num_stg_nodes);
    sgx_aes_gcm_128bit_key_t *ckey = &(clients[lcid].key);
    unsigned char *iv = msgbundle;
    msgbundle += SGX_AESGCM_IV_SIZE;

    uint16_t msg_size = g_teems_config.msg_size;
    uint32_t msgbundle_size;
    if(g_teems_config.token_channel) {
        msgbundle_size = num_msgs * (msg_size + TOKEN_SIZE);
    } else {
        msgbundle_size = num_msgs * msg_size;
    }
    unsigned char *dec_msgbundle = (unsigned char *) malloc (msgbundle_size);
    sgx_aes_gcm_128bit_tag_t *tag =
        (sgx_aes_gcm_128bit_tag_t*) (msgbundle + msgbundle_size);

    sgx_status_t ret = sgx_rijndael128GCM_decrypt(ckey, msgbundle,
        msgbundle_size, dec_msgbundle, iv, SGX_AESGCM_IV_SIZE,
        NULL, 0, tag);

    if(ret!=SGX_SUCCESS) {
        printf("Ingestion::processMsgBundle FAIL\n");
        printf("Error code: %d", (uint32_t) ret);
        free(dec_msgbundle);
        return false;
    }

    // Force the sender ids to be set properly.

    // In reality, we would also set the priorities here, but for the
    // benchmark, we'll let the clients choose them arbitrarily.

    unsigned char *sid_ptr = dec_msgbundle + sizeof(clientid_t);
    if (!g_teems_config.token_channel) {
        // Leave room for the priority
        sid_ptr += sizeof(uint32_t);
    }
    for (uint32_t k =0; k < num_msgs; k++) {
        *(clientid_t*)sid_ptr = global_client_id;
        sid_ptr += msg_size;
    }

    // Verify the tokens from end of the msgbundle

    if (g_teems_config.token_channel) {
        unsigned char token_body[TOKEN_SIZE];
        const sgx_aes_gcm_128bit_key_t *pTSK = &(g_teems_config.TSK);
        unsigned char *dm_ptr = dec_msgbundle;
        unsigned char *tkn_ptr = dm_ptr + (msg_size * num_msgs);
        sgx_cmac_128bit_tag_t token;

        for (uint32_t k = 0; k < num_msgs; k++) {
            ret = SGX_SUCCESS;
            memset(token_body, 0, TOKEN_SIZE);
            memcpy(token_body, dm_ptr, sizeof(clientid_t) * 2);
            memcpy(token_body + sizeof(clientid_t) * 2,
                (uint8_t*) &ingestion_epoch, sizeof(ingestion_epoch));
            ret = sgx_rijndael128_cmac_msg(pTSK, token_body, TOKEN_SIZE,
                (sgx_cmac_128bit_tag_t*) &token);
            if(ret!=SGX_SUCCESS) {
                printf("Ingestion::processMsgBundle: Token recreation FAIL\n");
                printf("ret = %x", ret);
                return false;
            }

            /*
            printf("Received token:\n");
            for(int l=0; l<SGX_CMAC_MAC_SIZE; l++) {
                printf("%x", tkn_ptr[l]);
            }
            printf("Verify created token:\n");
            for(int l=0; l<SGX_CMAC_MAC_SIZE; l++) {
                printf("%x", token[l]);
            }
            */

            int diff = memcmp(token, tkn_ptr, SGX_CMAC_MAC_SIZE);
            if(diff!=0) {
                printf("Ingestion::processMsgBundle: Tokens do not match FAIL\n");
                return false;
            }

            dm_ptr+= msg_size;
            tkn_ptr+=SGX_CMAC_MAC_SIZE;
        }
    }

    // Append msgbundle to g_ing.buffer

    MsgBuffer &msg_queue = *(g_ing.buffer);

    pthread_mutex_lock(&msg_queue.mutex);
    uint32_t head = msg_queue.reserved;
    if (head + num_msgs > g_ing.max_buffer_size) {
        pthread_mutex_unlock(&msg_queue.mutex);
        printf("Ingestions: Max %u messages exceeded\n",
            g_ing.max_buffer_size);
        return false;
    }
    msg_queue.reserved += num_msgs;
    pthread_mutex_unlock(&msg_queue.mutex);

    memmove(msg_queue.buf + head * msg_size,
        dec_msgbundle, num_msgs * msg_size);

    pthread_mutex_lock(&msg_queue.mutex);
    msg_queue.inserted += num_msgs;
    pthread_mutex_unlock(&msg_queue.mutex);

    free(dec_msgbundle);
    return true;
}

void Ingestion::generateClientKeys(sgx_aes_gcm_128bit_key_t &ESK)
{
    //printf("In Ingestion::genCK, num_clients = %d, client_start = %d, client_end = %d\n",
    //    cnum, cstart, cnum + cstart);
    uint32_t num_ing_nodes = g_teems_config.num_ingestion_nodes;

    for(uint32_t i=0; i<cnum; i++)
    {
        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);

        uint32_t client_num = cstart + (i * num_ing_nodes);
        memcpy(iv, (uint8_t*) (&client_num), sizeof(client_num));

        sgx_status_t ret = SGX_SUCCESS;
        ret = sgx_rijndael128GCM_encrypt((const sgx_aes_gcm_128bit_key_t *) (ESK),
            zeroes, SGX_AESGCM_KEY_SIZE, (uint8_t*) (clients[i].key), iv,
            SGX_AESGCM_IV_SIZE, NULL, 0, &tag);
        if(ret!=SGX_SUCCESS) {
            printf("Ingestion::GCM FAIL\n");
        }

        /*
        printf("Client %d Ingestion: (after Gen) ", client_num);
        for(int j=0;j<SGX_AESGCM_KEY_SIZE;j++) {
            printf("%x", (clients[i].key)[j]);
        }
        printf("\n");
        */

    }
}

sgx_aes_gcm_128bit_key_t* Ingestion::getClientKey(uint32_t lcid)
{
    return(&(clients[lcid].key));
}