#include "transaction_history_cache.h"

#include "../util/logger.h"
#include "../util/platform.h"

#include <nlohmann/json.hpp>
#include <sqlite3.h>
#include <sodium.h>

#include <algorithm>
#include <cstdio>
#include <filesystem>
#include <utility>

namespace fs = std::filesystem;
using json = nlohmann::json;

namespace dragonx {
namespace data {

namespace {

constexpr int kSchemaVersion = 1;
constexpr std::size_t kKeyBytes = 32;

struct Statement {
    sqlite3_stmt* handle = nullptr;

    Statement(sqlite3* db, const char* sql)
    {
        if (sqlite3_prepare_v2(db, sql, -1, &handle, nullptr) != SQLITE_OK) {
            handle = nullptr;
        }
    }

    ~Statement()
    {
        if (handle) sqlite3_finalize(handle);
    }

    Statement(const Statement&) = delete;
    Statement& operator=(const Statement&) = delete;
};

bool bindText(sqlite3_stmt* statement, int index, const std::string& value)
{
    return sqlite3_bind_text(statement, index, value.c_str(), -1, SQLITE_TRANSIENT) == SQLITE_OK;
}

bool bindBlob(sqlite3_stmt* statement, int index, const std::vector<unsigned char>& value)
{
    return sqlite3_bind_blob(statement, index, value.data(), static_cast<int>(value.size()), SQLITE_TRANSIENT) == SQLITE_OK;
}

std::vector<unsigned char> readBlob(sqlite3_stmt* statement, int index)
{
    const void* data = sqlite3_column_blob(statement, index);
    int bytes = sqlite3_column_bytes(statement, index);
    if (!data || bytes <= 0) return {};
    const auto* begin = static_cast<const unsigned char*>(data);
    return std::vector<unsigned char>(begin, begin + bytes);
}

std::string hexEncode(const unsigned char* bytes, std::size_t length)
{
    static constexpr char kHex[] = "0123456789abcdef";
    std::string output;
    output.resize(length * 2);
    for (std::size_t index = 0; index < length; ++index) {
        output[index * 2] = kHex[(bytes[index] >> 4) & 0x0F];
        output[index * 2 + 1] = kHex[bytes[index] & 0x0F];
    }
    return output;
}

json transactionToJson(const TransactionInfo& transaction)
{
    return json{
        {"txid", transaction.txid},
        {"type", transaction.type},
        {"amount", transaction.amount},
        {"timestamp", transaction.timestamp},
        {"confirmations", transaction.confirmations},
        {"address", transaction.address},
        {"from_address", transaction.from_address},
        {"memo", transaction.memo}
    };
}

TransactionInfo transactionFromJson(const json& source)
{
    TransactionInfo transaction;
    transaction.txid = source.value("txid", std::string());
    transaction.type = source.value("type", std::string());
    transaction.amount = source.value("amount", 0.0);
    transaction.timestamp = source.value("timestamp", static_cast<std::int64_t>(0));
    transaction.confirmations = source.value("confirmations", 0);
    transaction.address = source.value("address", std::string());
    transaction.from_address = source.value("from_address", std::string());
    transaction.memo = source.value("memo", std::string());
    return transaction;
}

std::string associatedDataForWallet(const std::string& walletHash)
{
    return std::string("obsidian-dragon-tx-history-v1:") + walletHash;
}

} // namespace

TransactionHistoryCache::TransactionHistoryCache()
    : TransactionHistoryCache(defaultDatabasePath())
{
}

TransactionHistoryCache::TransactionHistoryCache(std::string databasePath)
    : database_path_(std::move(databasePath))
{
    if (sodium_init() < 0) {
        DEBUG_LOGF("Failed to initialize libsodium for transaction history cache\n");
    }
}

TransactionHistoryCache::~TransactionHistoryCache()
{
    lockKey();
    close();
}

std::string TransactionHistoryCache::defaultDatabasePath()
{
    return (fs::path(util::Platform::getConfigDir()) / "transaction_history.sqlite").string();
}

std::string TransactionHistoryCache::walletIdentityFromAddresses(
    const std::vector<std::string>& shieldedAddresses,
    const std::vector<std::string>& transparentAddresses)
{
    std::vector<std::string> addresses;
    addresses.reserve(shieldedAddresses.size() + transparentAddresses.size());
    for (const auto& address : shieldedAddresses) {
        if (!address.empty()) addresses.push_back("z:" + address);
    }
    for (const auto& address : transparentAddresses) {
        if (!address.empty()) addresses.push_back("t:" + address);
    }
    if (addresses.empty()) return {};

    std::sort(addresses.begin(), addresses.end());
    std::string identity = "wallet-addresses-v1\n";
    for (const auto& address : addresses) {
        identity += address;
        identity += '\n';
    }
    return identity;
}

std::string TransactionHistoryCache::walletIdentityHash(const std::string& walletIdentity)
{
    unsigned char digest[crypto_generichash_BYTES];
    crypto_generichash(digest, sizeof(digest),
                       reinterpret_cast<const unsigned char*>(walletIdentity.data()),
                       walletIdentity.size(), nullptr, 0);
    return hexEncode(digest, sizeof(digest));
}

bool TransactionHistoryCache::ensureOpen()
{
    if (db_) return true;

    try {
        fs::path path(database_path_);
        if (!path.parent_path().empty()) fs::create_directories(path.parent_path());
    } catch (const std::exception& exception) {
        DEBUG_LOGF("Failed to create transaction history cache directory: %s\n", exception.what());
        return false;
    }

    sqlite3* openedDb = nullptr;
    if (sqlite3_open(database_path_.c_str(), &openedDb) != SQLITE_OK) {
        DEBUG_LOGF("Failed to open transaction history cache: %s\n",
                   openedDb ? sqlite3_errmsg(openedDb) : "unknown error");
        if (openedDb) sqlite3_close(openedDb);
        return false;
    }

    db_ = openedDb;
    sqlite3_busy_timeout(db_, 2000);
    exec("PRAGMA journal_mode=WAL");
    exec("PRAGMA synchronous=NORMAL");

    if (!createSchema()) {
        close();
        return false;
    }

    return true;
}

bool TransactionHistoryCache::unlockWithPassphrase(const std::string& walletIdentity,
                                                   const std::string& passphrase)
{
    if (walletIdentity.empty() || passphrase.empty() || !ensureOpen()) return false;

    std::string walletHash = walletIdentityHash(walletIdentity);
    std::vector<unsigned char> salt = getOrCreateSalt(walletHash);
    if (salt.empty()) return false;

    if (!deriveKey(passphrase, salt)) return false;
    unlocked_wallet_hash_ = std::move(walletHash);
    key_ready_ = true;
    return true;
}

void TransactionHistoryCache::lockKey()
{
    if (key_ready_) sodium_memzero(key_.data(), key_.size());
    key_ready_ = false;
    unlocked_wallet_hash_.clear();
}

bool TransactionHistoryCache::isUnlockedFor(const std::string& walletIdentity) const
{
    return key_ready_ && !walletIdentity.empty() &&
           unlocked_wallet_hash_ == walletIdentityHash(walletIdentity);
}

TransactionHistoryCache::LoadResult TransactionHistoryCache::load(
    const std::string& walletIdentity,
    int currentTipHeight,
    const std::string& currentTipHash)
{
    LoadResult result;
    if (!isUnlockedFor(walletIdentity) || !ensureOpen()) return result;

    std::string walletHash = walletIdentityHash(walletIdentity);
    int tipHeight = 0;
    std::string tipHash;
    std::time_t updatedAt = 0;
    std::vector<unsigned char> nonce;
    std::vector<unsigned char> cipherText;
    if (!readSnapshot(walletHash, tipHeight, tipHash, updatedAt, nonce, cipherText)) return result;

    if ((currentTipHeight > 0 && tipHeight > currentTipHeight) ||
        (currentTipHeight > 0 && tipHeight == currentTipHeight &&
         !currentTipHash.empty() && !tipHash.empty() && tipHash != currentTipHash)) {
        clearWalletByHash(walletHash);
        result.invalidated = true;
        return result;
    }

    std::string plainText;
    if (!decryptPayload(walletHash, nonce, cipherText, plainText)) return result;

    try {
        json payload = json::parse(plainText);
        if (payload.value("schema_version", 0) != kSchemaVersion) return result;
        if (payload.value("wallet_hash", std::string()) != walletHash) return result;
        if (!payload.contains("transactions") || !payload["transactions"].is_array()) return result;

        result.transactions.reserve(payload["transactions"].size());
        for (const auto& transactionJson : payload["transactions"]) {
            if (transactionJson.is_object()) {
                result.transactions.push_back(transactionFromJson(transactionJson));
            }
        }
        if (payload.contains("shielded_scan_heights") && payload["shielded_scan_heights"].is_object()) {
            for (auto it = payload["shielded_scan_heights"].begin();
                 it != payload["shielded_scan_heights"].end(); ++it) {
                if (!it.key().empty() && it.value().is_number_integer()) {
                    result.shieldedScanHeights[it.key()] = it.value().get<int>();
                }
            }
        }
        result.tipHeight = tipHeight;
        result.tipHash = tipHash;
        result.updatedAt = updatedAt;
        result.loaded = true;
    } catch (...) {
        result.transactions.clear();
    }

    sodium_memzero(plainText.data(), plainText.size());
    return result;
}

bool TransactionHistoryCache::replace(const std::string& walletIdentity,
                                      int tipHeight,
                                      const std::string& tipHash,
                                      const std::vector<TransactionInfo>& transactions,
                                      std::time_t updatedAt,
                                      const std::unordered_map<std::string, int>& shieldedScanHeights)
{
    if (!isUnlockedFor(walletIdentity) || !ensureOpen()) return false;

    std::string walletHash = walletIdentityHash(walletIdentity);
    json payload;
    payload["schema_version"] = kSchemaVersion;
    payload["wallet_hash"] = walletHash;
    payload["tip_height"] = tipHeight;
    payload["tip_hash"] = tipHash;
    payload["updated_at"] = static_cast<std::int64_t>(updatedAt);
    payload["transactions"] = json::array();
    for (const auto& transaction : transactions) {
        payload["transactions"].push_back(transactionToJson(transaction));
    }
    payload["shielded_scan_heights"] = json::object();
    for (const auto& [address, height] : shieldedScanHeights) {
        if (!address.empty() && height >= 0) {
            payload["shielded_scan_heights"][address] = height;
        }
    }

    std::string plainText = payload.dump();
    std::vector<unsigned char> nonce;
    std::vector<unsigned char> cipherText;
    bool encrypted = encryptPayload(walletHash, plainText, nonce, cipherText);
    sodium_memzero(plainText.data(), plainText.size());
    if (!encrypted) return false;

    Statement statement(db_,
        "INSERT OR REPLACE INTO transaction_history_snapshots "
        "(wallet_hash, schema_version, tip_height, tip_hash, updated_at, nonce, ciphertext) "
        "VALUES (?, ?, ?, ?, ?, ?, ?)");
    if (!statement.handle) return false;

    if (!bindText(statement.handle, 1, walletHash)) return false;
    sqlite3_bind_int(statement.handle, 2, kSchemaVersion);
    sqlite3_bind_int(statement.handle, 3, std::max(0, tipHeight));
    if (!bindText(statement.handle, 4, tipHash)) return false;
    sqlite3_bind_int64(statement.handle, 5, static_cast<sqlite3_int64>(updatedAt));
    if (!bindBlob(statement.handle, 6, nonce)) return false;
    if (!bindBlob(statement.handle, 7, cipherText)) return false;

    return sqlite3_step(statement.handle) == SQLITE_DONE;
}

void TransactionHistoryCache::clearWallet(const std::string& walletIdentity)
{
    if (walletIdentity.empty()) return;
    clearWalletByHash(walletIdentityHash(walletIdentity));
}

int TransactionHistoryCache::snapshotCount()
{
    if (!ensureOpen()) return 0;
    Statement statement(db_, "SELECT COUNT(*) FROM transaction_history_snapshots");
    if (!statement.handle || sqlite3_step(statement.handle) != SQLITE_ROW) return 0;
    return sqlite3_column_int(statement.handle, 0);
}

bool TransactionHistoryCache::exec(const char* sql)
{
    if (!db_) return false;
    char* error = nullptr;
    int result = sqlite3_exec(db_, sql, nullptr, nullptr, &error);
    if (result != SQLITE_OK) {
        DEBUG_LOGF("Transaction history cache SQL error: %s\n", error ? error : sqlite3_errmsg(db_));
        if (error) sqlite3_free(error);
        return false;
    }
    return true;
}

bool TransactionHistoryCache::createSchema()
{
    return exec("CREATE TABLE IF NOT EXISTS transaction_history_keys ("
                "wallet_hash TEXT PRIMARY KEY,"
                "salt BLOB NOT NULL)") &&
           exec("CREATE TABLE IF NOT EXISTS transaction_history_snapshots ("
                "wallet_hash TEXT PRIMARY KEY,"
                "schema_version INTEGER NOT NULL,"
                "tip_height INTEGER NOT NULL,"
                "tip_hash TEXT NOT NULL,"
                "updated_at INTEGER NOT NULL,"
                "nonce BLOB NOT NULL,"
                "ciphertext BLOB NOT NULL)");
}

std::vector<unsigned char> TransactionHistoryCache::getOrCreateSalt(const std::string& walletHash)
{
    if (!ensureOpen()) return {};

    {
        Statement statement(db_, "SELECT salt FROM transaction_history_keys WHERE wallet_hash = ?");
        if (!statement.handle) return {};
        if (!bindText(statement.handle, 1, walletHash)) return {};
        if (sqlite3_step(statement.handle) == SQLITE_ROW) {
            auto salt = readBlob(statement.handle, 0);
            if (salt.size() == crypto_pwhash_SALTBYTES) return salt;
        }
    }

    std::vector<unsigned char> salt(crypto_pwhash_SALTBYTES);
    randombytes_buf(salt.data(), salt.size());

    Statement insert(db_,
        "INSERT OR REPLACE INTO transaction_history_keys (wallet_hash, salt) VALUES (?, ?)");
    if (!insert.handle) return {};
    if (!bindText(insert.handle, 1, walletHash)) return {};
    if (!bindBlob(insert.handle, 2, salt)) return {};
    if (sqlite3_step(insert.handle) != SQLITE_DONE) return {};
    return salt;
}

bool TransactionHistoryCache::deriveKey(const std::string& passphrase,
                                        const std::vector<unsigned char>& salt)
{
    if (salt.size() != crypto_pwhash_SALTBYTES) return false;
    unsigned char derived[kKeyBytes];
    int result = crypto_pwhash(derived, sizeof(derived),
                               passphrase.c_str(), passphrase.size(),
                               salt.data(),
                               crypto_pwhash_OPSLIMIT_INTERACTIVE,
                               crypto_pwhash_MEMLIMIT_INTERACTIVE,
                               crypto_pwhash_ALG_ARGON2ID13);
    if (result != 0) return false;
    std::copy(derived, derived + sizeof(derived), key_.begin());
    sodium_memzero(derived, sizeof(derived));
    return true;
}

bool TransactionHistoryCache::encryptPayload(const std::string& walletHash,
                                             const std::string& plainText,
                                             std::vector<unsigned char>& nonce,
                                             std::vector<unsigned char>& cipherText) const
{
    if (!key_ready_) return false;
    nonce.resize(crypto_aead_xchacha20poly1305_ietf_NPUBBYTES);
    randombytes_buf(nonce.data(), nonce.size());

    std::string associatedData = associatedDataForWallet(walletHash);
    cipherText.resize(plainText.size() + crypto_aead_xchacha20poly1305_ietf_ABYTES);
    unsigned long long cipherLength = 0;
    int result = crypto_aead_xchacha20poly1305_ietf_encrypt(
        cipherText.data(), &cipherLength,
        reinterpret_cast<const unsigned char*>(plainText.data()), plainText.size(),
        reinterpret_cast<const unsigned char*>(associatedData.data()), associatedData.size(),
        nullptr, nonce.data(), key_.data());
    if (result != 0) return false;
    cipherText.resize(static_cast<std::size_t>(cipherLength));
    return true;
}

bool TransactionHistoryCache::decryptPayload(const std::string& walletHash,
                                             const std::vector<unsigned char>& nonce,
                                             const std::vector<unsigned char>& cipherText,
                                             std::string& plainText) const
{
    if (!key_ready_ || nonce.size() != crypto_aead_xchacha20poly1305_ietf_NPUBBYTES ||
        cipherText.size() < crypto_aead_xchacha20poly1305_ietf_ABYTES) {
        return false;
    }

    std::string associatedData = associatedDataForWallet(walletHash);
    std::vector<unsigned char> plain(cipherText.size() - crypto_aead_xchacha20poly1305_ietf_ABYTES);
    unsigned long long plainLength = 0;
    int result = crypto_aead_xchacha20poly1305_ietf_decrypt(
        plain.data(), &plainLength, nullptr,
        cipherText.data(), cipherText.size(),
        reinterpret_cast<const unsigned char*>(associatedData.data()), associatedData.size(),
        nonce.data(), key_.data());
    if (result != 0) return false;

    plainText.assign(reinterpret_cast<const char*>(plain.data()), static_cast<std::size_t>(plainLength));
    sodium_memzero(plain.data(), plain.size());
    return true;
}

bool TransactionHistoryCache::readSnapshot(const std::string& walletHash,
                                           int& tipHeight,
                                           std::string& tipHash,
                                           std::time_t& updatedAt,
                                           std::vector<unsigned char>& nonce,
                                           std::vector<unsigned char>& cipherText)
{
    Statement statement(db_,
        "SELECT tip_height, tip_hash, updated_at, nonce, ciphertext "
        "FROM transaction_history_snapshots WHERE wallet_hash = ?");
    if (!statement.handle) return false;
    if (!bindText(statement.handle, 1, walletHash)) return false;
    if (sqlite3_step(statement.handle) != SQLITE_ROW) return false;

    tipHeight = sqlite3_column_int(statement.handle, 0);
    const unsigned char* tipHashText = sqlite3_column_text(statement.handle, 1);
    tipHash = tipHashText ? reinterpret_cast<const char*>(tipHashText) : std::string();
    updatedAt = static_cast<std::time_t>(sqlite3_column_int64(statement.handle, 2));
    nonce = readBlob(statement.handle, 3);
    cipherText = readBlob(statement.handle, 4);
    return !nonce.empty() && !cipherText.empty();
}

void TransactionHistoryCache::clearWalletByHash(const std::string& walletHash)
{
    if (!ensureOpen()) return;
    Statement statement(db_, "DELETE FROM transaction_history_snapshots WHERE wallet_hash = ?");
    if (!statement.handle) return;
    if (!bindText(statement.handle, 1, walletHash)) return;
    sqlite3_step(statement.handle);
}

void TransactionHistoryCache::close()
{
    if (!db_) return;
    sqlite3_close(db_);
    db_ = nullptr;
}

} // namespace data
} // namespace dragonx