#include "duckdb/common/exception/conversion_exception.hpp"
#include "duckdb/common/encryption_key_manager.hpp"
#include "duckdb/common/encryption_functions.hpp"
#include "duckdb/main/attached_database.hpp"
#include "mbedtls_wrapper.hpp"
#include "duckdb/storage/storage_manager.hpp"
#include "duckdb/storage/storage_info.hpp"
#include "duckdb/common/typedefs.hpp"

namespace duckdb {

constexpr uint32_t AdditionalAuthenticatedData::INITIAL_AAD_CAPACITY;

AdditionalAuthenticatedData::~AdditionalAuthenticatedData() = default;

data_ptr_t AdditionalAuthenticatedData::data() const {
	return additional_authenticated_data->GetData();
}

idx_t AdditionalAuthenticatedData::size() const {
	return additional_authenticated_data->GetPosition();
}

void AdditionalAuthenticatedData::WriteStringData(const std::string &val) const {
	additional_authenticated_data->WriteData(reinterpret_cast<const_data_ptr_t>(val.data()), val.size());
}

EncryptionEngine::EncryptionEngine() {
}

EncryptionEngine::~EncryptionEngine() {
}

const_data_ptr_t EncryptionEngine::GetKeyFromCache(DatabaseInstance &db, const string &key_name) {
	auto &keys = EncryptionKeyManager::Get(db);
	return keys.GetKey(key_name);
}

bool EncryptionEngine::ContainsKey(DatabaseInstance &db, const string &key_name) {
	auto &keys = EncryptionKeyManager::Get(db);
	return keys.HasKey(key_name);
}

void EncryptionEngine::AddKeyToCache(DatabaseInstance &db, data_ptr_t key, const string &key_name, bool wipe) {
	auto &keys = EncryptionKeyManager::Get(db);
	if (!keys.HasKey(key_name)) {
		keys.AddKey(key_name, key);
	} else {
		duckdb_mbedtls::MbedTlsWrapper::AESStateMBEDTLS::SecureClearData(key,
		                                                                 MainHeader::DEFAULT_ENCRYPTION_KEY_LENGTH);
	}
}

string EncryptionEngine::AddKeyToCache(DatabaseInstance &db, data_ptr_t key) {
	auto &keys = EncryptionKeyManager::Get(db);
	const auto key_id = keys.GenerateRandomKeyID();

	if (!keys.HasKey(key_id)) {
		keys.AddKey(key_id, key);
	} else {
		duckdb_mbedtls::MbedTlsWrapper::AESStateMBEDTLS::SecureClearData(key,
		                                                                 MainHeader::DEFAULT_ENCRYPTION_KEY_LENGTH);
	}

	return key_id;
}

void EncryptionEngine::AddTempKeyToCache(DatabaseInstance &db) {
	//! Add a temporary key to the cache
	const auto length = MainHeader::DEFAULT_ENCRYPTION_KEY_LENGTH;
	data_t temp_key[length];

	// we cannot generate temporary keys with read-only enabled
	auto metadata = make_uniq<EncryptionStateMetadata>(EncryptionTypes::GCM, length, EncryptionTypes::V0_1);
	auto encryption_state = db.GetEncryptionUtil(false)->CreateEncryptionState(std::move(metadata));
	encryption_state->GenerateRandomData(temp_key, length);

	string key_id = "temp_key";
	AddKeyToCache(db, temp_key, key_id);
}

void EncryptionEngine::EncryptBlock(AttachedDatabase &attached_db, const string &key_id, FileBuffer &block,
                                    FileBuffer &temp_buffer_manager, uint64_t delta) {
	auto &db = attached_db.GetDatabase();
	data_ptr_t block_offset_internal = temp_buffer_manager.InternalBuffer();
	auto encrypt_key = GetKeyFromCache(db, key_id);
	auto version = attached_db.GetStorageManager().GetEncryptionVersion();
	auto cipher = attached_db.GetStorageManager().GetCipher();
	auto metadata = make_uniq<EncryptionStateMetadata>(cipher, MainHeader::DEFAULT_ENCRYPTION_KEY_LENGTH, version);
	auto encryption_state = db.GetEncryptionUtil(attached_db.IsReadOnly())->CreateEncryptionState(std::move(metadata));

	EncryptionTag tag;
	EncryptionNonce nonce(cipher, version);
	encryption_state->GenerateRandomData(nonce.data(), nonce.size());

	//! store the nonce at the start of the block
	memcpy(block_offset_internal, nonce.data(), nonce.size());
	encryption_state->InitializeEncryption(nonce, encrypt_key);

	auto checksum_offset = block.InternalBuffer() + delta;
	auto encryption_checksum_offset = block_offset_internal + delta;
	auto size = block.size + Storage::DEFAULT_BLOCK_HEADER_SIZE;

	//! encrypt the data including the checksum
	auto aes_res = encryption_state->Process(checksum_offset, size, encryption_checksum_offset, size);

	if (aes_res != size) {
		throw IOException("Block encryption failure: in- and output size differ (%llu/%llu)", size, aes_res);
	}

	//! Finalize and extract the tag
	encryption_state->Finalize(block.InternalBuffer() + delta, 0, tag.data(), tag.size());

	//! store the generated tag *behind* the nonce (but still at the beginning of the block)
	memcpy(block_offset_internal + nonce.size(), tag.data(), tag.size());
}

void EncryptionEngine::DecryptBlock(AttachedDatabase &attached_db, const string &key_id, data_ptr_t internal_buffer,
                                    uint64_t block_size, uint64_t delta) {
	//! initialize encryption state
	auto &db = attached_db.GetDatabase();
	auto version = attached_db.GetStorageManager().GetEncryptionVersion();
	auto cipher = attached_db.GetStorageManager().GetCipher();
	auto metadata = make_uniq<EncryptionStateMetadata>(cipher, MainHeader::DEFAULT_ENCRYPTION_KEY_LENGTH, version);
	auto decrypt_key = GetKeyFromCache(db, key_id);
	auto encryption_state = db.GetEncryptionUtil(attached_db.IsReadOnly())->CreateEncryptionState(std::move(metadata));

	//! load the stored nonce and tag
	EncryptionTag tag;
	EncryptionNonce nonce(cipher, version);
	memcpy(nonce.data(), internal_buffer, nonce.size());
	memcpy(tag.data(), internal_buffer + nonce.size(), tag.size());

	//! Initialize the decryption
	encryption_state->InitializeDecryption(nonce, decrypt_key);

	auto checksum_offset = internal_buffer + delta;
	auto size = block_size + Storage::DEFAULT_BLOCK_HEADER_SIZE;

	//! decrypt the block including the checksum
	auto aes_res = encryption_state->Process(checksum_offset, size, checksum_offset, size);

	if (aes_res != block_size + Storage::DEFAULT_BLOCK_HEADER_SIZE) {
		throw IOException("Block decryption failure: in- and output size differ (%llu/%llu)", size, aes_res);
	}

	//! check the tag
	encryption_state->Finalize(internal_buffer + delta, 0, tag.data(), tag.size());
}

void EncryptionEngine::EncryptTemporaryBuffer(DatabaseInstance &db, data_ptr_t buffer, idx_t buffer_size,
                                              data_ptr_t metadata) {
	if (!ContainsKey(db, "temp_key")) {
		AddTempKeyToCache(db);
	}

	auto temp_key = GetKeyFromCache(db, "temp_key");
	// we cannot encrypt temp buffers in read-only mode
	auto encryption_util = db.GetEncryptionUtil(false);
	// we hard-code GCM here for now, it's the safest and we don't know what is configured here
	auto state_metadata = make_uniq<EncryptionStateMetadata>(
	    EncryptionTypes::GCM, MainHeader::DEFAULT_ENCRYPTION_KEY_LENGTH, EncryptionTypes::V0_1);
	auto encryption_state = encryption_util->CreateEncryptionState(std::move(state_metadata));

	// zero-out the metadata buffer
	memset(metadata, 0, DEFAULT_ENCRYPTED_BUFFER_HEADER_SIZE);

	EncryptionTag tag;
	EncryptionNonce nonce(EncryptionTypes::CipherType::GCM, EncryptionTypes::V0_1);

	encryption_state->GenerateRandomData(nonce.data(), nonce.size());

	//! store the nonce at the start of metadata buffer
	memcpy(metadata, nonce.data(), nonce.size());
	encryption_state->InitializeEncryption(nonce, temp_key);

	auto aes_res = encryption_state->Process(buffer, buffer_size, buffer, buffer_size);

	if (aes_res != buffer_size) {
		throw IOException("Temporary buffer encryption failure: in- and output size differ (%llu/%llu)", buffer_size,
		                  aes_res);
	}

	//! Finalize and extract the tag
	encryption_state->Finalize(buffer, 0, tag.data(), tag.size());

	//! store the generated tag after consequetively the nonce
	memcpy(metadata + nonce.size(), tag.data(), tag.size());

	// check if tag is correctly stored
	D_ASSERT(memcmp(tag.data(), metadata + nonce.size(), tag.size()) == 0);
}

static void DecryptBuffer(EncryptionState &encryption_state, const_data_ptr_t temp_key, data_ptr_t buffer,
                          idx_t buffer_size, data_ptr_t metadata) {
	//! load the stored nonce and tag
	EncryptionTag tag;
	EncryptionNonce nonce(encryption_state.metadata->GetCipher(), encryption_state.metadata->GetVersion());
	memcpy(nonce.data(), metadata, nonce.size());
	memcpy(tag.data(), metadata + nonce.size(), tag.size());

	//! Initialize the decryption
	encryption_state.InitializeDecryption(nonce, temp_key);

	auto aes_res = encryption_state.Process(buffer, buffer_size, buffer, buffer_size);

	if (aes_res != buffer_size) {
		throw IOException("Buffer decryption failure: in- and output size differ (%llu/%llu)", buffer_size, aes_res);
	}

	//! check the tag
	encryption_state.Finalize(buffer, 0, tag.data(), tag.size());
}

void EncryptionEngine::DecryptTemporaryBuffer(DatabaseInstance &db, data_ptr_t buffer, idx_t buffer_size,
                                              data_ptr_t metadata) {
	//! initialize encryption state
	auto encryption_util = db.GetEncryptionUtil(false);
	auto temp_key = GetKeyFromCache(db, "temp_key");
	auto state_metadata = make_uniq<EncryptionStateMetadata>(
	    EncryptionTypes::GCM, MainHeader::DEFAULT_ENCRYPTION_KEY_LENGTH, EncryptionTypes::EncryptionVersion::V0_1);
	auto encryption_state = encryption_util->CreateEncryptionState(std::move(state_metadata));

	DecryptBuffer(*encryption_state, temp_key, buffer, buffer_size, metadata);
}

} // namespace duckdb
