// ObsidianDragon Wallet - ImGui Edition
// Copyright 2024-2026 The Hush Developers
// Released under the GPLv3

#include "secure_vault.h"
#include "platform.h"

#include <sodium.h>

#include <fstream>
#include <filesystem>
#include <cstring>
#include <algorithm>
#include "../util/logger.h"

#ifndef _WIN32
#include <fcntl.h>
#include <unistd.h>
#endif

namespace fs = std::filesystem;

namespace dragonx {
namespace util {

// Vault file format version
static constexpr uint8_t VAULT_VERSION = 0x01;

// Argon2id parameters — moderate: ~256 MB, ~3 ops
// Tuned for short PINs to make brute-force costly
static constexpr unsigned long long ARGON2_MEMLIMIT = crypto_pwhash_MEMLIMIT_MODERATE;
static constexpr unsigned long long ARGON2_OPSLIMIT = crypto_pwhash_OPSLIMIT_MODERATE;

SecureVault::SecureVault() {
    // Ensure libsodium is initialized
    if (sodium_init() < 0) {
        // sodium_init returns 0 on success, 1 if already initialized, -1 on failure
        // We'll proceed anyway — the functions will fail gracefully
    }
}

SecureVault::~SecureVault() = default;

std::string SecureVault::getVaultPath() {
    return (fs::path(Platform::getConfigDir()) / "vault.dat").string();
}

bool SecureVault::hasVault() const {
    std::error_code ec;
    return fs::exists(getVaultPath(), ec);
}

bool SecureVault::isValidPin(const std::string& pin) {
    if (pin.size() < 4 || pin.size() > 8) return false;
    return std::all_of(pin.begin(), pin.end(), ::isdigit);
}

bool SecureVault::deriveKey(const std::string& pin,
                            const uint8_t* salt, size_t saltLen,
                            uint8_t* key, size_t keyLen) const {
    if (saltLen < crypto_pwhash_SALTBYTES) return false;

    int rc = crypto_pwhash(
        key, keyLen,
        pin.c_str(), pin.size(),
        salt,
        ARGON2_OPSLIMIT,
        ARGON2_MEMLIMIT,
        crypto_pwhash_ALG_ARGON2ID13
    );
    return rc == 0;
}

void SecureVault::secureZero(void* ptr, size_t len) {
    sodium_memzero(ptr, len);
}

bool SecureVault::store(const std::string& pin, const std::string& passphrase) {
    if (!isValidPin(pin)) {
        DEBUG_LOGF("[SecureVault] Invalid PIN format\n");
        return false;
    }
    if (passphrase.empty()) {
        DEBUG_LOGF("[SecureVault] Empty passphrase\n");
        return false;
    }

    // Generate random salt and nonce
    uint8_t salt[crypto_pwhash_SALTBYTES];
    uint8_t nonce[crypto_secretbox_NONCEBYTES];
    randombytes_buf(salt, sizeof(salt));
    randombytes_buf(nonce, sizeof(nonce));

    // Derive key from PIN
    uint8_t key[crypto_secretbox_KEYBYTES];
    if (!deriveKey(pin, salt, sizeof(salt), key, sizeof(key))) {
        DEBUG_LOGF("[SecureVault] Key derivation failed\n");
        return false;
    }

    // Encrypt passphrase
    size_t plainLen = passphrase.size();
    size_t cipherLen = plainLen + crypto_secretbox_MACBYTES;
    std::vector<uint8_t> ciphertext(cipherLen);

    int rc = crypto_secretbox_easy(
        ciphertext.data(),
        reinterpret_cast<const uint8_t*>(passphrase.c_str()),
        plainLen,
        nonce,
        key
    );

    // Wipe key immediately
    secureZero(key, sizeof(key));

    if (rc != 0) {
        DEBUG_LOGF("[SecureVault] Encryption failed\n");
        return false;
    }

    // Write vault file: version + salt + nonce + ciphertext.
    // Atomic + owner-only (0600): the vault holds the PIN-encrypted passphrase, so it must
    // not be world-readable (a copy enables an offline brute-force of the short PIN), and a
    // crash mid-write must not leave a truncated vault.
    std::string vaultPath = getVaultPath();
    std::string content;
    content.reserve(1 + sizeof(salt) + sizeof(nonce) + cipherLen);
    content.push_back(static_cast<char>(VAULT_VERSION));
    content.append(reinterpret_cast<const char*>(salt), sizeof(salt));
    content.append(reinterpret_cast<const char*>(nonce), sizeof(nonce));
    content.append(reinterpret_cast<const char*>(ciphertext.data()), cipherLen);

    if (!Platform::writeFileAtomically(vaultPath, content, /*restrictPermissions=*/true)) {
        DEBUG_LOGF("[SecureVault] Write failed\n");
        return false;
    }

    DEBUG_LOGF("[SecureVault] Vault saved (%zu bytes ciphertext)\n", cipherLen);
    return true;
}

bool SecureVault::retrieve(const std::string& pin, std::string& passphrase) const {
    if (!isValidPin(pin)) return false;

    std::string vaultPath = getVaultPath();
    std::ifstream in(vaultPath, std::ios::binary);
    if (!in.is_open()) return false;

    // Read entire file
    in.seekg(0, std::ios::end);
    size_t fileSize = in.tellg();
    in.seekg(0);

    // Minimum size: 1 (version) + SALTBYTES + NONCEBYTES + MACBYTES + 1 (min plaintext)
    size_t minSize = 1 + crypto_pwhash_SALTBYTES + crypto_secretbox_NONCEBYTES
                   + crypto_secretbox_MACBYTES + 1;
    if (fileSize < minSize) {
        DEBUG_LOGF("[SecureVault] Vault file too small\n");
        return false;
    }

    std::vector<uint8_t> data(fileSize);
    in.read(reinterpret_cast<char*>(data.data()), fileSize);
    in.close();

    // Parse header
    size_t offset = 0;
    uint8_t version = data[offset++];
    if (version != VAULT_VERSION) {
        DEBUG_LOGF("[SecureVault] Unknown vault version: %d\n", version);
        return false;
    }

    const uint8_t* salt = &data[offset];
    offset += crypto_pwhash_SALTBYTES;

    const uint8_t* nonce = &data[offset];
    offset += crypto_secretbox_NONCEBYTES;

    const uint8_t* ciphertext = &data[offset];
    size_t cipherLen = fileSize - offset;
    size_t plainLen = cipherLen - crypto_secretbox_MACBYTES;

    // Derive key
    uint8_t key[crypto_secretbox_KEYBYTES];
    if (!deriveKey(pin, salt, crypto_pwhash_SALTBYTES, key, sizeof(key))) {
        DEBUG_LOGF("[SecureVault] Key derivation failed during retrieve\n");
        return false;
    }

    // Decrypt
    std::vector<uint8_t> plaintext(plainLen);
    int rc = crypto_secretbox_open_easy(
        plaintext.data(),
        ciphertext,
        cipherLen,
        nonce,
        key
    );

    secureZero(key, sizeof(key));

    if (rc != 0) {
        // Wrong PIN or corrupted data
        return false;
    }

    passphrase.assign(reinterpret_cast<const char*>(plaintext.data()), plainLen);
    secureZero(plaintext.data(), plainLen);

    return true;
}

bool SecureVault::changePin(const std::string& oldPin, const std::string& newPin) {
    std::string passphrase;
    if (!retrieve(oldPin, passphrase)) {
        return false;
    }

    bool ok = store(newPin, passphrase);
    secureZero(&passphrase[0], passphrase.size());
    return ok;
}

bool SecureVault::removeVault() {
    std::string vaultPath = getVaultPath();
    std::error_code ec;
    if (fs::exists(vaultPath, ec)) {
        // Overwrite with zeros before removing (best-effort secure delete)
        {
            std::ifstream probe(vaultPath, std::ios::binary | std::ios::ate);
            size_t sz = probe.is_open() ? (size_t)probe.tellg() : 0;
            probe.close();
            if (sz > 0) {
                std::vector<uint8_t> zeros(sz, 0);
                std::ofstream zap(vaultPath, std::ios::binary);
                if (zap.is_open()) {
                    zap.write(reinterpret_cast<const char*>(zeros.data()), sz);
                    zap.flush();
                    zap.close();
                    // Force the zeros to stable storage before unlinking — otherwise the
                    // write may never leave the OS cache (best-effort; not a guarantee on
                    // CoW/journaling filesystems or wear-leveling SSDs).
#ifndef _WIN32
                    int fd = ::open(vaultPath.c_str(), O_WRONLY);
                    if (fd >= 0) { ::fsync(fd); ::close(fd); }
#endif
                }
            }
        }
        return fs::remove(vaultPath, ec);
    }
    return true; // Didn't exist — that's fine
}

} // namespace util
} // namespace dragonx