#include "VaultManager.h"
#include <QDateTime>
#include <QDebug>
#include <QSqlRecord>
#include "DatabaseManager.h"

VaultManager::VaultManager() {
	if (sodium_init() < 0) {
		qCritical() << "Libsodium could not be initialized!";
	}
}

VaultManager::~VaultManager() {
	lock();
}

bool VaultManager::init() {
	return true;
}

bool VaultManager::isInitialized() {
	return !DatabaseManager::instance().getVaultMeta("encrypted_database_key").isEmpty();
}

bool VaultManager::setup(const QString &masterPassword) {
	QByteArray salt(crypto_pwhash_SALTBYTES, 0);
	randombytes_buf(salt.data(), salt.size());

	QByteArray dek(crypto_aead_xchacha20poly1305_ietf_KEYBYTES, 0);
	randombytes_buf(dek.data(), dek.size());

	unsigned long long opslimit = crypto_pwhash_OPSLIMIT_INTERACTIVE;
	unsigned long long memlimit = crypto_pwhash_MEMLIMIT_INTERACTIVE;

	QByteArray kek = deriveKEK(masterPassword, salt, memlimit, opslimit);

	QByteArray nonce(crypto_aead_xchacha20poly1305_ietf_NPUBBYTES, 0);
	randombytes_buf(nonce.data(), nonce.size());

	QByteArray encryptedDek(dek.size() + crypto_aead_xchacha20poly1305_ietf_ABYTES, 0);
	unsigned long long outLen;

	crypto_aead_xchacha20poly1305_ietf_encrypt((unsigned char *)encryptedDek.data(), &outLen, (unsigned char *)dek.data(), dek.size(), nullptr, 0, nullptr,
																						 (unsigned char *)nonce.data(), (unsigned char *)kek.data());

	DatabaseManager::instance().setVaultMeta("argon2_salt", salt);
	DatabaseManager::instance().setVaultMeta("encrypted_database_key", encryptedDek);
	DatabaseManager::instance().setVaultMeta("database_key_nonce", nonce);
	DatabaseManager::instance().setVaultMeta("argon2_memory_cost", QByteArray::number(memlimit));
	DatabaseManager::instance().setVaultMeta("argon2_time_cost", QByteArray::number(opslimit));

	databaseKey = dek;
	return true;
}

bool VaultManager::unlock(const QString &masterPassword) {
	QByteArray salt = DatabaseManager::instance().getVaultMeta("argon2_salt");
	QByteArray encryptedDek = DatabaseManager::instance().getVaultMeta("encrypted_database_key");
	QByteArray nonce = DatabaseManager::instance().getVaultMeta("database_key_nonce");
	unsigned long long memlimit = DatabaseManager::instance().getVaultMeta("argon2_memory_cost").toULongLong();
	unsigned long long opslimit = DatabaseManager::instance().getVaultMeta("argon2_time_cost").toULongLong();

	QByteArray kek = deriveKEK(masterPassword, salt, memlimit, opslimit);

	QByteArray dek(crypto_aead_xchacha20poly1305_ietf_KEYBYTES, 0);
	unsigned long long outLen;

	if (crypto_aead_xchacha20poly1305_ietf_decrypt((unsigned char *)dek.data(), &outLen, nullptr, (unsigned char *)encryptedDek.data(), encryptedDek.size(), nullptr, 0,
																								 (unsigned char *)nonce.data(), (unsigned char *)kek.data()) != 0) {
		return false;
	}

	databaseKey = dek;
	return true;
}

void VaultManager::lock() {
	sodium_memzero(databaseKey.data(), databaseKey.size());
	databaseKey.clear();
}

QByteArray VaultManager::deriveKEK(const QString &password, const QByteArray &salt, unsigned long long memlimit, unsigned long long opslimit) {
	QByteArray kek(crypto_aead_xchacha20poly1305_ietf_KEYBYTES, 0);
	QByteArray passwordBytes = password.toUtf8();

	if (crypto_pwhash((unsigned char *)kek.data(), kek.size(), passwordBytes.data(), passwordBytes.size(), (unsigned char *)salt.data(), opslimit, memlimit,
										crypto_pwhash_ALG_ARGON2ID13) != 0) {
		qCritical() << "Argon2id failed!";
	}
	return kek;
}

bool VaultManager::savePassword(const QString &origin, const QString &username, const QString &password) {
	if (!isUnlocked()) {
		return false;
	}

	QByteArray nonce(crypto_aead_xchacha20poly1305_ietf_NPUBBYTES, 0);
	randombytes_buf(nonce.data(), nonce.size());

	QByteArray passwordBytes = password.toUtf8();
	QByteArray ciphertext(passwordBytes.size() + crypto_aead_xchacha20poly1305_ietf_ABYTES, 0);
	unsigned long long outLen;

	crypto_aead_xchacha20poly1305_ietf_encrypt((unsigned char *)ciphertext.data(), &outLen, (unsigned char *)passwordBytes.data(), passwordBytes.size(), nullptr, 0, nullptr,
																						 (unsigned char *)nonce.data(), (unsigned char *)databaseKey.data());

	QSqlQuery query;
	// Check if it exists
	query.prepare("SELECT id FROM passwords WHERE origin = ? AND username = ?");
	query.addBindValue(origin);
	query.addBindValue(username);

	if (query.exec() && query.next()) {
		int id = query.value(0).toInt();
		query.prepare(
			"UPDATE passwords SET password_ciphertext = ?, "
			"password_nonce = ?, updated_at = ? WHERE id = ?");
		query.addBindValue(ciphertext);
		query.addBindValue(nonce);
		query.addBindValue(QDateTime::currentDateTime());
		query.addBindValue(id);
	} else {
		query.prepare(
			"INSERT INTO passwords (origin, username, password_ciphertext, "
			"password_nonce, created_at, updated_at) "
			"VALUES (?, ?, ?, ?, ?, ?)");
		query.addBindValue(origin);
		query.addBindValue(username);
		query.addBindValue(ciphertext);
		query.addBindValue(nonce);
		query.addBindValue(QDateTime::currentDateTime());
		query.addBindValue(QDateTime::currentDateTime());
	}

	return query.exec();
}

QList<PasswordEntry> VaultManager::getPasswords(const QString &origin) {
	QList<PasswordEntry> entries;
	if (!isUnlocked()) {
		return entries;
	}

	QSqlQuery query;
	if (origin.isEmpty()) {
		query.prepare(
			"SELECT id, origin, username, password_ciphertext, "
			"password_nonce, created_at, updated_at FROM passwords");
	} else {
		query.prepare(
			"SELECT id, origin, username, password_ciphertext, password_nonce, "
			"created_at, updated_at FROM passwords WHERE origin = ?");
		query.addBindValue(origin);
	}

	if (!query.exec()) {
		return entries;
	}

	while (query.next()) {
		QByteArray ciphertext = query.value(3).toByteArray();
		QByteArray nonce = query.value(4).toByteArray();
		QByteArray decrypted(ciphertext.size() - crypto_aead_xchacha20poly1305_ietf_ABYTES, 0);
		unsigned long long outLen;

		if (crypto_aead_xchacha20poly1305_ietf_decrypt((unsigned char *)decrypted.data(), &outLen, nullptr, (unsigned char *)ciphertext.data(), ciphertext.size(), nullptr, 0,
																									 (unsigned char *)nonce.data(), (unsigned char *)databaseKey.data()) == 0) {
			PasswordEntry entry;
			entry.id = query.value(0).toInt();
			entry.origin = query.value(1).toString();
			entry.username = query.value(2).toString();
			entry.password = QString::fromUtf8(decrypted);
			entry.created_at = query.value(5).toDateTime();
			entry.updated_at = query.value(6).toDateTime();
			entries.append(entry);
		}
	}

	return entries;
}