simple-object-server/src/database.cpp

#include <stdexcept>
#include <sstream>
#include <set>

#include "database.hpp"
#include "sqlite3/sqlite3.h"

namespace simple_object_storage {

bool Database::createObjectsTable() {
    const char* create_table_sql = 
        "CREATE TABLE IF NOT EXISTS objects ("
        "hash TEXT PRIMARY KEY,"
        "labeltags TEXT NOT NULL," // JSON array of label:tag pairs
        "metadata TEXT NOT NULL"
        ");";
    
    char* err_msg = nullptr;
    int rc = sqlite3_exec(db_, create_table_sql, nullptr, nullptr, &err_msg);
    if (rc != SQLITE_OK) {
        std::string error = err_msg;
        sqlite3_free(err_msg);
        return false;
    }
    return true;
}

bool Database::createVersionTable() {
    const char* sql = 
        "CREATE TABLE IF NOT EXISTS version_info ("
        "version INTEGER NOT NULL"
        ");";
    
    char* err_msg = nullptr;
    int rc = sqlite3_exec(db_, sql, nullptr, nullptr, &err_msg);
    if (rc != SQLITE_OK) {
        std::string error = err_msg;
        sqlite3_free(err_msg);
        return false;
    }

    // Check if we need to insert initial version
    sqlite3_stmt* stmt;
    sql = "SELECT COUNT(*) FROM version_info;";
    if (sqlite3_prepare_v2(db_, sql, -1, &stmt, nullptr) != SQLITE_OK) {
        return false;
    }

    bool needs_initial_version = false;
    if (sqlite3_step(stmt) == SQLITE_ROW) {
        needs_initial_version = sqlite3_column_int(stmt, 0) == 0;
    }
    sqlite3_finalize(stmt);

    if (needs_initial_version) {
        sql = "INSERT INTO version_info (version) VALUES (?);";
        if (sqlite3_prepare_v2(db_, sql, -1, &stmt, nullptr) != SQLITE_OK) {
            return false;
        }
        sqlite3_bind_int(stmt, 1, CURRENT_VERSION);
        bool success = sqlite3_step(stmt) == SQLITE_DONE;
        sqlite3_finalize(stmt);
        return success;
    }

    return true;
}

bool Database::getVersion(int& version) {
    const char* sql = "SELECT version FROM version_info;";
    sqlite3_stmt* stmt;
    
    if (sqlite3_prepare_v2(db_, sql, -1, &stmt, nullptr) != SQLITE_OK) {
        return false;
    }

    if (sqlite3_step(stmt) != SQLITE_ROW) {
        sqlite3_finalize(stmt);
        return false;
    }

    version = sqlite3_column_int(stmt, 0);
    sqlite3_finalize(stmt);
    return true;
}

bool Database::setVersion(int version) {
    const char* sql = "UPDATE version_info SET version = ?;";
    sqlite3_stmt* stmt;
    
    if (sqlite3_prepare_v2(db_, sql, -1, &stmt, nullptr) != SQLITE_OK) {
        return false;
    }

    sqlite3_bind_int(stmt, 1, version);
    bool success = sqlite3_step(stmt) == SQLITE_DONE;
    sqlite3_finalize(stmt);
    return success;
}

bool Database::migrate(int from_version, int to_version) {
    if (from_version < 3 && to_version == 3) {
        // Drop old table
        const char* drop_sql = "DROP TABLE IF EXISTS objects;";
        char* err_msg = nullptr;
        int rc = sqlite3_exec(db_, drop_sql, nullptr, nullptr, &err_msg);
        if (rc != SQLITE_OK) {
            std::string error = err_msg;
            sqlite3_free(err_msg);
            return false;
        }

        // Create new table with updated schema
        return createObjectsTable();
    }
    return false;
}

Database::Database(const std::filesystem::path& path) : path_(path) {
    int rc = sqlite3_open(path.string().c_str(), &db_);
    if (rc != SQLITE_OK) {
        throw std::runtime_error("Cannot open database: " + std::string(sqlite3_errmsg(db_)));
    }

    // Create version table and set initial version
    if (!createVersionTable()) {
        throw std::runtime_error("Failed to create version table");
    }

    // Check current version and migrate if needed
    int current_version;
    if (!getVersion(current_version)) {
        throw std::runtime_error("Failed to get database version");
    }

    if (current_version != CURRENT_VERSION) {
        if (!migrate(current_version, CURRENT_VERSION)) {
            throw std::runtime_error("Failed to migrate database");
        }
        if (!setVersion(CURRENT_VERSION)) {
            throw std::runtime_error("Failed to update database version");
        }
    }

    // Create objects table if it doesn't exist
    if (!createObjectsTable()) {
        throw std::runtime_error("Failed to create objects table");
    }
}

Database::~Database() {
    if (db_) {
        sqlite3_close(db_);
    }
}

// bool Database::remove(const std::string& labeltag) {
//     std::string sql = "DELETE FROM objects WHERE labeltag = ?;";
//     sqlite3_stmt* stmt;
    
//     if (sqlite3_prepare_v2(db_, sql.c_str(), -1, &stmt, nullptr) != SQLITE_OK) {
//         return false;
//     }

//     sqlite3_bind_text(stmt, 1, labeltag.c_str(), -1, SQLITE_STATIC);
//     bool success = sqlite3_step(stmt) == SQLITE_DONE;
//     sqlite3_finalize(stmt);
//     return success;
// }

bool Database::remove_by_hash(const std::string& hash) {
    std::string sql = "DELETE FROM objects WHERE hash = ?;";
    sqlite3_stmt* stmt;
    
    if (sqlite3_prepare_v2(db_, sql.c_str(), -1, &stmt, nullptr) != SQLITE_OK) {
        return false;
    }

    sqlite3_bind_text(stmt, 1, hash.c_str(), -1, SQLITE_STATIC);
    bool success = sqlite3_step(stmt) == SQLITE_DONE;
    sqlite3_finalize(stmt);
    return success;
}

bool Database::get(const std::string& hash_or_labeltag, dbEntry& entry) {
    std::string sql;
    if (hash_or_labeltag.find(':') != std::string::npos) {
        // Query by label:tag - search for exact match in the JSON array
        sql = "SELECT hash, labeltags, metadata FROM objects WHERE json_array_length(labeltags) > 0 AND EXISTS (SELECT 1 FROM json_each(labeltags) WHERE value = ?);";
    } else {
        // Query by hash
        sql = "SELECT hash, labeltags, metadata FROM objects WHERE hash = ?;";
    }
    
    sqlite3_stmt* stmt;
    if (sqlite3_prepare_v2(db_, sql.c_str(), -1, &stmt, nullptr) != SQLITE_OK) {
        return false;
    }

    sqlite3_bind_text(stmt, 1, hash_or_labeltag.c_str(), -1, SQLITE_STATIC);
    
    if (sqlite3_step(stmt) != SQLITE_ROW) {
        sqlite3_finalize(stmt);
        return false;
    }

    entry.hash = reinterpret_cast<const char*>(sqlite3_column_text(stmt, 0));
    std::string labeltags_str = reinterpret_cast<const char*>(sqlite3_column_text(stmt, 1));
    std::string metadata_str = reinterpret_cast<const char*>(sqlite3_column_text(stmt, 2));
    
    entry.labeltags = nlohmann::json::parse(labeltags_str).get<std::vector<std::string>>();
    entry.metadata = nlohmann::json::parse(metadata_str);

    sqlite3_finalize(stmt);
    return true;
}

bool Database::list(std::vector<dbEntry>& entries) {
    std::string sql = "SELECT hash, labeltags, metadata FROM objects;";
    sqlite3_stmt* stmt;
    
    if (sqlite3_prepare_v2(db_, sql.c_str(), -1, &stmt, nullptr) != SQLITE_OK) {
        return false;
    }

    entries.clear();
    while (sqlite3_step(stmt) == SQLITE_ROW) {
        dbEntry entry;
        entry.hash = reinterpret_cast<const char*>(sqlite3_column_text(stmt, 0));
        std::string labeltags_str = reinterpret_cast<const char*>(sqlite3_column_text(stmt, 1));
        std::string metadata_str = reinterpret_cast<const char*>(sqlite3_column_text(stmt, 2));
        
        entry.labeltags = nlohmann::json::parse(labeltags_str).get<std::vector<std::string>>();
        entry.metadata = nlohmann::json::parse(metadata_str);
        entries.push_back(entry);
    }

    sqlite3_finalize(stmt);
    return true;
}

bool Database::merge_existing_entry(const dbEntry& existing, const dbEntry& new_entry, dbEntry& merged) {
    // Merge label:tag pairs
    std::set<std::string> merged_labeltags(existing.labeltags.begin(), existing.labeltags.end());
    merged_labeltags.insert(new_entry.labeltags.begin(), new_entry.labeltags.end());
    
    // Create merged entry
    merged = new_entry; // Start with new entry's data
    merged.labeltags = std::vector<std::string>(merged_labeltags.begin(), merged_labeltags.end());
    
    // Update metadata - preserve fields from existing entry that aren't in new entry
    merged.metadata = existing.metadata; // Start with existing metadata
    for (const auto& [key, value] : new_entry.metadata.items()) {
        merged.metadata[key] = value; // Override with new values
    }
    
    // Ensure required fields are set correctly
    merged.metadata["labeltags"] = merged.labeltags;
    merged.metadata["hash"] = merged.hash;
    
    // Update database
    std::string sql = "UPDATE objects SET labeltags = ?, metadata = ? WHERE hash = ?;";
    sqlite3_stmt* stmt;
    
    if (sqlite3_prepare_v2(db_, sql.c_str(), -1, &stmt, nullptr) != SQLITE_OK) {
        return false;
    }

    std::string labeltags_str = nlohmann::json(merged.labeltags).dump();
    std::string metadata_str = merged.metadata.dump();
    
    sqlite3_bind_text(stmt, 1, labeltags_str.c_str(), -1, SQLITE_STATIC);
    sqlite3_bind_text(stmt, 2, metadata_str.c_str(), -1, SQLITE_STATIC);
    sqlite3_bind_text(stmt, 3, merged.hash.c_str(), -1, SQLITE_STATIC);

    bool success = sqlite3_step(stmt) == SQLITE_DONE;
    sqlite3_finalize(stmt);
    return success;
}

bool Database::insert_new_entry(const dbEntry& entry) {
    std::string sql = "INSERT INTO objects (hash, labeltags, metadata) VALUES (?, ?, ?);";
    sqlite3_stmt* stmt;
    
    if (sqlite3_prepare_v2(db_, sql.c_str(), -1, &stmt, nullptr) != SQLITE_OK) {
        return false;
    }

    // Update metadata to include labeltags and hash
    nlohmann::json metadata = entry.metadata;
    metadata["labeltags"] = entry.labeltags;
    metadata["hash"] = entry.hash;

    std::string labeltags_str = nlohmann::json(entry.labeltags).dump();
    std::string metadata_str = metadata.dump();
    
    sqlite3_bind_text(stmt, 1, entry.hash.c_str(), -1, SQLITE_STATIC);
    sqlite3_bind_text(stmt, 2, labeltags_str.c_str(), -1, SQLITE_STATIC);
    sqlite3_bind_text(stmt, 3, metadata_str.c_str(), -1, SQLITE_STATIC);

    bool success = sqlite3_step(stmt) == SQLITE_DONE;
    sqlite3_finalize(stmt);
    return success;
}

bool Database::handle_tag_conflicts(const dbEntry& entry) {
    for (const auto& labeltag : entry.labeltags) {
        // Find all entries with this exact label:tag pair
        std::string find_sql = "SELECT hash, labeltags, metadata FROM objects WHERE labeltags LIKE ?;";
        sqlite3_stmt* stmt;
        if (sqlite3_prepare_v2(db_, find_sql.c_str(), -1, &stmt, nullptr) != SQLITE_OK) {
            return false;
        }

        std::string labeltag_pattern = "%\"" + labeltag + "\"%";
        sqlite3_bind_text(stmt, 1, labeltag_pattern.c_str(), -1, SQLITE_STATIC);

        while (sqlite3_step(stmt) == SQLITE_ROW) {
            std::string other_hash = reinterpret_cast<const char*>(sqlite3_column_text(stmt, 0));
            if (other_hash == entry.hash) continue; // Skip our own entry

            std::string other_labeltags_str = reinterpret_cast<const char*>(sqlite3_column_text(stmt, 1));
            std::string other_metadata_str = reinterpret_cast<const char*>(sqlite3_column_text(stmt, 2));

            // Parse the other entry
            dbEntry other;
            other.hash = other_hash;
            other.labeltags = nlohmann::json::parse(other_labeltags_str).get<std::vector<std::string>>();
            other.metadata = nlohmann::json::parse(other_metadata_str);

            // Remove the exact label:tag pair
            std::vector<std::string> new_labeltags;
            for (const auto& other_labeltag : other.labeltags) {
                if (other_labeltag != labeltag) {
                    new_labeltags.push_back(other_labeltag);
                }
            }

            // Update the other entry if it had the label:tag pair removed
            if (new_labeltags.size() != other.labeltags.size()) {
                other.labeltags = new_labeltags;
                other.metadata["labeltags"] = new_labeltags; // Update metadata to match

                std::string update_sql = "UPDATE objects SET labeltags = ?, metadata = ? WHERE hash = ?;";
                sqlite3_stmt* update_stmt;
                if (sqlite3_prepare_v2(db_, update_sql.c_str(), -1, &update_stmt, nullptr) != SQLITE_OK) {
                    sqlite3_finalize(stmt);
                    return false;
                }

                std::string new_labeltags_str = nlohmann::json(new_labeltags).dump();
                std::string new_metadata_str = other.metadata.dump();
                
                sqlite3_bind_text(update_stmt, 1, new_labeltags_str.c_str(), -1, SQLITE_STATIC);
                sqlite3_bind_text(update_stmt, 2, new_metadata_str.c_str(), -1, SQLITE_STATIC);
                sqlite3_bind_text(update_stmt, 3, other.hash.c_str(), -1, SQLITE_STATIC);

                bool update_success = sqlite3_step(update_stmt) == SQLITE_DONE;
                sqlite3_finalize(update_stmt);
                if (!update_success) {
                    sqlite3_finalize(stmt);
                    return false;
                }
            }
        }
        sqlite3_finalize(stmt);
    }
    return true;
}

bool Database::update_or_insert(const dbEntry& entry) {
    // First try to get existing entry by hash
    dbEntry existing;
    bool exists = get(entry.hash, existing);
    
    if (exists) {
        // Merge with existing entry
        dbEntry merged;
        if (!merge_existing_entry(existing, entry, merged)) {
            return false;
        }
        
        // Handle tag conflicts
        return handle_tag_conflicts(merged);
    } else {
        // Insert new entry
        if (!insert_new_entry(entry)) {
            return false;
        }
        
        // Handle tag conflicts
        return handle_tag_conflicts(entry);
    }
}

} // namespace simple_object_storage