simple-object-server/src/server.cpp

#include "server.hpp"
#include "hash.hpp"
#include <filesystem>
#include <iostream>
#include <sstream>
#include <fstream>
#include <random>     // For random number generation
#include <chrono>     // For seeding random number generator
#include <vector> // For getAllKeys
#include <string_view> // For litecask values
#include <sqlite3.h> // Include SQLite
#include <stdexcept> // For std::runtime_error

namespace simple_object_storage {

// Simple RAII helper for file deletion
class ScopeFileDeleter {
public:
    ScopeFileDeleter(const std::filesystem::path& path) : path_(path), released_(false) {}
    ~ScopeFileDeleter() {
        if (!released_) {
            try {
                if (std::filesystem::exists(path_)) {
                    std::filesystem::remove(path_);
                }
            } catch (const std::filesystem::filesystem_error& e) {
                std::cerr << "Error deleting temp file: " << path_ << " - " << e.what() << std::endl;
            }
        }
    }
    void release() { released_ = true; }
private:
    std::filesystem::path path_;
    bool released_;
};

// Helper to execute SQL and check for errors
void execute_sql(sqlite3* db, const char* sql, const std::string& error_msg_prefix) {
    char* err_msg = nullptr;
    int rc = sqlite3_exec(db, sql, 0, 0, &err_msg);
    if (rc != SQLITE_OK) {
        std::string error_details = error_msg_prefix + ": " + (err_msg ? err_msg : "Unknown error");
        sqlite3_free(err_msg);
        throw std::runtime_error(error_details);
    }
}

bool Server::init_db() {
    db_path_ = config_.object_store_path / "index.db";
    int rc = sqlite3_open(db_path_.c_str(), &db_);
    if (rc != SQLITE_OK) {
        std::cerr << "Failed to open/create SQLite database '" << db_path_ << "': " << sqlite3_errmsg(db_) << std::endl;
        db_ = nullptr; // Ensure db_ is null if open failed
        return false;
    }

    try {
        // Enable WAL mode for better concurrency
        execute_sql(db_, "PRAGMA journal_mode=WAL;", "Failed to set WAL mode");

        // Create table if it doesn't exist
        const char* create_table_sql = 
            "CREATE TABLE IF NOT EXISTS objects (" 
            "label_tag TEXT PRIMARY KEY UNIQUE NOT NULL, " 
            "hash TEXT NOT NULL);";
        execute_sql(db_, create_table_sql, "Failed to create objects table");

    } catch (const std::runtime_error& e) {
        std::cerr << "Database initialization error: " << e.what() << std::endl;
        sqlite3_close(db_);
        db_ = nullptr;
        return false;
    }
    return true;
}

Server::Server(const ServerConfig& config)
    : config_(config), running_(false), db_(nullptr) {
    // Ensure object store directory exists
    try {
        std::filesystem::create_directories(config_.object_store_path);
    } catch (const std::filesystem::filesystem_error& e) {
        std::cerr << "Failed to create object store directory: " << config_.object_store_path << " - " << e.what() << std::endl;
        return;
    }

    // Initialize the database
    if (!init_db()) {
        // Error already printed in init_db
        // Consider throwing or setting an error state
    }
}

Server::~Server() {
    stop();
    if (db_) {
        sqlite3_close(db_); // Close the database connection
        db_ = nullptr;
    }
}

bool Server::start() {
    if (!db_) { // Check if DB initialization failed
        std::cerr << "Database is not initialized. Cannot start server." << std::endl;
        return false;
    }
    setup_routes();

    std::cout << "Server starting on " << config_.host << ":" << config_.port << std::endl;
    running_ = true;
    if (!server_.listen(config_.host.c_str(), config_.port)) {
        running_ = false;
        std::cerr << "Failed to listen on " << config_.host << ":" << config_.port << std::endl;
        return false;
    }
    return true; // Should not be reached if listen is blocking and successful
}

void Server::stop() {
    if (running_) {
        server_.stop();
        running_ = false;
        std::cout << "Server stopped." << std::endl;
    }
}

void Server::setup_routes() {
    const std::string welcome_page = "<html><body><h1>simple_object_storage Template Registry</h1></body></html>";
    // Welcome page
    server_.Get("/", [welcome_page](const httplib::Request&, httplib::Response& res) {
        res.set_content(welcome_page, "text/html");
    });

    server_.Get("/index.html", [welcome_page](const httplib::Request&, httplib::Response& res) {
        res.set_content(welcome_page, "text/html");
    });

    // Get object by hash or label:tag
    server_.Get("/object/(.*)", [this](const httplib::Request& req, httplib::Response& res) {
        handle_get_object(req, res);
    });

    // Get hash for label:tag
    server_.Get("/hash/(.*)", [this](const httplib::Request& req, httplib::Response& res) {
        handle_get_hash(req, res);
    });

    // Get directory listing
    server_.Get("/dir", [this](const httplib::Request& req, httplib::Response& res) {
        handle_get_directory(req, res);
    });

    // Upload object
    server_.Put("/([^/]+)/(.*)", [this](const httplib::Request& req, httplib::Response& res) { // Adjusted regex slightly for label:tag
        handle_put_object(req, res);
    });
}

void Server::handle_get_object(const httplib::Request& req, httplib::Response& res) {
    const auto& key = req.matches[1].str();
    std::string hash_str;

    // Check if the key looks like a hash (numeric)
    bool is_hash_lookup = true;
    for (char c : key) {
        if (!std::isdigit(c)) {
            is_hash_lookup = false;
            break;
        }
    }

    if (!is_hash_lookup) {
        // Lookup by label:tag in the SQLite database
        sqlite3_stmt* stmt = nullptr;
        const char* sql = "SELECT hash FROM objects WHERE label_tag = ?;";
        int rc = sqlite3_prepare_v2(db_, sql, -1, &stmt, nullptr);

        if (rc != SQLITE_OK) {
            std::cerr << "Failed to prepare statement (get hash): " << sqlite3_errmsg(db_) << std::endl;
            res.status = 500;
            res.set_content("Database error preparing statement", "text/plain");
            return;
        }

        sqlite3_bind_text(stmt, 1, key.c_str(), -1, SQLITE_STATIC);

        rc = sqlite3_step(stmt);
        if (rc == SQLITE_ROW) {
            const unsigned char* text = sqlite3_column_text(stmt, 0);
            if (text) {
                hash_str = reinterpret_cast<const char*>(text);
            }
        } else if (rc == SQLITE_DONE) {
            // Not found
            sqlite3_finalize(stmt);
            res.status = 404;
            res.set_content("Object not found (label:tag)", "text/plain");
            return;
        } else {
            std::cerr << "Failed to execute statement (get hash): " << sqlite3_errmsg(db_) << std::endl;
            sqlite3_finalize(stmt);
            res.status = 500;
            res.set_content("Database error executing statement", "text/plain");
            return;
        }
        sqlite3_finalize(stmt);

    } else {
        // Lookup directly by hash
        hash_str = key;
    }

    if (hash_str.empty()) {
        // Should have been caught earlier if not found, but as a safeguard
         res.status = 404;
         res.set_content("Object hash could not be determined", "text/plain");
         return;
    }

    // Construct the file path using the hash string
    std::filesystem::path file_path = config_.object_store_path / hash_str;
    if (!std::filesystem::exists(file_path) || !std::filesystem::is_regular_file(file_path)) {
        res.status = 404;
        res.set_content("Object file not found for hash: " + hash_str, "text/plain");
        return;
    }

    // Send file using Response::set_file_content
    std::string content_type = "application/octet-stream"; // Basic default
    res.set_file_content(file_path.string(), content_type);
}

void Server::handle_get_hash(const httplib::Request& req, httplib::Response& res) {
    const auto& label_tag = req.matches[1].str();

    sqlite3_stmt* stmt = nullptr;
    const char* sql = "SELECT hash FROM objects WHERE label_tag = ?;";
    int rc = sqlite3_prepare_v2(db_, sql, -1, &stmt, nullptr);

    if (rc != SQLITE_OK) {
        std::cerr << "Failed to prepare statement (get hash direct): " << sqlite3_errmsg(db_) << std::endl;
        res.status = 500;
        res.set_content("Database error preparing statement", "text/plain");
        return;
    }

    sqlite3_bind_text(stmt, 1, label_tag.c_str(), -1, SQLITE_STATIC);

    rc = sqlite3_step(stmt);
    if (rc == SQLITE_ROW) {
        const unsigned char* text = sqlite3_column_text(stmt, 0);
        if (text) {
            res.set_content(reinterpret_cast<const char*>(text), "text/plain");
        }
    } else if (rc == SQLITE_DONE) {
        res.status = 404;
        res.set_content("Label:tag not found", "text/plain");
    } else {
        std::cerr << "Failed to execute statement (get hash direct): " << sqlite3_errmsg(db_) << std::endl;
        res.status = 500;
        res.set_content("Database error executing statement", "text/plain");
    }
    sqlite3_finalize(stmt);
}

void Server::handle_get_directory(const httplib::Request& /*req*/, httplib::Response& res) {
    std::stringstream ss;
    sqlite3_stmt* stmt = nullptr;
    const char* sql = "SELECT label_tag, hash FROM objects;";
    int rc = sqlite3_prepare_v2(db_, sql, -1, &stmt, nullptr);

    if (rc != SQLITE_OK) {
        std::cerr << "Failed to prepare statement (get dir): " << sqlite3_errmsg(db_) << std::endl;
        res.status = 500;
        res.set_content("Database error preparing statement", "text/plain");
        return;
    }

    while ((rc = sqlite3_step(stmt)) == SQLITE_ROW) {
        const unsigned char* label_tag_text = sqlite3_column_text(stmt, 0);
        const unsigned char* hash_text = sqlite3_column_text(stmt, 1);
        if (label_tag_text && hash_text) {
             ss << reinterpret_cast<const char*>(label_tag_text) << "," 
                << reinterpret_cast<const char*>(hash_text) << "\n";
        }
    }

    if (rc != SQLITE_DONE) {
         std::cerr << "Failed to execute/iterate statement (get dir): " << sqlite3_errmsg(db_) << std::endl;
         // Don't overwrite potential results, but log error
         if (ss.str().empty()) { // Only send error if no data was retrieved
             res.status = 500;
             res.set_content("Database error executing statement", "text/plain");
         }
    }

    sqlite3_finalize(stmt);
    res.set_content(ss.str(), "text/plain");
}

void Server::handle_put_object(const httplib::Request& req, httplib::Response& res) {
    const auto& token = req.matches[1].str();
    const auto& label_tag = req.matches[2].str();

    if (!validate_write_token(token)) {
        res.status = 403;
        res.set_content("Invalid write token", "text/plain");
        return;
    }

    auto [label, tag] = parse_label_tag(label_tag);
    if (label.empty() || tag.empty()) {
        res.status = 400;
        res.set_content("Invalid label:tag format", "text/plain");
        return;
    }

    // Generate a random number for the temporary filename
    std::mt19937_64 rng(std::chrono::high_resolution_clock::now().time_since_epoch().count());
    std::uniform_int_distribution<uint64_t> dist;
    uint64_t random_num = dist(rng);
    std::string temp_filename = "temp_" + std::to_string(random_num);

    // Create temporary file
    std::filesystem::path temp_path = config_.object_store_path / temp_filename;
    std::ofstream temp_file(temp_path, std::ios::binary);
    if (!temp_file.is_open()) {
        res.status = 500;
        res.set_content("Failed to create temporary file", "text/plain");
        return;
    }

    // Write request body to temporary file
    temp_file.write(req.body.data(), req.body.size());
    temp_file.close();

    // Ensure the temporary file is removed even if errors occur
    ScopeFileDeleter temp_file_deleter(temp_path);

    // Calculate hash
    uint64_t hash = hash_file(temp_path.string());
    if (hash == 0) {
        res.status = 500;
        res.set_content("Failed to calculate hash", "text/plain");
        return;
    }

    // Move file to final location
    std::string hash_str = std::to_string(hash);
    std::filesystem::path final_path = config_.object_store_path / hash_str;
    if (!std::filesystem::exists(final_path)) {
        try {
            std::filesystem::rename(temp_path, final_path);
            temp_file_deleter.release();
        } catch (const std::filesystem::filesystem_error& e) {
            std::cerr << "Error renaming temp file: " << e.what() << std::endl;
            res.status = 500;
            res.set_content("Failed to store object file", "text/plain");
            return;
        }
    }

    // Update SQLite index (INSERT OR REPLACE)
    sqlite3_stmt* stmt = nullptr;
    const char* sql = "INSERT OR REPLACE INTO objects (label_tag, hash) VALUES (?, ?);";
    int rc = sqlite3_prepare_v2(db_, sql, -1, &stmt, nullptr);
    if (rc != SQLITE_OK) {
        std::cerr << "Failed to prepare statement (put object): " << sqlite3_errmsg(db_) << std::endl;
        res.status = 500;
        res.set_content("Database error preparing statement", "text/plain");
        // Attempt to clean up the moved file if index fails
        try { if (std::filesystem::exists(final_path)) std::filesystem::remove(final_path); } catch(...) {};
        return;
    }

    sqlite3_bind_text(stmt, 1, label_tag.c_str(), -1, SQLITE_STATIC);
    sqlite3_bind_text(stmt, 2, hash_str.c_str(), -1, SQLITE_STATIC);

    rc = sqlite3_step(stmt);
    sqlite3_finalize(stmt);

    if (rc != SQLITE_DONE) {
        std::cerr << "Failed to execute statement (put object): " << sqlite3_errmsg(db_) << std::endl;
        res.status = 500;
        res.set_content("Failed to update database index", "text/plain");
        // Attempt to clean up the moved file if index fails
        try { if (std::filesystem::exists(final_path)) std::filesystem::remove(final_path); } catch(...) {};
        return;
    }

    res.set_content(hash_str, "text/plain");
}

bool Server::validate_write_token(const std::string& token) const {
    return std::find(config_.write_tokens.begin(), config_.write_tokens.end(), token) != config_.write_tokens.end();
}

std::pair<std::string, std::string> Server::parse_label_tag(const std::string& label_tag) const {
    size_t colon_pos = label_tag.find(':');
    if (colon_pos == std::string::npos || colon_pos == 0 || colon_pos == label_tag.length() - 1) {
        return {"", ""};
    }
    return {label_tag.substr(0, colon_pos), label_tag.substr(colon_pos + 1)};
}

} // namespace simple_object_storage