/**
 * @file dbHandler.cpp
 * @brief DBHandler class that manages all interactions with the SQLite database.
 *
 * Responsibilities:
 * - Insert job records and metadata into relevant tables.
 * - Maintain and cache job ID mappings for fast lookup.
 * - Track which parts of history/epoch files have been processed.
 * - Provide user-level analytics (e.g., jobs per user).
 *
 * Notes:
 * - All common DB statements are prepared once at startup and reused.
 * - Most write functions use INSERT OR IGNORE + SELECT fallback.
 * - A small in-memory LRU cache is used to avoid repeated lookups for {ClusterId, ProcId}.
 * - I plan to split this into different files/modules
 * - I also plan to write a transaction.cpp file to wrap most transactions safely
 */

#include "condor_common.h"
#include "condor_config.h"
#include "condor_debug.h"

#include <iostream>
#include <sqlite3.h>
#include <vector>
#include <string>
#include <fstream>
#include <sstream> 
#include <unordered_map>
#include <functional>
#include <filesystem>
#include <list>
#include <set>
#include <thread>
#include <chrono>

#include "JobRecord.h"
#include "dbHandler.h"
#include "archiveMonitor.h"
#include "SavedQueries.h"

namespace conf = LibrarianConfigOptions;
namespace fs = std::filesystem;


#define ROLLBACK_AND_RETURN()\
    char* err; \
    if (sqlite3_exec(db_, "ROLLBACK;", nullptr, nullptr, &err) != SQLITE_OK) { \
        dprintf(D_ERROR, "Error: Failed to rollback: %s\n", err); \
        sqlite3_free(err); \
    } \
    return false;

/* Macro to do schema updates on a per version basis provided a raw string of statements
#define UPDATE_SCHEMA(details)\
    char* err; \
    if (sqlite3_exec(db_, details.c_str(), nullptr, nullptr, &err) != SQLITE_OK) { \
        printf("Error: Failed to execute schema update: %s\n", err); \
        sqlite3_free(err); \
        ROLLBACK_AND_RETURN(); \
    }
*/

// -------------------------
// DBHandler Constructor / Destructor
// -------------------------

bool DBHandler::initialize() {
    jobIdCache_.setLimit((size_t)(config[conf::i::DBMaxJobCacheSize]));

    if (sqlite3_open(config[conf::str::DBPath].c_str(), &db_) != SQLITE_OK) {
        std::string err = sqlite3_errmsg(db_);
        dprintf(D_ERROR, "Failed to open SQLite Data Base: %s\n", err.c_str());
        sqlite3_close(db_); // technically redundant if db is already bad, but harmless
        return false;
    }

    char* errMsg = nullptr;
    int rc;

    // Enable foreign key constraints
    rc = sqlite3_exec(db_, "PRAGMA foreign_keys = ON;", nullptr, nullptr, &errMsg);
    if (rc != SQLITE_OK) {
        dprintf(D_ERROR, "Failed to enable foreign keys: %s\n", errMsg ? errMsg : "Unknown");
        sqlite3_free(errMsg);
        return false;
    }

    int version = getSchemaVersion();
    if (version == -1) {
        dprintf(D_ERROR, "Failed to get current schema version.\n");
        return false;
    }

    dprintf(D_ALWAYS, "Database Schema Version: %d\n", version);

    // Begin transaction for atomic schema creation
    rc = sqlite3_exec(db_, "BEGIN IMMEDIATE TRANSACTION;", nullptr, nullptr, &errMsg);
    if (rc != SQLITE_OK) {
        dprintf(D_ERROR, "Failed to begin transaction: %s\n", errMsg ? errMsg : "Unknown");
        sqlite3_free(errMsg);
        return false;
    }

    // Execute the schema SQL
    rc = sqlite3_exec(db_, SavedQueries::SCHEMA_SQL.c_str(), nullptr, nullptr, &errMsg);
    if (rc != SQLITE_OK) {
        dprintf(D_ERROR, "Failed to initialize DB schema: %s\n", errMsg ? errMsg : "Unknown");
        sqlite3_free(errMsg);
        ROLLBACK_AND_RETURN();
    }

    // Switch to run version modification statments of database
    // Note: Use case fall throughs adding newer versions lower
    //       Also, don't forget to change the new version in the database!!!
    constexpr int SCHEMA_VERSION = 1;

    // If the detected database version is newer then what we can handle fail
    // since we are likely to be incompatable and fail later or muck up the database
    if (version > SCHEMA_VERSION) {
        dprintf(D_ALWAYS, "Database schema version (%d) is newer then my version (%d).\n", version, SCHEMA_VERSION);
        ROLLBACK_AND_RETURN();
    } else if (version < SCHEMA_VERSION) { // Database schema is older so update databse appropriately
        std::string version_stmt;
        formatstr(version_stmt, "PRAGMA user_version = %d;", SCHEMA_VERSION);

        switch (version) {
            case 0: // Initializing new database
                [[fallthrough]];
            default:
                if (sqlite3_exec(db_, version_stmt.c_str(), nullptr, nullptr, &errMsg) != SQLITE_OK) {
                    dprintf(D_ERROR, "Failed to set new schema version in database: %s\n", errMsg ? errMsg : "Unknown");
                    sqlite3_free(errMsg);
                    ROLLBACK_AND_RETURN();
                }
        }
    }

    // Commit the transaction
    rc = sqlite3_exec(db_, "COMMIT;", nullptr, nullptr, &errMsg);
    if (rc != SQLITE_OK) {
        dprintf(D_ERROR, "Failed to commit schema transaction: %s\n", errMsg ? errMsg : "Unknown");
        sqlite3_free(errMsg);
        ROLLBACK_AND_RETURN();
    }

    // TODO: Make these a map(name:stmt)?
    const char* jobLookupSQL = "SELECT JobId, JobListId FROM Jobs WHERE ClusterId = ? AND ProcId = ?";
    if (sqlite3_prepare_v2(db_, jobLookupSQL, -1, &jobIdLookupStmt_, nullptr) != SQLITE_OK) {
        dprintf(D_ERROR, "Failed to prepare jobIdLookupStmt_: %s \n ", sqlite3_errmsg(db_));
        jobIdLookupStmt_ = nullptr;
    }

    const char* userInsertSQL = "INSERT OR IGNORE INTO Users (UserName) VALUES (?)";
    if (sqlite3_prepare_v2(db_, userInsertSQL, -1, &userInsertStmt_, nullptr) != SQLITE_OK) {
        dprintf(D_ERROR, "Failed to prepare userInsertStmt: %s\n", sqlite3_errmsg(db_));
        userInsertStmt_ = nullptr;
    }

    const char* userLookupSQL = "SELECT UserId FROM Users WHERE UserName = ?";
    if (sqlite3_prepare_v2(db_, userLookupSQL, -1, &userSelectStmt_, nullptr) != SQLITE_OK) {
        dprintf(D_ERROR, "Failed to prepare userSelectStmt: %s\n", sqlite3_errmsg(db_));
        userSelectStmt_ = nullptr;
    }

    const char* jobListInsertSQL = "INSERT OR IGNORE INTO JobLists (ClusterId, UserId) VALUES (?, ?)";
    if (sqlite3_prepare_v2(db_, jobListInsertSQL, -1, &jobListInsertStmt_, nullptr) != SQLITE_OK) {
        dprintf(D_ERROR, "Failed to prepare jobListInsertStmt: %s\n", sqlite3_errmsg(db_));
        jobListInsertStmt_ = nullptr;
    }

    const char* jobListLookupSQL = "SELECT JobListId FROM JobLists WHERE ClusterId = ? AND UserId = ?";
    if (sqlite3_prepare_v2(db_, jobListLookupSQL, -1, &jobListSelectStmt_, nullptr) != SQLITE_OK) {
        dprintf(D_ERROR, "Failed to prepare jobListSelectStmt: %s\n", sqlite3_errmsg(db_));
        jobListSelectStmt_ = nullptr;
    }

    const char* jobInsertSQL = "INSERT OR IGNORE INTO Jobs (ClusterId, ProcId, UserId, JobListId, TimeOfCreation) VALUES (?, ?, ?, ?, ?)";
    if (sqlite3_prepare_v2(db_, jobInsertSQL, -1, &jobInsertStmt_, nullptr) != SQLITE_OK) {
        dprintf(D_ERROR, "Failed to prepare jobInsertStmt: %s\n", sqlite3_errmsg(db_));
        jobInsertStmt_ = nullptr;
    }

    const char* jobSelectSQL = "SELECT JobId FROM Jobs WHERE ClusterId = ? AND ProcId = ?";
    if (sqlite3_prepare_v2(db_, jobSelectSQL, -1, &jobSelectStmt_, nullptr) != SQLITE_OK) {
        dprintf(D_ERROR, "Failed to prepare jobSelectStmt: %s\n", sqlite3_errmsg(db_));
        jobSelectStmt_ = nullptr;
    }

    return true;
}

DBHandler::~DBHandler() {
    if (jobIdLookupStmt_) {
        sqlite3_finalize(jobIdLookupStmt_);
        jobIdLookupStmt_ = nullptr;
    }

    if (userInsertStmt_) {
        sqlite3_finalize(userInsertStmt_);
        userInsertStmt_ = nullptr;
    }   

    if (userSelectStmt_) {
        sqlite3_finalize(userSelectStmt_);
        userSelectStmt_ = nullptr;
    }

    if (jobListInsertStmt_) {
        sqlite3_finalize(jobListInsertStmt_);
        jobListInsertStmt_ = nullptr;
    }

    if (jobListSelectStmt_) {
        sqlite3_finalize(jobListSelectStmt_);
        jobListSelectStmt_ = nullptr;
    }

    if (jobInsertStmt_) {
        sqlite3_finalize(jobInsertStmt_);
        jobInsertStmt_ = nullptr;
    }

    if (jobSelectStmt_) {
        sqlite3_finalize(jobSelectStmt_);
        jobSelectStmt_ = nullptr;
    }

    if (db_) {
        int rc = sqlite3_close(db_);
        if (rc != SQLITE_OK) {
            dprintf(D_ERROR, "Failed to close db: %s \n", sqlite3_errmsg(db_));
        }
        db_ = nullptr;
    }
}


// For Librarian to check whether DBHandler has been initialized and works properly upon startup
// Check if the database connection is working properly
bool DBHandler::testDatabaseConnection() {
    // 1. Check database handle exists
    if (!db_) {
        dprintf(D_ERROR, "Database handle is null\n");
        return false;
    }

    // 2. Test basic connectivity
    int result = sqlite3_exec(db_, "SELECT 1;", nullptr, nullptr, nullptr);
    if (result != SQLITE_OK) {
        dprintf(D_ERROR, "Basic query failed: %s\n", sqlite3_errmsg(db_));
        return false;
    }

    return true;
} 

// Verify that the database schema matches expected structure
bool DBHandler::verifyDatabaseSchema(const std::string& schemaSQL) {
    // Count CREATE TABLE statements in the schema SQL
    std::string searchStr = "CREATE TABLE";
    int expectedTableCount = 0;
    size_t pos = 0;

    while ((pos = schemaSQL.find(searchStr, pos)) != std::string::npos) {
        expectedTableCount++;
        pos += searchStr.length();
    }

    if (expectedTableCount == 0) {
        dprintf(D_ERROR, "No CREATE TABLE statements found in schema SQL\n");
        return false;
    }

    // Count actual tables in the database
    const char* countTablesQuery = R"(
        SELECT COUNT(*) FROM sqlite_master 
        WHERE type='table' AND name NOT LIKE 'sqlite_%';
    )";

    sqlite3_stmt* stmt = nullptr;
    int result = sqlite3_prepare_v2(db_, countTablesQuery, -1, &stmt, nullptr);
    if (result != SQLITE_OK) {
        dprintf(D_ERROR, "Table count prepare failed: %s\n",
                sqlite3_errmsg(db_));
        return false;
    }

    int actualTableCount = 0;
    if (sqlite3_step(stmt) == SQLITE_ROW) {
        actualTableCount = sqlite3_column_int(stmt, 0);
    }
    sqlite3_finalize(stmt);

    if (actualTableCount != expectedTableCount) {
        dprintf(D_ERROR, "Table count mismatch. Expected: %d, Found: %d\n",
                expectedTableCount, actualTableCount);
        return false;
    }

    return true;
}


int DBHandler::getSchemaVersion() {
    if ( ! db_) {
        dprintf(D_ERROR, "Error: Database connection is not created.\n");
        return -1;
    }

    sqlite3_stmt* stmt = nullptr;
    if (sqlite3_prepare_v2(db_, "PRAGMA user_version;", -1, &stmt, nullptr) != SQLITE_OK) {
        dprintf(D_ERROR, "Error: Failed to prepare version query: %s\n", sqlite3_errmsg(db_));
        return -1;
    }

    int version = 0;
    if (sqlite3_step(stmt) == SQLITE_ROW) {
        version = sqlite3_column_int(stmt, 0);
    }

    sqlite3_finalize(stmt);
    return version;
}


// -------------------------
// JobRecord Insertion
// -------------------------

// Checks cache for {clusterId, procID} pair before checking database
// Returns {-1, -1} if there was an error, 
// or if it was not in the cache or the database
std::pair<int, int> DBHandler::jobIdLookup(int clusterId, int procId) {
    std::pair<int, int> key = {clusterId, procId};

    // Step 1: Try cache
    if (jobIdCache_.contains(key)) {
        return jobIdCache_[key]; 
    }

    // Step 2: Fallback to DB query
    if (!jobIdLookupStmt_) {
        dprintf(D_ERROR, "No prepared job id lookup statement\n");
        return {-1, -1};
    }

    sqlite3_reset(jobIdLookupStmt_);
    sqlite3_clear_bindings(jobIdLookupStmt_);
    sqlite3_bind_int(jobIdLookupStmt_, 1, clusterId);
    sqlite3_bind_int(jobIdLookupStmt_, 2, procId);

    // If found in db and not cache, insert into cache for future entries 
    int rc = sqlite3_step(jobIdLookupStmt_);
    if (rc == SQLITE_ROW) {
        int jobId = sqlite3_column_int(jobIdLookupStmt_, 0);
        int jobListId = sqlite3_column_int(jobIdLookupStmt_, 1);
        std::pair<int, int> value = {jobId, jobListId};
        jobIdCache_.insert(key, value);
        return {jobId, jobListId};
    }

    if (rc != SQLITE_DONE) {
        dprintf(D_ERROR, "sqlite3_step failed for ClusterId= %d, ProcId= %d - %s \n",
        clusterId, procId, sqlite3_errmsg(db_));
    }

    return {-1, -1};
}


int getOrInsertIdPrepared(
    sqlite3_stmt* insertStmt,
    sqlite3_stmt* selectStmt,
    std::function<void(sqlite3_stmt*)> bindInsert,
    std::function<void(sqlite3_stmt*)> bindSelect,
    std::function<int(sqlite3_stmt*)> extractId
) {
    int id = -1;

    // STEP 1: Try SELECT first
    bindSelect(selectStmt);
    int rc = sqlite3_step(selectStmt);
    if (rc == SQLITE_ROW) {
        // Found existing record
        id = extractId(selectStmt);
        sqlite3_reset(selectStmt);
        sqlite3_clear_bindings(selectStmt);
        return id;
    } else if (rc != SQLITE_DONE) {
        // SELECT failed for some reason other than "no rows found"
        dprintf(D_ERROR, "Select failed: %s\n", sqlite3_errmsg(sqlite3_db_handle(selectStmt)));
        sqlite3_reset(selectStmt);
        sqlite3_clear_bindings(selectStmt);
        return -1;
    }

    // Clean up SELECT statement
    sqlite3_reset(selectStmt);
    sqlite3_clear_bindings(selectStmt);

    // STEP 2: Record doesn't exist, try INSERT
    bindInsert(insertStmt);
    rc = sqlite3_step(insertStmt);

    if (rc == SQLITE_DONE) {
        // Insert succeeded
        id = static_cast<int>(sqlite3_last_insert_rowid(sqlite3_db_handle(insertStmt)));
    } else if (rc == SQLITE_CONSTRAINT) {
        // Race condition: someone else inserted between our SELECT and INSERT
        // Do SELECT again to get the ID
        sqlite3_reset(insertStmt);
        sqlite3_clear_bindings(insertStmt);

        bindSelect(selectStmt);
        rc = sqlite3_step(selectStmt);
        if (rc == SQLITE_ROW) {
            id = extractId(selectStmt);
        } else {
            dprintf(D_ERROR, "Race condition SELECT fallback failed: %s\n",
                    sqlite3_errmsg(sqlite3_db_handle(selectStmt)));
        }
        sqlite3_reset(selectStmt);
        sqlite3_clear_bindings(selectStmt);
    } else {
        // Some other insert error
        dprintf(D_ERROR, "Insert failed: %s\n", sqlite3_errmsg(sqlite3_db_handle(insertStmt)));
    }

    sqlite3_reset(insertStmt);
    sqlite3_clear_bindings(insertStmt);

    return id;
}

/**
 * Insert an unseen job and its associated metadata into the database and cache
 * Inserts entries into 'User', 'JobLists', and 'Job' tables in database
 */
bool DBHandler::insertUnseenJob(const std::string& owner, int clusterId, int procId, int64_t timeOfCreation) {
    // 1. Insert User and get userId
    int userId = getOrInsertIdPrepared(
        userInsertStmt_,
        userSelectStmt_,
        [&](sqlite3_stmt* stmt) {
            sqlite3_bind_text(stmt, 1, owner.c_str(), -1, SQLITE_TRANSIENT);
        },
        [&](sqlite3_stmt* stmt) {
            sqlite3_bind_text(stmt, 1, owner.c_str(), -1, SQLITE_TRANSIENT);
        },
        [](sqlite3_stmt* stmt) {
            return sqlite3_column_int(stmt, 0);
        }
    );
    if(userId == -1) {
        dprintf(D_ERROR, "Failed to insert/lookup user: %s\n", owner.c_str());
        return false;
    }

    // 2. Insert JobList and get jobListId
    int jobListId = getOrInsertIdPrepared(
        jobListInsertStmt_,
        jobListSelectStmt_,
        [&](sqlite3_stmt* stmt) {
            sqlite3_bind_int(stmt, 1, clusterId);
            sqlite3_bind_int(stmt, 2, userId);
        },
        [&](sqlite3_stmt* stmt) {
            sqlite3_bind_int(stmt, 1, clusterId);
            sqlite3_bind_int(stmt, 2, userId);
        },
        [](sqlite3_stmt* stmt) {
            return sqlite3_column_int(stmt, 0);
        }
    );
    if(jobListId == -1) {
        dprintf(D_ERROR, "Failed to insert/lookup jobListId for {cluster, user} : {%d, %d}\n", clusterId, userId);
        return false;
    }

    // 3. Insert Job and get jobId
    int jobId = getOrInsertIdPrepared(
        jobInsertStmt_,
        jobSelectStmt_,
        [&](sqlite3_stmt* stmt) {
            sqlite3_bind_int(stmt, 1, clusterId);
            sqlite3_bind_int(stmt, 2, procId);
            sqlite3_bind_int(stmt, 3, userId);
            sqlite3_bind_int(stmt, 4, jobListId);
            sqlite3_bind_int64(stmt, 5, timeOfCreation);  // assume column for time
        },
        [&](sqlite3_stmt* stmt) {
            sqlite3_bind_int(stmt, 1, clusterId);
            sqlite3_bind_int(stmt, 2, procId);
        },
        [](sqlite3_stmt* stmt) {
            return sqlite3_column_int(stmt, 0);
        }
    );
    if(jobId == -1) {
        dprintf(D_ERROR, "Failed to insert/lookup jobId for {cluster, proc} : {%d, %d}\n", clusterId, procId);
        return false;
    }

    // 4. Cache the result
    std::pair<int, int> key = {clusterId, procId};
    std::pair<int, int> value = {jobId, jobListId};
    jobIdCache_.insert(key, value);
    return true;
}


/**
 * Batch insert multiple JobRecords in a single transaction.
 */
bool DBHandler::batchInsertJobRecords(const std::vector<JobRecord>& jobs) {
    if (jobs.empty()) {
        dprintf(D_FULLDEBUG, "No records to batch insert\n");
        return false;
    }

    const char* sql = R"(
        INSERT INTO JobRecords (Offset, CompletionDate, JobId, FileId, JobListId)
        VALUES (?, ?, ?, ?, ?);
    )";

    sqlite3_stmt* stmt = nullptr;
    int rc = sqlite3_prepare_v2(db_, sql, -1, &stmt, nullptr);
    if (rc != SQLITE_OK) {
        dprintf(D_ERROR, "Failed to prepare JobRecords insert: %s\n", sqlite3_errmsg(db_));
        return false;
    }

    int jobId, jobListId;

    for (const auto& job : jobs) {

        std::tie(jobId, jobListId) = jobIdLookup(job.ClusterId, job.ProcId);

        if (jobId == -1) { 
            bool insertSuccess = insertUnseenJob(job.Owner, job.ClusterId, job.ProcId, job.CompletionDate);
            if(!insertSuccess){
                dprintf(D_ERROR, "Failed to create needed job information tables for %d.%d\n", job.ClusterId, job.ProcId);
                return false;
            }
            
            // Lookup freshly inserted info
            std::tie(jobId, jobListId) = jobIdLookup(job.ClusterId, job.ProcId);
        }

        sqlite3_bind_int64(stmt, 1, job.Offset);
        sqlite3_bind_int64(stmt, 2, job.CompletionDate);
        sqlite3_bind_int(stmt, 3, jobId);
        sqlite3_bind_int(stmt, 4, job.FileId);
        sqlite3_bind_int(stmt, 5, jobListId);

        rc = sqlite3_step(stmt);
        if (rc != SQLITE_DONE) {
            dprintf(D_ERROR, "Failed to insert JobRecord for ClusterId=%d, ProcId=%d: %s\n",
                job.ClusterId, job.ProcId, sqlite3_errmsg(db_)); 
        }

        sqlite3_reset(stmt);
        sqlite3_clear_bindings(stmt);
    }

    sqlite3_finalize(stmt);
    return true;
}

// To convert DateOfRotation and DateOfDeletion strings to timestamps for database insert
int64_t convertRotationStringToTimestamp(const std::string& rotationStr) {
    if (rotationStr.empty()) return 0;

    // Parse "20241215T143022" format into a timestamp
    std::tm tm = {};
    std::istringstream ss(rotationStr);
    ss >> std::get_time(&tm, "%Y%m%dT%H%M%S");

    return static_cast<int64_t>(std::mktime(&tm));
}


/**
 * Insert job records and update file processing status atomically
 * This wrapper ensures both job records insertion and file info update
 * happen in a single transaction
 */
bool DBHandler::insertJobFileRecords(const std::vector<JobRecord>& jobs, const FileInfo& fileInfo) {

    // Begin outer transaction for atomicity
    int result = sqlite3_exec(db_, "BEGIN TRANSACTION;", nullptr, nullptr, nullptr);
    if (result != SQLITE_OK) {
        dprintf(D_ERROR, "Failed to begin transaction for job file records insertion\n");
        return false;
    }

    // Insert job records using existing function (now transaction-less)
    bool jobInsertSuccess = batchInsertJobRecords(jobs);
    if (!jobInsertSuccess) return false;

    // Update the Files table - always update LastOffset and FullyRead, conditionally update DateOfRotation
    std::string updateFilesSql;
    if (!fileInfo.DateOfRotation.empty()) {
        updateFilesSql = "UPDATE Files SET LastOffset = ?, DateOfRotation = ?, FullyRead = ? WHERE FileId = ?;";
    } else {
        updateFilesSql = "UPDATE Files SET LastOffset = ?, FullyRead = ? WHERE FileId = ?;";
    }

    sqlite3_stmt* stmt = nullptr;
    result = sqlite3_prepare_v2(db_, updateFilesSql.c_str(), -1, &stmt, nullptr);
    if (result != SQLITE_OK) {
        dprintf(D_ERROR, "Failed to prepare Files update statement: %s\n", sqlite3_errmsg(db_));
        ROLLBACK_AND_RETURN();
    }

    // Bind parameters
    sqlite3_bind_int64(stmt, 1, static_cast<sqlite3_int64>(fileInfo.LastOffset));

    if (!fileInfo.DateOfRotation.empty()) {
        int64_t rotationTimestamp = convertRotationStringToTimestamp(fileInfo.DateOfRotation);
        sqlite3_bind_int64(stmt, 2, rotationTimestamp);
        sqlite3_bind_int(stmt, 3, fileInfo.FullyRead ? 1 : 0);
        sqlite3_bind_int(stmt, 4, fileInfo.FileId);
    } else {
        sqlite3_bind_int(stmt, 2, fileInfo.FullyRead ? 1 : 0);
        sqlite3_bind_int(stmt, 3, fileInfo.FileId);
    }

    result = sqlite3_step(stmt);
    sqlite3_finalize(stmt);

    if (result != SQLITE_DONE) {
        dprintf(D_ERROR, "Failed to update Files table for FileId %ld: %s\n",
               fileInfo.FileId, sqlite3_errmsg(db_));
        ROLLBACK_AND_RETURN();
    }

    // Commit the transaction
    result = sqlite3_exec(db_, "COMMIT;", nullptr, nullptr, nullptr);
    if (result != SQLITE_OK) {
        dprintf(D_ERROR, "Failed to commit transaction: %s\n", sqlite3_errmsg(db_));
        ROLLBACK_AND_RETURN();
    }

    return true;
}

// -------------------------
// File Info Insertion / Update
// ------------------------

/**
 * Insert or fetch metadata about a file (FileId and LastOffset)
 * Populate data into FileInfo struct 
 */
void DBHandler::writeFileInfo(FileInfo &info) {
    const char *insertSql = R"(
        INSERT OR IGNORE INTO Files(FileName, FileInode, FileHash, LastOffset)
        VALUES (?, ?, ?, ?)
    )";

    sqlite3_stmt *insertStmt = nullptr;

    if (sqlite3_prepare_v2(db_, insertSql, -1, &insertStmt, nullptr) != SQLITE_OK) {
        dprintf(D_ERROR, "Failed to prepare insert statement: %s\n", sqlite3_errmsg(db_));
        return;
    }

    sqlite3_bind_text(insertStmt, 1, info.FileName.c_str(), -1, SQLITE_TRANSIENT);
    sqlite3_bind_int64(insertStmt, 2, info.FileInode);
    sqlite3_bind_text(insertStmt, 3, info.FileHash.c_str(), -1, SQLITE_TRANSIENT);
    sqlite3_bind_int64(insertStmt, 4, info.LastOffset);

    if (sqlite3_step(insertStmt) != SQLITE_DONE) {
        dprintf(D_ERROR, "Insert (or ignore) failed: %s\n", sqlite3_errmsg(db_));
        sqlite3_finalize(insertStmt);
        return;
    }

    sqlite3_finalize(insertStmt);

    // Now fetch FileId and LastOffset (existing or newly inserted)
    const char *selectSql = R"(
        SELECT FileId, LastOffset FROM Files WHERE FileInode = ? AND FileHash = ?
    )";

    sqlite3_stmt *selectStmt = nullptr;
    if (sqlite3_prepare_v2(db_, selectSql, -1, &selectStmt, nullptr) != SQLITE_OK) {
        dprintf(D_ERROR, "Failed to prepare select statement: %s\n", sqlite3_errmsg(db_));
        return;
    }

    sqlite3_bind_int64(selectStmt, 1, info.FileInode);
    sqlite3_bind_text(selectStmt, 2, info.FileHash.c_str(), -1, SQLITE_TRANSIENT);

    if (sqlite3_step(selectStmt) == SQLITE_ROW) {
        info.FileId = sqlite3_column_int(selectStmt, 0);
        info.LastOffset = sqlite3_column_int64(selectStmt, 1);
    } else {
        dprintf(D_ERROR, "Failed to fetch FileId/LastOffset after insert/ignore.\n");
    }

    sqlite3_finalize(selectStmt);
}

/**
 * Update LastOffset for the history and epoch files after parsing.
 * Change offset to match progress in reading Epoch + History files
 */
void DBHandler::updateFileInfo(FileInfo historyFile) {
    const char* updateSQL = R"(
        UPDATE Files
        SET LastOffset = ?
        WHERE FileName = ?
    )";

    sqlite3_stmt* stmt = nullptr;
    if (sqlite3_prepare_v2(db_, updateSQL, -1, &stmt, nullptr) != SQLITE_OK) {
        dprintf(D_ERROR, "Failed to prepare update statement: %s\n", sqlite3_errmsg(db_));
        return;
    }

    sqlite3_bind_int64(stmt, 1, historyFile.LastOffset);
    sqlite3_bind_text(stmt, 2, historyFile.FileName.c_str(), -1, SQLITE_TRANSIENT);

    if (sqlite3_step(stmt) != SQLITE_DONE) {
        dprintf(D_ERROR, "Failed to update offset for file %s: %s\n", historyFile.FileName.c_str(), sqlite3_errmsg(db_));
    }

    sqlite3_finalize(stmt);
}


// -------------------------
// Status Tracking
// -------------------------

/**
 * Insert a status snapshot into the 'Status' table and update the single row in 'StatusData' table
 * Includes Status but also running averages of various performance stats
 */
/**
 * Insert a status snapshot into the 'Status' table and update the single row in 'StatusData' table
 * Includes Status but also running averages of various performance stats
 */
bool DBHandler::writeStatusAndData(const Status& status, const StatusData& statusData) {
    const char* statusSql = R"(
        INSERT INTO Status (
            TimeOfUpdate,
            FileIdLastRead,
            FileOffsetLastRead,
            TotalRecordsRead,
            DurationMs,
            JobBacklogEstimate,
            HitMaxIngestLimit,
            GarbageCollectionRun
        ) VALUES (?, ?, ?, ?, ?, ?, ?, ?);
    )";

    // Use INSERT OR REPLACE to handle missing row - will insert if doesn't exist, update if it does
    const char* statusDataSqlArrivalHz = R"(
        INSERT OR REPLACE INTO StatusData (
            StatusDataId,
            AvgAdsIngestedPerCycle,
            AvgIngestDurationMs,
            MeanIngestHz,
            MeanArrivalHz,
            MeanBacklogEstimate,
            TotalCycles,
            TotalAdsIngested,
            HitMaxIngestLimitRate,
            LastRunLeftBacklog,
            TimeOfLastUpdate
        ) VALUES (1, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?);
    )";

    // Start transaction
    char* errMsg = nullptr;
    if (sqlite3_exec(db_, "BEGIN TRANSACTION;", nullptr, nullptr, &errMsg) != SQLITE_OK) {
        dprintf(D_ERROR, "Failed to begin transaction: %s\n", errMsg);
        sqlite3_free(errMsg);
        return false;
    }

    sqlite3_stmt* stmt = nullptr;
    bool success = false;

    // Insert into Status
    if (sqlite3_prepare_v2(db_, statusSql, -1, &stmt, nullptr) == SQLITE_OK) {
        sqlite3_bind_int64(stmt, 1, status.TimeOfUpdate);
        sqlite3_bind_int(stmt, 2, status.HistoryFileIdLastRead);
        sqlite3_bind_int64(stmt, 3, status.HistoryFileOffsetLastRead);
        sqlite3_bind_int(stmt, 4, status.TotalJobsRead);
        sqlite3_bind_int(stmt, 5, status.DurationMs);
        sqlite3_bind_int(stmt, 6, status.JobBacklogEstimate);
        sqlite3_bind_int(stmt, 7, status.HitMaxIngestLimit ? 1 : 0);
        sqlite3_bind_int(stmt, 8, status.GarbageCollectionRun ? 1 : 0);

        if (sqlite3_step(stmt) != SQLITE_DONE) {
            dprintf(D_ERROR, "Failed to insert into Status: %s\n", sqlite3_errmsg(db_));
        } else {
            success = true;
        }
    } else {
        dprintf(D_ERROR, "Failed to prepare Status insert: %s\n", sqlite3_errmsg(db_));
    }
    sqlite3_finalize(stmt);

    // Insert/Update StatusData (if Status insert succeeded)
    if (success) {
        const char* sql = statusDataSqlArrivalHz;

        if (sqlite3_prepare_v2(db_, sql, -1, &stmt, nullptr) == SQLITE_OK) {
            int idx = 1;
            sqlite3_bind_double(stmt, idx++, statusData.AvgAdsIngestedPerCycle);
            sqlite3_bind_double(stmt, idx++, statusData.AvgIngestDurationMs);
            sqlite3_bind_double(stmt, idx++, statusData.MeanIngestHz);
            sqlite3_bind_double(stmt, idx++, statusData.MeanArrivalHz);
            sqlite3_bind_double(stmt, idx++, statusData.MeanBacklogEstimate);
            sqlite3_bind_int64(stmt, idx++, statusData.TotalCycles);
            sqlite3_bind_int64(stmt, idx++, statusData.TotalAdsIngested);
            sqlite3_bind_double(stmt, idx++, statusData.HitMaxIngestLimitRate);
            sqlite3_bind_int(stmt, idx++, statusData.LastRunLeftBacklog ? 1 : 0);
            sqlite3_bind_int64(stmt, idx++, statusData.TimeOfLastUpdate);

            if (sqlite3_step(stmt) != SQLITE_DONE) {
                dprintf(D_ERROR, "Failed to insert/update StatusData: %s\n", sqlite3_errmsg(db_));
                success = false;
            }
        } else {
            dprintf(D_ERROR, "Failed to prepare StatusData insert/update: %s\n", sqlite3_errmsg(db_));
            success = false;
        }
        sqlite3_finalize(stmt);
    }

    // Finalize transaction
    if (success) {
        if (sqlite3_exec(db_, "COMMIT;", nullptr, nullptr, &errMsg) != SQLITE_OK) {
            dprintf(D_ERROR, "Failed to commit transaction: %s\n", errMsg);
            sqlite3_free(errMsg);
            return false;
        }
        return true;
    } else {
        ROLLBACK_AND_RETURN();
    }
}

/**
 * Recovers the librarian's local cache information from the database if the process
 * caches are empty or uninitialized (e.g., after a restart or crash). This function
 * repopulates file tracking state, status data, and reading positions to resume
 * processing from where it left off.
 * NOTE: Currently this function unconditionally runs recovery if database is not empty,
 * but this behavior might change after the system is daemonized
 */
bool DBHandler::maybeRecoverStatusAndFiles(FileSet& historyFileSet_, StatusData& statusData_)
{
    // Check if database is empty -> if so, no recovery needed, this is the first run
    const char* checkDatabaseEmptySql = "SELECT COUNT(*) FROM (SELECT 1 FROM Files LIMIT 1);";
    sqlite3_stmt* stmt = nullptr;

    if (sqlite3_prepare_v2(db_, checkDatabaseEmptySql, -1, &stmt, nullptr) != SQLITE_OK) {
        dprintf(D_ERROR, "Database was empty, starting first startup protocol\n");
        return false;
    }

    sqlite3_step(stmt);
    int count = sqlite3_column_int(stmt, 0);
    sqlite3_finalize(stmt);

    if (count == 0) {
        dprintf(D_ERROR, "No existing data found, skipping recovery\n");
        return false;  // No data in database, this is the first startup
    }

	char* errMsg = nullptr;
	int rc = sqlite3_exec(db_, "BEGIN;", nullptr, nullptr, &errMsg);
	if (rc != SQLITE_OK) {
		dprintf(D_ERROR, "Failed to begin transaction: %s\n", errMsg ? errMsg : "unknown");
		if (errMsg) sqlite3_free(errMsg);
		return false;
	}

	// 1. Recover StatusData - now reads from the single row
	const std::string statusDataSql =
		"SELECT AvgAdsIngestedPerCycle, AvgIngestDurationMs, MeanIngestHz, MeanArrivalHz, "
		"MeanBacklogEstimate, TotalCycles, TotalAdsIngested, HitMaxIngestLimitRate, LastRunLeftBacklog, "
		"TimeOfLastUpdate FROM StatusData WHERE StatusDataId = 1;";

	sqlite3_stmt* stmtStatusData = nullptr;
	rc = sqlite3_prepare_v2(db_, statusDataSql.c_str(), -1, &stmtStatusData, nullptr);
	if (rc != SQLITE_OK) {
		dprintf(D_ERROR, "Failed to prepare StatusData query\n");
		ROLLBACK_AND_RETURN();
	}
	rc = sqlite3_step(stmtStatusData);
	if (rc == SQLITE_ROW) {
		statusData_.AvgAdsIngestedPerCycle = sqlite3_column_double(stmtStatusData, 0);
		statusData_.AvgIngestDurationMs = sqlite3_column_double(stmtStatusData, 1);
		statusData_.MeanIngestHz = sqlite3_column_double(stmtStatusData, 2);
		statusData_.MeanArrivalHz = sqlite3_column_double(stmtStatusData, 3);
		statusData_.MeanBacklogEstimate = sqlite3_column_double(stmtStatusData, 4);
		statusData_.TotalCycles = sqlite3_column_int64(stmtStatusData, 5);
		statusData_.TotalAdsIngested = sqlite3_column_int64(stmtStatusData, 6);
		statusData_.HitMaxIngestLimitRate = sqlite3_column_double(stmtStatusData, 7);
		statusData_.LastRunLeftBacklog = sqlite3_column_int(stmtStatusData, 8) != 0;
		statusData_.TimeOfLastUpdate = sqlite3_column_int64(stmtStatusData, 9);
	}
	sqlite3_finalize(stmtStatusData);

	// 2. Recover Status
	const std::string statusSql =
		"SELECT TimeOfUpdate, FileIdLastRead, FileOffsetLastRead FROM Status ORDER BY TimeOfUpdate DESC LIMIT 1;";

	sqlite3_stmt* stmtStatus = nullptr;
	rc = sqlite3_prepare_v2(db_, statusSql.c_str(), -1, &stmtStatus, nullptr);
	if (rc != SQLITE_OK) {
		dprintf(D_ERROR, "Failed to prepare Status query\n");
		ROLLBACK_AND_RETURN();
	}
	rc = sqlite3_step(stmtStatus);
	if (rc == SQLITE_ROW) {
		int64_t timeOfUpdate = sqlite3_column_int64(stmtStatus, 0);
		int historyFileId = sqlite3_column_int(stmtStatus, 1);
		std::ignore = sqlite3_column_int64(stmtStatus, 2);

		statusData_.TimeOfLastUpdate = timeOfUpdate;

		historyFileSet_.lastFileReadId = historyFileId;
		historyFileSet_.lastStatusTime = timeOfUpdate;
	}
	sqlite3_finalize(stmtStatus);

	// 3. Recover Files
	const std::string filesSql =
		"SELECT FileId, FileName, FileInode, FileHash, LastOffset FROM Files WHERE DateOfRotation IS NULL;";

	sqlite3_stmt* stmtFiles = nullptr;
	rc = sqlite3_prepare_v2(db_, filesSql.c_str(), -1, &stmtFiles, nullptr);
	if (rc != SQLITE_OK) {
		dprintf(D_ERROR, "Failed to prepare Files query\n");
		ROLLBACK_AND_RETURN();
	}
	while ((rc = sqlite3_step(stmtFiles)) == SQLITE_ROW) {
		FileInfo fileInfo{};
		fileInfo.FileId = sqlite3_column_int64(stmtFiles, 0);
		fileInfo.FileName = reinterpret_cast<const char*>(sqlite3_column_text(stmtFiles, 1));
		fileInfo.FileInode = sqlite3_column_int64(stmtFiles, 2);
		fileInfo.FileHash = reinterpret_cast<const char*>(sqlite3_column_text(stmtFiles, 3));
		fileInfo.LastOffset = sqlite3_column_int64(stmtFiles, 4);
		fileInfo.FileSize = -1;      // No size in table
		fileInfo.LastModified = 0;   // Not available

		if (fileInfo.FileName.rfind(historyFileSet_.GetFileName(), 0) == 0) {
			historyFileSet_.fileMap[fileInfo.FileId] = fileInfo;
		}
	}
	sqlite3_finalize(stmtFiles);

	rc = sqlite3_exec(db_, "COMMIT;", nullptr, nullptr, &errMsg);
	if (rc != SQLITE_OK) {
		dprintf(D_ERROR, "Failed to commit transaction: %s\n", errMsg ? errMsg : "unknown");
		if (errMsg) sqlite3_free(errMsg);
		// Try rollback if commit failed
		ROLLBACK_AND_RETURN();
	}

	return true;
}

// -------------------------
// ArchiveChange Handlers
// -------------------------

bool DBHandler::insertNewFilesAndMarkOldOnes(ArchiveChange& fileSetChange) {

    std::pair<int, std::string>& rotatedFile = fileSetChange.rotatedFile;
    std::vector<FileInfo>& newFiles = fileSetChange.newFiles;
    std::vector<int>& deletedFileIds = fileSetChange.deletedFileIds;

    sqlite3_exec(db_, "BEGIN TRANSACTION;", nullptr, nullptr, nullptr);

    bool success = true;

    // 1. Handle rotated file
    if (!rotatedFile.second.empty()) {
        int oldFileId = rotatedFile.first;
        const std::string& newName = rotatedFile.second;

        std::string sql = "UPDATE Files SET FileName = ?, DateOfRotation = ? WHERE FileId = ?";
        sqlite3_stmt* stmt = nullptr;

        if (sqlite3_prepare_v2(db_, sql.c_str(), -1, &stmt, nullptr) != SQLITE_OK) {
            dprintf(D_ERROR, "Failed to prepare update statement: %s\n", sqlite3_errmsg(db_));
            success = false;
        } else {
            sqlite3_bind_text(stmt, 1, newName.c_str(), -1, SQLITE_STATIC);

            // Set DateOfRotation to current UNIX timestamp
            std::time_t now = std::time(nullptr);
            sqlite3_bind_int64(stmt, 2, static_cast<sqlite3_int64>(now));

            sqlite3_bind_int(stmt, 3, oldFileId);

            if (sqlite3_step(stmt) != SQLITE_DONE) {
                dprintf(D_ERROR, "Failed to update rotated file: %s\n", sqlite3_errmsg(db_));
                success = false;
            }

            sqlite3_finalize(stmt);
        }

        if (!success) {
            ROLLBACK_AND_RETURN();
        }
    }

    // 2. Insert new files and record FileIds back into FileInfo objects
    const char* insertSQL = R"(
        INSERT OR IGNORE INTO Files (FileName, FileInode, FileHash, LastOffset)
        VALUES (?, ?, ?, ?)
    )";

    sqlite3_stmt* insertStmt = nullptr;

    if (sqlite3_prepare_v2(db_, insertSQL, -1, &insertStmt, nullptr) != SQLITE_OK) {
        dprintf(D_ERROR, "Failed to prepare insert statement: %s\n", sqlite3_errmsg(db_));
        ROLLBACK_AND_RETURN();
    }

    for (FileInfo& newFile : newFiles) {
        sqlite3_bind_text(insertStmt, 1, newFile.FileName.c_str(), -1, SQLITE_TRANSIENT);
        sqlite3_bind_int64(insertStmt, 2, newFile.FileInode);
        sqlite3_bind_text(insertStmt, 3, newFile.FileHash.c_str(), -1, SQLITE_TRANSIENT);
        sqlite3_bind_int64(insertStmt, 4, newFile.LastOffset);

        if (sqlite3_step(insertStmt) != SQLITE_DONE) {
            dprintf(D_ERROR, "Insert (or ignore) failed: %s\n", sqlite3_errmsg(db_));
            sqlite3_finalize(insertStmt);
            ROLLBACK_AND_RETURN();
        }

        sqlite3_reset(insertStmt);
    }

    sqlite3_finalize(insertStmt);

    // Now fetch FileIds for all inserted/existing files
    const char* selectSQL = R"(
        SELECT FileId FROM Files WHERE FileInode = ? AND FileHash = ?
    )";

    sqlite3_stmt* selectStmt = nullptr;
    if (sqlite3_prepare_v2(db_, selectSQL, -1, &selectStmt, nullptr) != SQLITE_OK) {
        dprintf(D_ERROR, "Failed to prepare select statement: %s\n", sqlite3_errmsg(db_));
        ROLLBACK_AND_RETURN();
    }

    for (FileInfo& newFile : newFiles) {
        sqlite3_bind_int64(selectStmt, 1, newFile.FileInode);
        sqlite3_bind_text(selectStmt, 2, newFile.FileHash.c_str(), -1, SQLITE_TRANSIENT);

        if (sqlite3_step(selectStmt) == SQLITE_ROW) {
            newFile.FileId = sqlite3_column_int(selectStmt, 0);
        } else {
            dprintf(D_ERROR, "Failed to fetch FileId after insert/ignore for file: %s\n", newFile.FileName.c_str());
            sqlite3_finalize(selectStmt);
            ROLLBACK_AND_RETURN();
        }

        sqlite3_reset(selectStmt);
    }

    sqlite3_finalize(selectStmt);

    // 3. Mark deleted files with deletion timestamp
    if (!deletedFileIds.empty()) {
       std::string deleteSQL = "UPDATE Files SET DateOfDeletion = ? WHERE FileId = ?";
       sqlite3_stmt* deleteStmt = nullptr;

       if (sqlite3_prepare_v2(db_, deleteSQL.c_str(), -1, &deleteStmt, nullptr) != SQLITE_OK) {
           dprintf(D_ERROR, "Failed to prepare delete marking statement: %s\n", sqlite3_errmsg(db_));
           ROLLBACK_AND_RETURN();
       }

       std::time_t now = std::time(nullptr);
       
       for (int fileId : deletedFileIds) {
           sqlite3_bind_int64(deleteStmt, 1, static_cast<sqlite3_int64>(now));
           sqlite3_bind_int(deleteStmt, 2, fileId);

           if (sqlite3_step(deleteStmt) != SQLITE_DONE) {
               dprintf(D_ERROR, "Failed to mark file as deleted: %s\n", sqlite3_errmsg(db_));
               sqlite3_finalize(deleteStmt);
               ROLLBACK_AND_RETURN();
           }

           sqlite3_reset(deleteStmt);
       }

       sqlite3_finalize(deleteStmt);
    }


    sqlite3_exec(db_, "COMMIT;", nullptr, nullptr, nullptr);

    dprintf(D_FULLDEBUG, "Inserted %zu new files and marked old ones.\n", newFiles.size());
    return true;
}


// -------------------------
// Garbage Collection
// -------------------------

// NOTE: This function would not work if more of the statements had parameters that need to be bound
// It currently only accepts bindings for the very first statement, which is where we indicate 
// the n number of jobs that we need to delete to get to the low watermark
bool DBHandler::runGarbageCollection(const std::string &gcQuerySQL, int fileLimit) {
    if (!db_) return false;

    char *errMsg = nullptr;

    if (sqlite3_exec(db_, "BEGIN TRANSACTION;", nullptr, nullptr, &errMsg) != SQLITE_OK) {
        dprintf(D_ERROR, "Garbage collection begin transaction failed: %s\n", errMsg);
        sqlite3_free(errMsg);
        return false;
    }

    size_t pos = gcQuerySQL.find(';');
    if (pos == std::string::npos) {
        dprintf(D_ERROR, "Invalid Garbage Collection SQL - no statement terminator\n");
        ROLLBACK_AND_RETURN();
    }

    std::string step1Sql = gcQuerySQL.substr(0, pos + 1);
    std::string step2Sql = gcQuerySQL.substr(pos + 1);

    // Prepare step1, bind ?, execute
    sqlite3_stmt *stmt = nullptr;
    if (sqlite3_prepare_v2(db_, step1Sql.c_str(), -1, &stmt, nullptr) != SQLITE_OK) {
        dprintf(D_ERROR, "Garbage collection prepare step1 failed: %s\n", sqlite3_errmsg(db_));
        ROLLBACK_AND_RETURN();
    }
    if (sqlite3_bind_int(stmt, 1, fileLimit) != SQLITE_OK) {
        dprintf(D_ERROR, "Garbage collection bind failed: %s\n", sqlite3_errmsg(db_));
        sqlite3_finalize(stmt);
        ROLLBACK_AND_RETURN();
    }
    if (sqlite3_step(stmt) != SQLITE_DONE) {
        dprintf(D_ERROR, "Garbage collection execute step1 failed: %s\n", sqlite3_errmsg(db_));
        sqlite3_finalize(stmt);
        ROLLBACK_AND_RETURN();
    }
    sqlite3_finalize(stmt);

    // Exec the rest
    if (sqlite3_exec(db_, step2Sql.c_str(), nullptr, nullptr, &errMsg) != SQLITE_OK) {
        dprintf(D_ERROR, "Garbage collection execute step2 failed: %s\n", errMsg);
        sqlite3_free(errMsg);
        ROLLBACK_AND_RETURN();
    }

    if (sqlite3_exec(db_, "COMMIT;", nullptr, nullptr, &errMsg) != SQLITE_OK) {
        dprintf(D_ERROR, "Garbage collection commit failed: %s\n", errMsg);
        sqlite3_free(errMsg);
        ROLLBACK_AND_RETURN();
    }

    dprintf(D_FULLDEBUG, "Garbage collection successful.\n");
    return true;
}

