#include <openssl/evp.h>
#include <openssl/err.h>
#include <string>
#include <iostream>
#include <fstream>

#include "cryptography.h"
#include "buffer.h"

Cryptography::Cryptography(const char* password)
{
	OpenSSL_add_all_algorithms();
	ERR_load_crypto_strings();

	if (!generate_key_and_iv_from_password(password))
	{
		std::cout << "Error generating key and IV from password\n";
		return;
	}
}

Cryptography::~Cryptography()
{
	ERR_free_strings();
	EVP_cleanup();
}

int Cryptography::encrypt(uint8_t* input_buffer, int32_t input_buffer_size)
{
	// Allocate memory for the ciphertext
	int ciphertext_len = 0;
	
	Buffer encrypted_buffer(input_buffer_size + EVP_CIPHER_block_size(EVP_aes_256_cbc()));

	// Create and initialize the context for encryption
	ctx = EVP_CIPHER_CTX_new();
	if (!ctx) handleErrors();

	if (1 != EVP_EncryptInit_ex(ctx, EVP_aes_256_cbc(), NULL, key, iv)) handleErrors();

	// Perform encryption
	if (1 != EVP_EncryptUpdate(ctx, encrypted_buffer.buffer, &ciphertext_len, input_buffer, input_buffer_size)) handleErrors();

	int final_len;
	if (1 != EVP_EncryptFinal_ex(ctx, encrypted_buffer.buffer + ciphertext_len, &final_len)) handleErrors();
	ciphertext_len += final_len;

	// Clean up the context 
	EVP_CIPHER_CTX_free(ctx);


	// store the encrypted buffer to file

	std::ofstream file("encrypted_file.bin", std::ios::binary);
	if (file) 
	{
		file.write((char*)encrypted_buffer.buffer, ciphertext_len);
		file.close();
	}



	return ciphertext_len;
}

int Cryptography::decrypt(uint8_t* out_buffer)
{
	// read the encrypted buffer from file
	


	std::ifstream file("encrypted_file.bin", std::ios::binary);
	if (!file) return 0;

	file.seekg(0, std::ios::end);
	int encrypted_buffer_len = file.tellg();
	file.seekg(0, std::ios::beg);

	Buffer encrypted_buffer(encrypted_buffer_len);

	file.read((char*)encrypted_buffer.buffer, encrypted_buffer_len);
	file.close();
	
	


	


	ctx = EVP_CIPHER_CTX_new();
	if (!ctx) handleErrors();

	if (1 != EVP_DecryptInit_ex(ctx, EVP_aes_256_cbc(), NULL, key, iv)) handleErrors();

	int decrypted_len = 0;
	if (1 != EVP_DecryptUpdate(ctx, out_buffer, &decrypted_len, encrypted_buffer.buffer, encrypted_buffer_len)) handleErrors();
	int final_len;
	if (1 != EVP_DecryptFinal_ex(ctx, out_buffer + decrypted_len, &final_len)) handleErrors();
	decrypted_len += final_len;

	// Clean up the context
	EVP_CIPHER_CTX_free(ctx);

	return decrypted_len;
}

bool Cryptography::generate_key_and_iv_from_password(const char* password)
{
	int iterations = 10000;

	// Derive key and IV using PBKDF2
	if (1 != PKCS5_PBKDF2_HMAC(password, strlen(password), nullptr, 0, iterations, EVP_sha256(), 32, key)) return false;
	if (1 != PKCS5_PBKDF2_HMAC(password, strlen(password), nullptr, 0, iterations, EVP_sha256(), 16, iv)) return false;
	return true;
}

void Cryptography::handleErrors()
{
	EVP_CIPHER_CTX_free(ctx);

	while (auto error = ERR_get_error())
	{
		char* error_string = ERR_error_string(error, nullptr);
		std::cout << error_string << "\n";
	}
	abort();
}