#include <botan/auto_rng.h>
#include <botan/rsa.h>
#include <botan/pkcs8.h>
#include <botan/x509_key.h>
#include <botan/pubkey.h>
#include <botan/data_src.h>
#include <fstream>
#include <iostream>

namespace CryptoLib {
	std::pair<std::unique_ptr<Botan::Private_Key>, std::unique_ptr<Botan::Public_Key>> 
	generate_keys(int key_size = 2048);

	std::pair<std::unique_ptr<Botan::Private_Key>, std::unique_ptr<Botan::Public_Key>> 
	generate_keys(int key_size) {
		Botan::AutoSeeded_RNG rng;
		std::unique_ptr<Botan::RSA_PrivateKey> private_key(new Botan::RSA_PrivateKey(rng, key_size));
		std::unique_ptr<Botan::RSA_PublicKey> public_key(new Botan::RSA_PublicKey(*private_key));
	
		return std::make_pair(std::move(private_key), std::move(public_key));
	}
	
	std::vector<uint8_t> encrypt(const Botan::Public_Key& public_key, const std::string& data) {
	    std::vector<uint8_t> data_vector(data.begin(), data.end());
	    Botan::AutoSeeded_RNG rng;
	    Botan::PK_Encryptor_EME encryptor(public_key, rng, "OAEP(SHA-256)");
	    return encryptor.encrypt(data_vector, rng);
	}
	
	std::string decrypt(const Botan::Private_Key& private_key, const std::vector<uint8_t>& encrypted_data) {
	    std::vector<uint8_t> encrypted_vector(encrypted_data.begin(), encrypted_data.end());
	    Botan::AutoSeeded_RNG rng;
	    Botan::PK_Decryptor_EME decryptor(private_key, rng, "OAEP(SHA-256)");
	    auto result = decryptor.decrypt(encrypted_vector);
	    return std::string(result.begin(), result.end());
	}
	
	std::string public_key_to_string(const Botan::Public_Key& public_key) {
		return Botan::X509::PEM_encode(public_key);
	}
	
	std::unique_ptr<Botan::Public_Key> public_key_from_string(const std::string& key_string) {
		Botan::DataSource_Memory ds(key_string);
		return std::unique_ptr<Botan::Public_Key>(Botan::X509::load_key(ds));
	}
	
	std::string private_key_to_string(const Botan::Private_Key& private_key) {
		return Botan::PKCS8::PEM_encode(private_key);
	}
	
	std::unique_ptr<Botan::Private_Key> private_key_from_string(const std::string& key_string) {
		Botan::DataSource_Memory ds(key_string);
		return std::unique_ptr<Botan::Private_Key>(Botan::PKCS8::load_key(ds));
	}
	
	void save_private_key(const Botan::Private_Key& private_key, const std::string& filename) {
		std::ofstream file(filename);
		if (file.is_open()) {
			file << Botan::PKCS8::PEM_encode(private_key);
			file.close();
		} else {
			throw std::runtime_error("Cannot open file for writing: " + filename);
		}
	}
	
	void save_public_key(const Botan::Public_Key& public_key, const std::string& filename) {
		std::ofstream file(filename);
		if (file.is_open()) {
			file << Botan::X509::PEM_encode(public_key);
			file.close();
		} else {
			throw std::runtime_error("Cannot open file for writing: " + filename);
		}
	}
	
	std::unique_ptr<Botan::Private_Key> load_private_key(const std::string& filename) {
		std::ifstream file(filename);
		if (file.is_open()) {
			std::string content((std::istreambuf_iterator<char>(file)),
							   std::istreambuf_iterator<char>());
			file.close();
			Botan::DataSource_Memory ds(content);
			return std::unique_ptr<Botan::Private_Key>(Botan::PKCS8::load_key(ds));
		} else {
			throw std::runtime_error("Cannot open file for reading: " + filename);
		}
	}
	
	std::unique_ptr<Botan::Public_Key> load_public_key(const std::string& filename) {
		std::ifstream file(filename);
		if (file.is_open()) {
			std::string content((std::istreambuf_iterator<char>(file)),
							   std::istreambuf_iterator<char>());
			file.close();
			Botan::DataSource_Memory ds(content);
			return std::unique_ptr<Botan::Public_Key>(Botan::X509::load_key(ds));
		} else {
			throw std::runtime_error("Cannot open file for reading: " + filename);
		}
	}
	
	std::string bytes_to_base64(const std::vector<unsigned char>& data) {
		static const char* const chars = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
		std::string result;
		size_t i = 0;
	
		while (i < data.size()) {
			uint32_t octet_a = i < data.size() ? data[i++] : 0;
			uint32_t octet_b = i < data.size() ? data[i++] : 0;
			uint32_t octet_c = i < data.size() ? data[i++] : 0;
		
			uint32_t triple = (octet_a << 16) + (octet_b << 8) + octet_c;
	
			result += chars[(triple >> 18) & 63];
			result += chars[(triple >> 12) & 63];
			result += chars[(triple >> 6) & 63];
			result += chars[triple & 63];
		}
		
		// Add padding
		switch (data.size() % 3) {
			case 1:
				result[result.size() - 2] = '=';
				result[result.size() - 1] = '=';
				break;
			case 2:
				result[result.size() - 1] = '=';
				break;
		}
		
		return result;
	}
		
	std::vector<unsigned char> base64_to_bytes(const std::string& input) {
		static const std::string chars = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
		std::vector<unsigned char> result;
	
		int val = 0, valb = -8;
		for (unsigned char c : input) {
			if (c == '=') break;
			if (chars.find(c) == std::string::npos) continue;
		
			val = (val << 6) + chars.find(c);
			valb += 6;
			if (valb >= 0) {
				result.push_back((val >> valb) & 0xFF);
				valb -= 8;
			}
		}
		return result;
	}
};