#include "EncryptionManager.h"

EncryptionManager::EncryptionManager()
    :keyReady(false)
{
	if (sodium_init() < 0) {
		throw std::runtime_error("libsodium init failed");
	}

    keyReady = false;
    generateKeys();
}

std::vector<char> EncryptionManager::encrypt(const std::vector<char>& raw)
{
	std::vector<char> crypt(raw.size());
    
    // Generowanie kluczy
    // generateKeys();

    if (crypto_stream_chacha20_ietf_xor_ic(
        reinterpret_cast<unsigned char*>(crypt.data()),
        reinterpret_cast<const unsigned char*>(raw.data()),
        static_cast<unsigned long long>(raw.size()),
        nonce.data(), 0, key.data()) != 0)
    {
        throw std::runtime_error("crypto_stream_chacha20_ietf_xor_ic failed");
    }

    return crypt;
}

void EncryptionManager::generateKeys()
{
    if (keyReady) return;

    //randombytes_buf(key.data(), key.size());
    crypto_stream_chacha20_ietf_keygen(key.data());
    randombytes_buf(nonce.data(), nonce.size());

    keyReady = true;
}

void EncryptionManager::saveKey(const std::string& path, bool hpp)
{
    const int sig = SIGNATURE_KEY_FILE;
    const short ver = VERSION;

    // Wygeneruj time stamp
    std::time_t now = std::time(nullptr);
    const int time = static_cast<int>(now);

    // Wygeneruj crc kluczy
    std::vector<char> keyVec(reinterpret_cast<const char*>(key.data()),
        reinterpret_cast<const char*>(key.data()) + key.size());

    std::vector<char> nonceVec(reinterpret_cast<const char*>(nonce.data()),
        reinterpret_cast<const char*>(nonce.data()) + nonce.size());

    const uint64_t crcKey = XXH64(keyVec.data(), keyVec.size(), VERSION);
    const uint64_t crcNonce = XXH64(nonceVec.data(), nonceVec.size(), VERSION);

    // Zapisz ten śmietnik do pliku KEY
    std::ofstream file(path + ".key", std::ios::binary);
    if (!file) { std::cout << "Failed to save encryption key to file" << std::endl; }

    file.write(reinterpret_cast<const char*>(&sig), sizeof(sig));
    file.write(reinterpret_cast<const char*>(&ver), sizeof(ver));
    file.write(reinterpret_cast<const char*>(&time), sizeof(time));
    file.write(reinterpret_cast<const char*>(keyVec.data()), keyVec.size());
    file.write(reinterpret_cast<const char*>(&crcKey), sizeof(crcKey));
    file.write(reinterpret_cast<const char*>(nonceVec.data()), nonceVec.size());
    file.write(reinterpret_cast<const char*>(&crcNonce), sizeof(crcNonce));

    file.close();

    if (hpp) {saveCppHeadFile(path);}
}

// Generowanie pliku nagłówkowego CPP z kluczem i nonce
void EncryptionManager::saveCppHeadFile(const std::string& path)
{
    std::vector<unsigned char> keyVec(key.begin(), key.end());
    std::vector<unsigned char> nonceVec(nonce.begin(), nonce.end());

    const std::string headerText =
        "// Plik wygenerowany przy wykorzystaniu exPAK\n\n"
        "// Klucz deszyfrujący\n"
        "const std::array<unsigned char, crypto_stream_chacha20_ietf_KEYBYTES> key{" + toHex(keyVec) + "};\n\n"
        "// Ciąg nonce\n"
        "const std::array<unsigned char, crypto_stream_chacha20_ietf_NONCEBYTES> nonce{" + toHex(nonceVec) + "}; ";

    std::ofstream file(path + ".hh");
    file << headerText;
    file.close();
}

std::string EncryptionManager::toHex(const std::vector<unsigned char>& data)
{
    std::string bytes;
    int sk = data.size();
    int skp = 1;

    for (const auto& b : data)
    {
        std::stringstream ss;
        ss << "0x" << std::hex << std::uppercase << std::setw(2) << std::setfill('0')
            << static_cast<int>(b);

        bytes += "'";
        bytes += ss.str();
        bytes += "'";

        if (skp < sk)
        {
            bytes += ", ";
            skp++;
        }
    }

    return bytes;
}

// Wczytaj klucz
void EncryptionManager::loadKey(const std::string& path)
{
    std::ifstream file(path + ".key", std::ios::binary);

    int sig;
    short ver;
    int time;

    // Wczytaj
    file.read(reinterpret_cast<char*>(&sig), sizeof(sig));
    file.read(reinterpret_cast<char*>(&ver), sizeof(ver));

    // Sprawdź czy plik klucza jest poprawny
    if (sig != SIGNATURE_KEY_FILE || ver != VERSION)
    {
        throw std::runtime_error("Invalid key file!");
    }

    std::vector<char> keyVec(key.size());
    std::vector<char> nonceVec(nonce.size());
    uint64_t crcKey;
    uint64_t crcNonce;

    file.read(reinterpret_cast<char*>(&time), sizeof(time));
    file.read(keyVec.data(), keyVec.size());
    file.read(reinterpret_cast<char*>(&crcKey), sizeof(crcKey));
    file.read(nonceVec.data(), nonceVec.size());
    file.read(reinterpret_cast<char*>(&crcNonce), sizeof(crcNonce));

    // Sprawdź integralność klucza
    if (XXH64(keyVec.data(), keyVec.size(), VERSION) != crcKey
        || XXH64(nonceVec.data(), nonceVec.size(), VERSION) != crcNonce)
    {
        throw std::runtime_error("Key integrity error!");
    }

    file.close();
    
    // Przekonwertuj vector na array
    key = toArray<crypto_stream_chacha20_ietf_KEYBYTES>(keyVec);
    nonce = toArray<crypto_stream_chacha20_ietf_NONCEBYTES>(nonceVec);
}

// Deszyfracja
std::vector<char> EncryptionManager::decrypt(const std::vector<char>& crypt)
{
    std::vector<char> raw(crypt.size());

    if (crypto_stream_chacha20_ietf_xor(
        reinterpret_cast<unsigned char*>(raw.data()),
        reinterpret_cast<const unsigned char*>(crypt.data()),
        static_cast<unsigned long long>(crypt.size()),
        nonce.data(), key.data()) != 0)
    {
        throw std::runtime_error("Data decryption error!");
    }
    
    return raw;
}