#include "encrypt_wrapper.h"
#include <stdio.h>
#include <string.h>

#ifdef _MSC_VER
#define __bswap_16 _byteswap_ushort
#define __bswap_32 _byteswap_ulong
#elif defined(__APPLE__)
#include <libkern/OSByteOrder.h>
#define __bswap_16 OSSwapInt16
#define __bswap_32 OSSwapInt32
#define __bswap_64 OSSwapInt64
#else
#include <byteswap.h>
#endif
#include <stdlib.h>
#include "romulus_m.h"
#include "api.h"

#include <rand.h>

#ifdef _DEBUG
#define DEBUG_PRINT printf
#else
#define DEBUG_PRINT
#endif

#define MIN(a,b) ((a) < (b) ? (a) : (b))

randctx rCtx = {0};

int encrypt(packet *packet, unsigned char *outBuf, UINT16 *len, unsigned char *key)
{
    DEBUG_PRINT("line %d\n", __LINE__);
    if (sizeof(UINT16) + sizeof(UINT16) + packet->Size > *len)
    {
        return 1;
    }

    int packetSize = packet->Size;
    // enc
    // ciphertext - out, MAX 16 bytes larger than plaintext - defined as CRYPTO_ABYTES
    unsigned long long clen = packetSize + CRYPTO_ABYTES;
    unsigned char *c = (unsigned char *)calloc(clen, 1);

    // plaintext - in
    int mlen = packetSize;
    unsigned char *m = packet->Data;

    // additional text
    int adlen = sizeof(packet->DataType);
    unsigned char *ad = (unsigned char *)&packet->DataType;

    // n = nonce, CRYPTO_NPUBBYTES size 16
    unsigned char npub[CRYPTO_NPUBBYTES];
    DEBUG_PRINT("line %d\n", __LINE__);
    fill_random(npub, CRYPTO_NPUBBYTES);

    DEBUG_PRINT("line %d\n", __LINE__);
    int ret = romulus_m_encrypt(c, &clen, m, mlen, ad, adlen, 0, &npub[0], key);

    if (clen > *len + CRYPTO_NPUBBYTES || ret != 0)
    {
        free(c);
        return 1;
    }

    // type 0 1
    // Swap after adding it as additional data
    packet->DataType = __bswap_16(packet->DataType);
    int loc = 0;
    memcpy(&outBuf[loc], &packet->DataType, sizeof(UINT16));
    loc += sizeof(UINT16);

    DEBUG_PRINT("line %d\n", __LINE__);
    // data size 2 3
    UINT16 clen16 = (UINT16)clen;
    // we copy data and nonce
    clen16 += CRYPTO_NPUBBYTES;
    clen16 = __bswap_16(clen16);
    memcpy(&outBuf[loc], &clen16, sizeof(UINT16));
    loc += sizeof(UINT16);

    memcpy(&outBuf[loc], npub, CRYPTO_NPUBBYTES);
    loc += CRYPTO_NPUBBYTES;

    memcpy(&outBuf[loc], c, clen);
    loc += clen;

    *len = loc;

    DEBUG_PRINT("line %d\n", __LINE__);
    free(c);

    return 0;
}

int decrypt(packetEx *packet, unsigned char *inBuf, UINT16 len, unsigned char *key)
{
    int loc = 0;
    if (inBuf == 0)
    {
        return 1;
    }
    DEBUG_PRINT("line %d\n", __LINE__);
    memcpy(&packet->DataType, &inBuf[loc], sizeof(UINT16));
    loc += sizeof(UINT16);
    memcpy(&packet->Id, &inBuf[loc], sizeof(UINT32));
    loc += sizeof(UINT32);
    memcpy(&packet->Size, &inBuf[loc], sizeof(UINT16));
    loc += sizeof(UINT16);

    packet->DataType = __bswap_16(packet->DataType);
    packet->Size = __bswap_16(packet->Size);
    packet->Id = __bswap_32(packet->Id);

    if ((len - loc) < 0 || packet->Size < CRYPTO_NPUBBYTES || packet->Size > 1000)
    {
        return 1;
    }
    DEBUG_PRINT("line %d\n", __LINE__);

    // Data
    // Nonce|Data
    // dec
    // n = nonce, CRYPTO_NPUBBYTES size 16
    unsigned char *npub = &inBuf[loc];
    loc += CRYPTO_NPUBBYTES;

    // ciphertext - in, MAX 16 bytes larger than plaintext - defined as CRYPTO_ABYTES
    int clen = len - loc;
    unsigned char *c = &inBuf[loc];
    DEBUG_PRINT("line %d len:%d loc:%d\n", __LINE__, len, loc);

    // plaintext - out
    unsigned long long mlen = clen;
    unsigned char *m = (unsigned char *)calloc(clen, 1);

    // additional text
    int adlen = sizeof(packet->DataType);
    unsigned char *ad = (unsigned char *)&packet->DataType;

    DEBUG_PRINT("line %d mlen:%d clen:%d\n", __LINE__, mlen, clen);
    int ret = romulus_m_decrypt(m, &mlen, 0, c, clen, ad, adlen, npub, key);
    DEBUG_PRINT("line %d packet:%d, mlen %d\n", __LINE__, packet->Size, mlen);
    if (ret == 0)
    {
        DEBUG_PRINT("line %d\n", __LINE__);
        memcpy(packet->Data, m, mlen);
        packet->Size = (UINT16)mlen;
    }

    free(m);
    DEBUG_PRINT("line %d\n", __LINE__);
    return ret;
}

void fill_random(unsigned char* buffer, int length){
    DEBUG_PRINT("line %d\n", __LINE__);
    if(rCtx.randcnt == 0){
        randinit(&rCtx, RANDSIZL);
    }
    UINT32 secRandom;
    for (size_t i = 0; i < length; i += sizeof(UINT32))
    {
        secRandom = rand(&rCtx);
        memcpy(&buffer[i], &secRandom, MIN(sizeof(UINT32), length-i));
    }
}