#include "wallet_store.h"
#include <sodium.h>
#include <stdexcept>

namespace {
constexpr int KeyBytes = crypto_aead_xchacha20poly1305_ietf_KEYBYTES;
constexpr int MaxWalletBytes = 8 * 1024 * 1024;
const QByteArray Aad = "Nerbus encrypted wallet v1|NRB|" + NrbGenesis.toLatin1();
constexpr unsigned long long KdfOps = 3;
constexpr size_t KdfMemory = 64 * 1024 * 1024;
QByteArray readLimited(const QString &path) {
    QFile file(path);
    if (!file.open(QIODevice::ReadOnly) || file.size() < 1 || file.size() > MaxWalletBytes)
        walletError("Cannot read wallet file.");
    return file.readAll();
}
QJsonObject walletHeader(const QByteArray &file) {
    QJsonParseError parse;
    const auto doc = QJsonDocument::fromJson(file, &parse);
    const auto header = doc.object();
    if (parse.error != QJsonParseError::NoError || !doc.isObject() ||
        header["format"] != "NerbusWallet" || header["genesis"].toString() != NrbGenesis ||
        (header["version"].toInt() != 1 &&
         !(header["version"].toInt() == 2 && header["protection"] == "none")))
        walletError("Unsupported wallet file.");
    return header;
}
QJsonObject walletContents(const QJsonObject &data) {
    if (data["genesis"].toString() != NrbGenesis || !data["keys"].isArray() ||
        data["keys"].toArray().size() > 1000)
        walletError("Invalid wallet contents.");
    return data;
}
void atomicWrite(const QString &path, const QByteArray &data) {
    QSaveFile file(path);
    file.setDirectWriteFallback(false);
    if (!file.open(QIODevice::WriteOnly) ||
        !file.setPermissions(QFileDevice::ReadOwner | QFileDevice::WriteOwner) ||
        file.write(data) != data.size() || !file.commit())
        walletError("Cannot save wallet file. Check the folder permissions and free space.");
}
unsigned char *newSecret() {
    auto *p = static_cast<unsigned char *>(sodium_malloc(KeyBytes));
    if (!p) walletError("Cannot allocate secure wallet memory.");
    return p;
}
void derive(QByteArray &password, const QByteArray &salt, unsigned char *key) {
    const int result = crypto_pwhash(key, KeyBytes, password.constData(), password.size(),
        reinterpret_cast<const unsigned char *>(salt.constData()), KdfOps, KdfMemory, crypto_pwhash_ALG_ARGON2ID13);
    wipe(password);
    if (result != 0) walletError("Not enough memory to unlock the wallet.");
}
}

void wipe(QByteArray &bytes) {
    if (!bytes.isEmpty()) sodium_memzero(bytes.data(), bytes.size());
    bytes.clear();
}
void walletError(const QString &message) { throw std::runtime_error(message.toUtf8().constData()); }

QJsonObject offlineCore(QJsonObject request) {
    QProcess process;
    QString executable = QCoreApplication::applicationDirPath() + "/nerbus-wallet-core";
#ifdef Q_OS_WIN
    executable += ".exe";
#endif
    if (!QFileInfo::exists(executable)) walletError("The wallet signing component is missing. Extract the complete wallet package.");
    process.setProgram(executable);
    process.setArguments({"-wallet-json"});
    process.setProcessChannelMode(QProcess::SeparateChannels);
    process.start();
    if (!process.waitForStarted(10000)) walletError("Cannot start the wallet signing component.");
    QByteArray payload = QJsonDocument(request).toJson(QJsonDocument::Compact);
    request = {};
    const qint64 expected = payload.size();
    const qint64 written = process.write(payload);
    wipe(payload);
    process.closeWriteChannel();
    if (written != expected || !process.waitForFinished(60000)) {
        process.kill(); process.waitForFinished(3000);
        walletError("The wallet signing component did not finish.");
    }
    QByteArray output = process.readAllStandardOutput();
    QByteArray errors = process.readAllStandardError();
    wipe(errors); // Never display process diagnostics that could contain private input.
    QJsonParseError parse;
    QJsonDocument doc = QJsonDocument::fromJson(output, &parse);
    wipe(output);
    if (parse.error != QJsonParseError::NoError || !doc.isObject())
        walletError("The wallet signing component returned an invalid response.");
    QJsonObject result = doc.object();
    if (process.exitStatus() != QProcess::NormalExit || process.exitCode() != 0 || result.contains("error"))
        walletError(result.value("error").toString("Wallet operation failed."));
    return result;
}

WalletStore::WalletStore(QString path) : filename(std::move(path)) {}
WalletStore::~WalletStore() { lock(); }
bool WalletStore::exists() const { return QFileInfo::exists(filename); }
bool WalletStore::fileRequiresPassword(const QString &path) {
    QByteArray file = readLimited(path);
    try {
        const bool protectedFile = walletHeader(file)["version"].toInt() == 1;
        wipe(file);
        return protectedFile;
    } catch (...) { wipe(file); throw; }
}

QJsonObject WalletStore::decrypt(const QByteArray &file, QByteArray &password,
                                unsigned char *outKey, QByteArray &outSalt) {
    QJsonObject header = walletHeader(file);
    if (header["version"].toInt() == 2) {
        wipe(password);
        outSalt.clear();
        return walletContents(header["data"].toObject());
    }
    if (!outKey || header["kdf"] != "argon2id13" ||
        header["ops"].toInt() != int(KdfOps) || header["memory"].toInt() != int(KdfMemory))
        walletError("Unsupported wallet file.");
    outSalt = QByteArray::fromBase64(header["salt"].toString().toLatin1());
    QByteArray nonce = QByteArray::fromBase64(header["nonce"].toString().toLatin1());
    QByteArray cipher = QByteArray::fromBase64(header["ciphertext"].toString().toLatin1());
    if (outSalt.size() != crypto_pwhash_SALTBYTES || nonce.size() != crypto_aead_xchacha20poly1305_ietf_NPUBBYTES ||
        cipher.size() < crypto_aead_xchacha20poly1305_ietf_ABYTES)
        walletError("Incomplete wallet file.");
    derive(password, outSalt, outKey);
    QByteArray clear(cipher.size(), Qt::Uninitialized);
    unsigned long long length = 0;
    int rc = crypto_aead_xchacha20poly1305_ietf_decrypt(
        reinterpret_cast<unsigned char *>(clear.data()), &length, nullptr,
        reinterpret_cast<const unsigned char *>(cipher.constData()), cipher.size(),
        reinterpret_cast<const unsigned char *>(Aad.constData()), Aad.size(),
        reinterpret_cast<const unsigned char *>(nonce.constData()), outKey);
    if (rc != 0) { wipe(clear); walletError("Incorrect password or damaged wallet file."); }
    clear.resize(qsizetype(length));
    QJsonParseError parse;
    QJsonDocument payload = QJsonDocument::fromJson(clear, &parse);
    wipe(clear);
    auto data = payload.object();
    if (parse.error != QJsonParseError::NoError)
        walletError("Invalid wallet contents.");
    return walletContents(data);
}

void WalletStore::save(const QJsonObject &data) {
    if (!unlocked()) walletError("Open the wallet first.");
    QByteArray clear = QJsonDocument(data).toJson(QJsonDocument::Compact);
    if (plainOpen) {
        const QJsonObject header{{"format", "NerbusWallet"}, {"version", 2},
            {"genesis", NrbGenesis}, {"protection", "none"}, {"data", data}};
        QByteArray file = QJsonDocument(header).toJson(QJsonDocument::Compact);
        try { atomicWrite(filename, file); }
        catch (...) { wipe(file); wipe(clear); throw; }
        wipe(file);
        wipe(plaintext);
        plaintext = std::move(clear);
        return;
    }
    QByteArray nonce(crypto_aead_xchacha20poly1305_ietf_NPUBBYTES, Qt::Uninitialized);
    randombytes_buf(nonce.data(), nonce.size());
    QByteArray cipher(clear.size() + crypto_aead_xchacha20poly1305_ietf_ABYTES, Qt::Uninitialized);
    unsigned long long length = 0;
    crypto_aead_xchacha20poly1305_ietf_encrypt(reinterpret_cast<unsigned char *>(cipher.data()), &length,
        reinterpret_cast<const unsigned char *>(clear.constData()), clear.size(),
        reinterpret_cast<const unsigned char *>(Aad.constData()), Aad.size(), nullptr,
        reinterpret_cast<const unsigned char *>(nonce.constData()), key);
    QJsonObject header{{"format", "NerbusWallet"}, {"version", 1}, {"genesis", NrbGenesis},
        {"kdf", "argon2id13"}, {"ops", int(KdfOps)}, {"memory", int(KdfMemory)},
        {"salt", QString::fromLatin1(salt.toBase64())}, {"nonce", QString::fromLatin1(nonce.toBase64())},
        {"ciphertext", QString::fromLatin1(cipher.toBase64())}};
    try { atomicWrite(filename, QJsonDocument(header).toJson(QJsonDocument::Compact)); }
    catch (...) { wipe(clear); throw; }
    wipe(plaintext);
    plaintext = std::move(clear);
}

void WalletStore::create(QByteArray password) {
    if (exists()) { wipe(password); walletError("A wallet already exists here."); }
    if (!QDir().mkpath(QFileInfo(filename).absolutePath())) walletError("Cannot create the wallet folder.");
    lock();
    try {
        plainOpen = password.isEmpty();
        if (!plainOpen) {
            key = newSecret();
            salt.resize(crypto_pwhash_SALTBYTES);
            randombytes_buf(salt.data(), salt.size());
            derive(password, salt, key);
        }
        save({{"genesis", NrbGenesis}, {"keys", QJsonArray{}},
            {"created", QDateTime::currentDateTimeUtc().toString(Qt::ISODate)}});
    } catch (...) { wipe(password); lock(); throw; }
}

void WalletStore::unlock(QByteArray password) {
    lock();
    QByteArray file;
    try {
        file = readLimited(filename);
        plainOpen = walletHeader(file)["version"].toInt() == 2;
        if (!plainOpen) key = newSecret();
        QJsonObject data = decrypt(file, password, key, salt);
        plaintext = QJsonDocument(data).toJson(QJsonDocument::Compact);
        wipe(file);
    } catch (...) { wipe(file); wipe(password); lock(); throw; }
}
void WalletStore::changePassword(QByteArray password) {
    if (!unlocked()) { wipe(password); walletError("Open the wallet first."); }
    unsigned char *candidate = nullptr;
    QByteArray candidateSalt;
    try {
        if (!password.isEmpty()) {
            candidate = newSecret();
            candidateSalt.resize(crypto_pwhash_SALTBYTES);
            randombytes_buf(candidateSalt.data(), candidateSalt.size());
            derive(password, candidateSalt, candidate);
        }
    } catch (...) { wipe(password); if (candidate) sodium_free(candidate); throw; }
    auto *oldKey = key;
    const QByteArray oldSalt = salt;
    const bool wasPlain = plainOpen;
    key = candidate; salt = candidateSalt; plainOpen = candidate == nullptr;
    try { save(QJsonDocument::fromJson(plaintext).object()); }
    catch (...) {
        if (candidate) sodium_free(candidate);
        key = oldKey; salt = oldSalt; plainOpen = wasPlain;
        throw;
    }
    if (oldKey) sodium_free(oldKey);
}
void WalletStore::lock() {
    wipe(plaintext);
    if (key) { sodium_free(key); key = nullptr; }
    plainOpen = false;
    salt.clear();
}
QJsonArray WalletStore::entries() const {
    if (!unlocked()) walletError("Open the wallet first.");
    return QJsonDocument::fromJson(plaintext).object()["keys"].toArray();
}
QString WalletStore::addKey(const QJsonObject &pair, const QString &label) {
    if (!unlocked()) walletError("Open the wallet first.");
    QString address = pair["address"].toString();
    if (address.isEmpty() || pair["privateKey"].toString().isEmpty()) walletError("Invalid key response.");
    QJsonObject data = QJsonDocument::fromJson(plaintext).object();
    QJsonArray keys = data["keys"].toArray();
    for (const auto &entry : keys) if (entry.toObject()["address"] == address) return address;
    if (keys.size() >= 1000) walletError("This wallet already contains 1,000 addresses.");
    keys.append(QJsonObject{{"address", address}, {"privateKey", pair["privateKey"]},
        {"label", label.left(200)}, {"created", QDateTime::currentDateTimeUtc().toString(Qt::ISODate)}});
    data["keys"] = keys;
    save(data); // Persist the private key before returning an address to the UI.
    return address;
}
void WalletStore::backup(const QString &destination) const {
    if (QFileInfo(destination).absoluteFilePath() == QFileInfo(filename).absoluteFilePath())
        walletError("Choose a different location for the backup.");
    QByteArray file = readLimited(filename);
    try { atomicWrite(destination, file); }
    catch (...) { wipe(file); throw; }
    wipe(file);
}
void WalletStore::restore(const QString &source, QByteArray password) {
    if (exists()) { wipe(password); walletError("Restore into a new wallet folder to preserve the current wallet."); }
    QByteArray file = readLimited(source);
    auto *candidate = newSecret();
    QByteArray candidateSalt;
    try {
        const auto data = decrypt(file, password, candidate, candidateSalt);
        if (!QDir().mkpath(QFileInfo(filename).absolutePath())) walletError("Cannot create the wallet folder.");
        atomicWrite(filename, file);
        lock();
        plainOpen = walletHeader(file)["version"].toInt() == 2;
        if (plainOpen) sodium_free(candidate);
        else { key = candidate; salt = candidateSalt; }
        candidate = nullptr;
        plaintext = QJsonDocument(data).toJson(QJsonDocument::Compact);
    } catch (...) { wipe(file); wipe(password); if (candidate) sodium_free(candidate); throw; }
    wipe(file);
    wipe(password);
}