#include "AES.h"
#include <assert.h>

// write data from src to array and increment array data pointer
static inline void __Array_writeNext(Array(u8)* dst, u8* src, size_t size){
    memcpy(dst->data, src, size);
    *dst = Array_sliceFrom(*dst, size);
}

// read data from array to dst and increment array data pointer
static inline void __Array_readNext(u8* dst, Array(u8)* src, size_t size){
    memcpy(dst, src->data, size);
    *src = Array_sliceFrom(*src, size);
}


//////////////////////////////////////////////////////////////////////////////
//                            AESBlockEncryptor                             //
//////////////////////////////////////////////////////////////////////////////

void AESBlockEncryptor_construct(AESBlockEncryptor* ptr,
    Array(u8) key, const br_block_cbcenc_class* enc_class)
{
    assert(key.size == 16 || key.size == 24 || key.size == 32);
    
    ptr->enc_class = enc_class;
    ptr->enc_class->init((void*)ptr->enc_keys, key.data, key.size);
    
    ptr->rng_ctx.vtable = &br_hmac_drbg_vtable;
    rng_init_sha256_seedFromSystem(&ptr->rng_ctx.vtable);
}

void AESBlockEncryptor_changeKey(AESBlockEncryptor* ptr, Array(u8) key)
{
    assert(key.size == 16 || key.size == 24 || key.size == 32);
    ptr->enc_class->init((void*)ptr->enc_keys, key.data, key.size);
}

Result(u32) AESBlockEncryptor_encrypt(AESBlockEncryptor* ptr,
    Array(u8) src, Array(u8) dst)
{
    Deferral(4);
    u32 encrypted_size = AESBlockEncryptor_calcDstSize(src.size);
    try_assert(dst.size >= encrypted_size);
    
    // generate random initial vector
    br_hmac_drbg_generate(&ptr->rng_ctx, ptr->iv, __AES_BLOCK_IV_SIZE);
    // write IV to the beginning of dst
    __Array_writeNext(&dst, ptr->iv, __AES_BLOCK_IV_SIZE);

    const EncryptedBlockHeader header = { .padding_size = 16 - src.size % 16 };
    // write header to buffer
    memcpy(ptr->buf, &header, sizeof(header));
    // encrypt header
    ptr->enc_class->run((void*)ptr->enc_keys, ptr->iv, ptr->buf, sizeof(header));
    // write encrypted header to dst
    __Array_writeNext(&dst, ptr->buf, sizeof(header));

    // encrypt full AESBlockEncryptor buffers
    while(src.size > __AES_BUFFER_SIZE){
        __Array_readNext(ptr->buf, &src, __AES_BUFFER_SIZE);
        ptr->enc_class->run((void*)ptr->enc_keys,
            ptr->iv,
            ptr->buf, __AES_BUFFER_SIZE);
        __Array_writeNext(&dst, ptr->buf, __AES_BUFFER_SIZE);
    }

    // encrypt buffer with remaining data
    if(src.size > 0){
        memcpy(ptr->buf, src.data, src.size);
        u32 src_size_padded = src.size + header.padding_size;
        memset(ptr->buf + src.size, 0, header.padding_size);
        ptr->enc_class->run((void*)ptr->enc_keys,
            ptr->iv,
            ptr->buf, src_size_padded);
        memcpy(dst.data, ptr->buf, src_size_padded);
    }

    Return RESULT_VALUE(u, encrypted_size);
}


//////////////////////////////////////////////////////////////////////////////
//                            AESBlockDecryptor                             //
//////////////////////////////////////////////////////////////////////////////

void AESBlockDecryptor_construct(AESBlockDecryptor* ptr,
    Array(u8) key, const br_block_cbcdec_class* dec_class)
{
    assert(key.size == 16 || key.size == 24 || key.size == 32);
    
    ptr->dec_class = dec_class;
    ptr->dec_class->init((void*)ptr->dec_keys, key.data, key.size);
}

void AESBlockDecryptor_changeKey(AESBlockDecryptor* ptr, Array(u8) key)
{
    assert(key.size == 16 || key.size == 24 || key.size == 32);
    ptr->dec_class->init((void*)ptr->dec_keys, key.data, key.size);
}

Result(u32) AESBlockDecryptor_decrypt(AESBlockDecryptor* ptr,
    Array(u8) src, Array(u8) dst)
{
    Deferral(4);
    u32 overhead_size = AESBlockEncryptor_calcDstSize(0);
    try_assert(src.size >= overhead_size);
    try_assert(src.size % 16 == 0 && "src must be array of 16-byte blocks");
    try_assert(dst.size >= src.size - overhead_size);
    
    // read IV from the beginning of src
    __Array_readNext(ptr->iv, &src, __AES_BLOCK_IV_SIZE);

    EncryptedBlockHeader header;
    // read encrypted header from src
    __Array_readNext((void*)&header, &src, sizeof(header));
    // decrypt header
    ptr->dec_class->run((void*)ptr->dec_keys, ptr->iv, &header, sizeof(header));
    // size of decrypted data without padding
    u32 decrypted_size = src.size - header.padding_size;
    const u32 src_size_padded = src.size;
    src.size = decrypted_size;

    // decrypt full buffers
    while(src.size > __AES_BUFFER_SIZE){
        __Array_readNext(ptr->buf, &src, __AES_BUFFER_SIZE);
        ptr->dec_class->run((void*)ptr->dec_keys, 
            ptr->iv, 
            ptr->buf, __AES_BUFFER_SIZE);
        __Array_writeNext(&dst, ptr->buf, __AES_BUFFER_SIZE);
    }
    
    // decrypt buffer with remaining data
    if(src.size > 0){
        memcpy(ptr->buf, src.data, src.size);
        ptr->dec_class->run((void*)ptr->dec_keys, 
            ptr->iv,
            ptr->buf, src_size_padded);
        memcpy(dst.data, ptr->buf, src.size);
    }

    Return RESULT_VALUE(u, decrypted_size);
}


//////////////////////////////////////////////////////////////////////////////
//                            AESStreamEncryptor                            //
//////////////////////////////////////////////////////////////////////////////

void AESStreamEncryptor_construct(AESStreamEncryptor* ptr, 
    Array(u8) key, const br_block_ctr_class* ctr_class)
{
    assert(key.size == 16 || key.size == 24 || key.size == 32);
    
    ptr->ctr_class = ctr_class;
    ptr->ctr_class->init((void*)ptr->ctr_keys, key.data, key.size);
    
    br_hmac_drbg_context rng_ctx;
    rng_ctx.vtable = &br_hmac_drbg_vtable;
    rng_init_sha256_seedFromSystem(&rng_ctx.vtable);
    br_hmac_drbg_generate(&rng_ctx, ptr->iv, __AES_STREAM_IV_SIZE);

    ptr->block_counter = 0;
}

void AESStreamEncryptor_changeKey(AESStreamEncryptor* ptr, Array(u8) key)
{
    assert(key.size == 16 || key.size == 24 || key.size == 32);
    ptr->ctr_class->init((void*)ptr->ctr_keys, key.data, key.size);
}

Result(u32) AESStreamEncryptor_encrypt(AESStreamEncryptor* ptr,
    Array(u8) src, Array(u8) dst)
{
    Deferral(4);

    u32 encrypted_size = src.size;
    // if it is the beginning of the stream, write IV
    if(ptr->block_counter == 0){
        __Array_writeNext(&dst, ptr->iv, __AES_STREAM_IV_SIZE);
        encrypted_size = AESStreamEncryptor_calcDstSize(encrypted_size);
    }
    try_assert(dst.size >= encrypted_size);
    
    // encrypt full buffers
    while(src.size > __AES_BUFFER_SIZE){
        __Array_readNext(ptr->buf, &src, __AES_BUFFER_SIZE);
        ptr->ctr_class->run((void*)ptr->ctr_keys, 
            ptr->iv, ptr->block_counter,
            ptr->buf, __AES_BUFFER_SIZE);
        __Array_writeNext(&dst, ptr->buf, __AES_BUFFER_SIZE);
    }

    // encrypt remaining data
    if(src.size > 0){
        memcpy(ptr->buf, src.data, src.size);
        ptr->ctr_class->run((void*)ptr->ctr_keys,
            ptr->iv, ptr->block_counter,
            ptr->buf, src.size);
        memcpy(dst.data, ptr->buf, src.size);
    }

    ptr->block_counter++;
    Return RESULT_VALUE(u, encrypted_size);
}


//////////////////////////////////////////////////////////////////////////////
//                            AESStreamDecryptor                            //
//////////////////////////////////////////////////////////////////////////////

void AESStreamDecryptor_construct(AESStreamDecryptor* ptr,
    Array(u8) key, const br_block_ctr_class* ctr_class)
{
    assert(key.size == 16 || key.size == 24 || key.size == 32);
    
    ptr->ctr_class = ctr_class;
    ptr->ctr_class->init((void*)ptr->ctr_keys, key.data, key.size);

    ptr->block_counter = 0;
}

void AESStreamDecryptor_changeKey(AESStreamDecryptor* ptr, Array(u8) key)
{
    assert(key.size == 16 || key.size == 24 || key.size == 32);
    ptr->ctr_class->init((void*)ptr->ctr_keys, key.data, key.size);
}

Result(u32) AESStreamDecryptor_decrypt(AESStreamDecryptor* ptr,
    Array(u8) src, Array(u8) dst)
{
    Deferral(4);

    // if it is the beginning of the stream, read IV
    if(ptr->block_counter == 0){
        __Array_readNext(ptr->iv, &src, __AES_STREAM_IV_SIZE);
    }
    // size without IV
    u32 decrypted_size = src.size;
    try_assert(dst.size >= decrypted_size);

    // decrypt full buffers
    while(src.size > __AES_BUFFER_SIZE){
        __Array_readNext(ptr->buf, &src, __AES_BUFFER_SIZE);
        ptr->ctr_class->run((void*)ptr->ctr_keys, 
            ptr->iv, ptr->block_counter,
            ptr->buf, __AES_BUFFER_SIZE);
        __Array_writeNext(&dst, ptr->buf, __AES_BUFFER_SIZE);
    }

    // decrypt remaining data
    if(src.size > 0){
        memcpy(ptr->buf, src.data, src.size);
        ptr->ctr_class->run((void*)ptr->ctr_keys,
            ptr->iv, ptr->block_counter,
            ptr->buf, src.size);
        memcpy(dst.data, ptr->buf, src.size);
    }

    ptr->block_counter++;
    Return RESULT_VALUE(u, decrypted_size);
}