/*
 * This file is part of VoidArchiveTool.
 *
 * Copyright (C) 2025 Yanczi
 *
 * Void Archive Toolis free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

#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::array<uint8_t, crypto_stream_chacha20_ietf_NONCEBYTES> nonce_local;
    randombytes_buf(nonce_local.data(), nonce_local.size());

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

    std::vector<char> output;
    output.insert(output.end(), 
        reinterpret_cast<const char*>(nonce_local.data()),
        reinterpret_cast<const char*>(nonce_local.data()) + nonce_local.size());

    output.insert(output.end(), tmp.begin(), tmp.end());

    return output;
}

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

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

    keyReady = true;
}

void EncryptionManager::saveKey(const std::string& path, bool hpp)
{
    const std::string sig = SIGNATURE_KEY_FILE;
    const int8_t 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());

    const uint64_t crcKey = XXH64(keyVec.data(), keyVec.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(sig.data(), sig.length());
    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.close();

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

// Generowanie pliku nagłówkowego CPP z kluczem i nonce
void EncryptionManager::saveCppHeadFile(const std::string& path)
{
    const uint32_t keySize = crypto_stream_chacha20_ietf_KEYBYTES;

    std::ofstream file(path + ".hpp");

    file << "// Plik wygenerowany przez " << PROGRAM_TITLE << " " << PROGRAM_VERSION << std::endl;
    file << std::endl;
    file << std::endl;
    file << "#pragma once" << std::endl;
    file << "#include <array>" << std::endl;
    file << "#include <cstdint>" << std::endl;
    file << std::endl;
    file << "namespace enc" << std::endl;
    file << "{" << std::endl;
    file << "   // Klucz deszyfrujący" << std::endl;
    file << "   const std::array<uint8_t, " << keySize << "> key{" << std::endl;
    file << "       " << toHex(key.data(), key.size()) << std::endl;
    file << "   };" << std::endl;
    file << std::endl;
    file << "} //namespace" << std::endl;

    file.close();
}

std::string EncryptionManager::toHex(const unsigned char* data, size_t len)
{
    std::ostringstream oss;
    oss << std::hex << std::setfill('0');
    for (size_t i = 0; i < len; ++i) {
        oss << "0x" << std::setw(2) << static_cast<int>(data[i]);
        if (i + 1 != len) oss << ", ";
        if ((i + 1) % 12 == 0 && i + 1 != len) oss << "\n        ";
    }
    return oss.str();
}

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

    const std::string signature = SIGNATURE_KEY_FILE;
    std::vector<char> sig(signature.size());
    int8_t ver;
    int time;

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

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

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

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

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

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

// Deszyfracja
std::vector<char> EncryptionManager::decrypt(const std::vector<char>& crypt)
{
    const size_t cryptoSize = crypto_stream_chacha20_ietf_NONCEBYTES;

    std::array<uint8_t, cryptoSize> nonce_local;
    std::memcpy(nonce_local.data(), 
        reinterpret_cast<const uint8_t*>(crypt.data()), cryptoSize);


    const size_t rawSize = crypt.size() - cryptoSize;
    std::vector<char> tmp(rawSize);
    std::memcpy(tmp.data(), crypt.data() + cryptoSize, rawSize);

    std::vector<char> raw(rawSize);

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