#include "../include/odhtdb/Encryption.hpp"
#include <sodium/crypto_aead_xchacha20poly1305.h>
#include <sodium/randombytes.h>
#include <cstring>

namespace odhtdb
{
    static_assert(ENCRYPTION_CHECKSUM_BYTE_SIZE == crypto_aead_xchacha20poly1305_ietf_ABYTES, "Encryption checksum key size has changed for some reason, oops...");
    
    Encryption::Encryption(const DataView &data, const DataView &_key)
    {
        cipherTextLength = crypto_aead_xchacha20poly1305_ietf_ABYTES + data.size;
        cipherText = new unsigned char[cipherTextLength];   
        
        if(_key.data)
        {
            if(_key.size != ENCRYPTION_KEY_BYTE_SIZE)
                throw EncryptionException("Encryption key is wrong size");
            memcpy(key, _key.data, _key.size);
        }
        else
            generateKey(key);
        
        randombytes_buf(nonce, ENCRYPTION_NONCE_BYTE_SIZE);
        if(crypto_aead_xchacha20poly1305_ietf_encrypt(cipherText, &cipherTextLength, (const unsigned char*)data.data, data.size, nullptr, 0, nullptr, nonce, key) < 0)
            throw EncryptionException("Failed to encrypt data");
    }
    
    Encryption::~Encryption()
    {
        delete[](cipherText);
    }
    
    DataView Encryption::getKey() const
    {
        return DataView((void*)key, ENCRYPTION_KEY_BYTE_SIZE);
    }
    
    DataView Encryption::getNonce() const
    {
        return DataView((void*)nonce, ENCRYPTION_NONCE_BYTE_SIZE);
    }
    
    DataView Encryption::getCipherText() const
    {
        return DataView((void*)cipherText, cipherTextLength);
    }
    
    void Encryption::generateKey(unsigned char *output)
    {
        crypto_aead_xchacha20poly1305_ietf_keygen(output);
    }
    
    Decryption::Decryption(const DataView &data, const DataView &nonce, const DataView &key)
    {
        decryptedTextLength = data.size;
        decryptedText = new unsigned char[decryptedTextLength];
        
        if(nonce.size < ENCRYPTION_NONCE_BYTE_SIZE)
            throw DecryptionException("Nonce is not big enough");
        
        if(key.size < ENCRYPTION_KEY_BYTE_SIZE)
            throw DecryptionException("Key is not big enough");
        
        if(crypto_aead_xchacha20poly1305_ietf_decrypt(decryptedText, &decryptedTextLength, nullptr, (const unsigned char*)data.data, data.size, nullptr, 0, (const unsigned char*)nonce.data, (const unsigned char*)key.data) < 0)
            throw DecryptionException("Failed to decrypt data");
    }
    
    Decryption::Decryption(Decryption &&other)
    {
        decryptedText = other.decryptedText;
        decryptedTextLength = other.decryptedTextLength;
        
        other.decryptedText = nullptr;
        other.decryptedTextLength = 0;
    }
    
    Decryption& Decryption::operator=(Decryption &&other)
    {
        decryptedText = other.decryptedText;
        decryptedTextLength = other.decryptedTextLength;
        
        other.decryptedText = nullptr;
        other.decryptedTextLength = 0;
        return *this;
    }
    
    Decryption::~Decryption()
    {
        delete[](decryptedText);
    }
    
    DataView Decryption::getDecryptedText() const
    {
        return DataView((void*)decryptedText, decryptedTextLength);
    }
}