smartbotic-database/cli/main.cpp

#include <smartbotic/database/client.hpp>
#include <nlohmann/json.hpp>

#include "storage/subdb_placement.hpp"

#include <openssl/bio.h>
#include <openssl/bn.h>
#include <openssl/err.h>
#include <openssl/evp.h>
#include <openssl/pem.h>
#include <openssl/rand.h>
#include <openssl/x509.h>
#include <openssl/x509v3.h>

#include <arpa/inet.h>
#include <fcntl.h>
#include <sys/random.h>
#include <sys/stat.h>
#include <unistd.h>

#include <cerrno>
#include <cstdio>
#include <cstring>
#include <fstream>
#include <iostream>
#include <sstream>
#include <string>
#include <vector>

using json = nlohmann::json;

namespace {

struct Args {
    std::string address = "localhost:9004";
    std::string command;
    std::vector<std::string> params;
};

// ANSI colors
constexpr auto C_RESET = "\033[0m";
constexpr auto C_BOLD  = "\033[1m";
constexpr auto C_DIM   = "\033[2m";
constexpr auto C_CYAN  = "\033[36m";
constexpr auto C_GREEN = "\033[32m";
constexpr auto C_RED   = "\033[31m";
constexpr auto C_YELLOW = "\033[33m";

void printJson(const json& j) {
    std::cout << j.dump(2) << "\n";
}

void printError(const std::string& msg) {
    std::cerr << C_RED << "error: " << C_RESET << msg << "\n";
}

void printUsage() {
    std::cout << C_BOLD << "smartbotic-db-cli" << C_RESET << " — admin tool for smartbotic-database\n\n"
              << C_BOLD << "Usage:" << C_RESET << "\n"
              << "  smartbotic-db-cli [--address HOST:PORT] <command> [args...]\n"
              << "  smartbotic-db-cli [--address HOST:PORT]   " << C_DIM << "# interactive mode" << C_RESET << "\n\n"
              << C_BOLD << "Commands:" << C_RESET << "\n"
              << "  " << C_CYAN << "collections" << C_RESET << "                          List all collections\n"
              << "  " << C_CYAN << "info" << C_RESET << " <collection>                    Collection info\n"
              << "  " << C_CYAN << "find" << C_RESET << " <collection>                    List documents\n"
              << "  " << C_CYAN << "get" << C_RESET << " <collection> <id>                Get a document\n"
              << "  " << C_CYAN << "upsert" << C_RESET << " <collection> <id> '<json>'    Insert or update\n"
              << "  " << C_CYAN << "remove" << C_RESET << " <collection> <id>             Delete a document\n"
              << "  " << C_CYAN << "count" << C_RESET << " <collection>                   Count documents\n"
              << "  " << C_CYAN << "lock" << C_RESET << "                                 Lock database (read-only)\n"
              << "  " << C_CYAN << "unlock" << C_RESET << "                               Unlock database (accept writes)\n"
              << "  " << C_CYAN << "status" << C_RESET << "                               Show read-only + recovery status\n"
              << "  " << C_CYAN << "help" << C_RESET << "                                 Show this help\n\n"
              << C_BOLD << "Offline helpers" << C_RESET << " " << C_DIM << "(no server connection required)" << C_RESET << ":\n"
              << "  " << C_CYAN << "generate-auth-key" << C_RESET << "                    Emit a base64 32-byte API key\n"
              << "  " << C_CYAN << "generate-tls-cert" << C_RESET << " --bind <addr>      Generate a self-signed TLS cert + key\n"
              << "    " << C_DIM << "[--out-cert PATH] [--out-key PATH] [--days N]" << C_RESET << "\n"
              << "  " << C_CYAN << "verify-subdbs" << C_RESET << " --env <path>          Report documents filed under the wrong collection\n"
              << "    " << C_DIM << "[--project NAME]   read-only; safe against a running server" << C_RESET << "\n"
              << "  " << C_CYAN << "reconcile-subdbs" << C_RESET << " --env <path>       Repair misfiled documents (dry run by default)\n"
              << "    " << C_DIM << "[--project NAME] [--apply]   STOP THE SERVICE and back up before --apply" << C_RESET << "\n\n"
              << C_BOLD << "Options:" << C_RESET << "\n"
              << "  --address HOST:PORT    Database address (default: localhost:9004)\n";
}

// ---------------------------------------------------------------------------
// v2.4 Stage F: offline helpers (no server connection required)
// ---------------------------------------------------------------------------

// Base64-encode raw bytes. Uses OpenSSL's EVP_EncodeBlock which emits the
// standard base64 alphabet (no newlines) and pads with '='.
std::string base64Encode(const unsigned char* data, size_t len) {
    if (len == 0) return {};
    // EVP_EncodeBlock writes ((len + 2) / 3) * 4 bytes + a NUL terminator.
    const size_t out_len = 4 * ((len + 2) / 3);
    std::string out(out_len, '\0');
    int written = EVP_EncodeBlock(
        reinterpret_cast<unsigned char*>(out.data()),
        data, static_cast<int>(len));
    if (written < 0) return {};
    out.resize(static_cast<size_t>(written));
    return out;
}

// Pull 32 cryptographically random bytes from getentropy() (preferred) or
// /dev/urandom (fallback). Returns true on success.
bool fillRandomBytes(unsigned char* buf, size_t len) {
    // getentropy() is limited to 256 bytes per call; our use-case is 32.
    if (len <= 256) {
        if (getentropy(buf, len) == 0) return true;
    }
    // Fallback: /dev/urandom.
    int fd = ::open("/dev/urandom", O_RDONLY | O_CLOEXEC);
    if (fd < 0) return false;
    size_t got = 0;
    while (got < len) {
        ssize_t n = ::read(fd, buf + got, len - got);
        if (n <= 0) {
            if (errno == EINTR) continue;
            ::close(fd);
            return false;
        }
        got += static_cast<size_t>(n);
    }
    ::close(fd);
    return true;
}

int cmdGenerateAuthKey() {
    unsigned char key[32];
    if (!fillRandomBytes(key, sizeof(key))) {
        std::fprintf(stderr, "error: failed to obtain entropy for key\n");
        return 1;
    }
    auto encoded = base64Encode(key, sizeof(key));
    if (encoded.empty()) {
        std::fprintf(stderr, "error: base64 encoding failed\n");
        return 1;
    }
    std::cout << encoded << "\n";
    return 0;
}

namespace {

// Print the topmost OpenSSL error to stderr with a prefix.
void printOpenSslError(const char* prefix) {
    unsigned long e = ERR_get_error();
    char buf[256] = {0};
    if (e != 0) {
        ERR_error_string_n(e, buf, sizeof(buf));
        std::fprintf(stderr, "error: %s: %s\n", prefix, buf);
    } else {
        std::fprintf(stderr, "error: %s\n", prefix);
    }
}

// Returns true if `s` parses as an IPv4 or IPv6 literal.
bool looksLikeIp(const std::string& s) {
    unsigned char buf[16];
    if (inet_pton(AF_INET, s.c_str(), buf) == 1) return true;
    if (inet_pton(AF_INET6, s.c_str(), buf) == 1) return true;
    return false;
}

// Generate an RSA private key as an EVP_PKEY using the modern (3.x) API.
EVP_PKEY* generateRsaKey(int bits) {
    EVP_PKEY_CTX* ctx = EVP_PKEY_CTX_new_id(EVP_PKEY_RSA, nullptr);
    if (!ctx) return nullptr;
    EVP_PKEY* pkey = nullptr;
    if (EVP_PKEY_keygen_init(ctx) <= 0) {
        EVP_PKEY_CTX_free(ctx);
        return nullptr;
    }
    if (EVP_PKEY_CTX_set_rsa_keygen_bits(ctx, bits) <= 0) {
        EVP_PKEY_CTX_free(ctx);
        return nullptr;
    }
    if (EVP_PKEY_keygen(ctx, &pkey) <= 0) {
        EVP_PKEY_CTX_free(ctx);
        return nullptr;
    }
    EVP_PKEY_CTX_free(ctx);
    return pkey;
}

// Write a string to disk with the given file mode. Truncates existing files.
bool writeFileWithMode(const std::string& path, const std::string& contents, mode_t mode) {
    int fd = ::open(path.c_str(),
                    O_WRONLY | O_CREAT | O_TRUNC | O_CLOEXEC,
                    mode);
    if (fd < 0) {
        std::fprintf(stderr, "error: open %s: %s\n", path.c_str(), std::strerror(errno));
        return false;
    }
    // Re-assert mode in case the file pre-existed with a wider mode and
    // O_CREAT was a no-op (open(2) only applies the mode on create).
    if (fchmod(fd, mode) != 0) {
        std::fprintf(stderr, "error: fchmod %s: %s\n", path.c_str(), std::strerror(errno));
        ::close(fd);
        return false;
    }
    size_t off = 0;
    while (off < contents.size()) {
        ssize_t n = ::write(fd, contents.data() + off, contents.size() - off);
        if (n <= 0) {
            if (errno == EINTR) continue;
            std::fprintf(stderr, "error: write %s: %s\n", path.c_str(), std::strerror(errno));
            ::close(fd);
            return false;
        }
        off += static_cast<size_t>(n);
    }
    if (::close(fd) != 0) {
        std::fprintf(stderr, "error: close %s: %s\n", path.c_str(), std::strerror(errno));
        return false;
    }
    return true;
}

// Serialize an X509* to a PEM string.
std::string pemEncodeCert(X509* cert) {
    BIO* bio = BIO_new(BIO_s_mem());
    if (!bio) return {};
    if (PEM_write_bio_X509(bio, cert) != 1) {
        BIO_free(bio);
        return {};
    }
    BUF_MEM* mem = nullptr;
    BIO_get_mem_ptr(bio, &mem);
    std::string out(mem->data, mem->length);
    BIO_free(bio);
    return out;
}

// Serialize an EVP_PKEY* to an unencrypted PEM string (PKCS#8 format via
// PEM_write_bio_PrivateKey).
std::string pemEncodeKey(EVP_PKEY* key) {
    BIO* bio = BIO_new(BIO_s_mem());
    if (!bio) return {};
    if (PEM_write_bio_PrivateKey(bio, key, nullptr, nullptr, 0, nullptr, nullptr) != 1) {
        BIO_free(bio);
        return {};
    }
    BUF_MEM* mem = nullptr;
    BIO_get_mem_ptr(bio, &mem);
    std::string out(mem->data, mem->length);
    BIO_free(bio);
    return out;
}

} // anonymous namespace

int cmdGenerateTlsCert(const std::vector<std::string>& params) {
    std::string bind;
    std::string out_cert = "./server.pem";
    std::string out_key  = "./server.key";
    int days = 3650;

    for (size_t i = 0; i < params.size(); ++i) {
        const auto& p = params[i];
        auto need = [&](const char* flag) -> const std::string* {
            if (i + 1 >= params.size()) {
                std::fprintf(stderr, "error: %s requires a value\n", flag);
                return nullptr;
            }
            return &params[++i];
        };
        if (p == "--bind") {
            auto* v = need("--bind"); if (!v) return 2; bind = *v;
        } else if (p == "--out-cert") {
            auto* v = need("--out-cert"); if (!v) return 2; out_cert = *v;
        } else if (p == "--out-key") {
            auto* v = need("--out-key"); if (!v) return 2; out_key = *v;
        } else if (p == "--days") {
            auto* v = need("--days"); if (!v) return 2;
            try { days = std::stoi(*v); }
            catch (...) {
                std::fprintf(stderr, "error: --days must be an integer\n");
                return 2;
            }
            if (days <= 0) {
                std::fprintf(stderr, "error: --days must be > 0\n");
                return 2;
            }
        } else {
            std::fprintf(stderr, "error: unknown argument: %s\n", p.c_str());
            return 2;
        }
    }

    if (bind.empty()) {
        std::fprintf(stderr,
                     "usage: generate-tls-cert --bind <addr> "
                     "[--out-cert PATH] [--out-key PATH] [--days N]\n");
        return 2;
    }

    // 1. RSA 4096-bit key.
    EVP_PKEY* pkey = generateRsaKey(4096);
    if (!pkey) {
        printOpenSslError("RSA key generation failed");
        return 1;
    }

    // 2. X.509 certificate.
    X509* cert = X509_new();
    if (!cert) {
        printOpenSslError("X509_new failed");
        EVP_PKEY_free(pkey);
        return 1;
    }

    // Version 3 (the integer field encodes v3 as 2).
    if (X509_set_version(cert, 2) != 1) {
        printOpenSslError("X509_set_version failed");
        X509_free(cert); EVP_PKEY_free(pkey);
        return 1;
    }

    // Random 64-bit serial number.
    {
        unsigned char serial_bytes[8];
        if (RAND_bytes(serial_bytes, sizeof(serial_bytes)) != 1) {
            printOpenSslError("RAND_bytes for serial failed");
            X509_free(cert); EVP_PKEY_free(pkey);
            return 1;
        }
        // Mask the top bit so the BIGNUM is positive.
        serial_bytes[0] &= 0x7F;
        BIGNUM* bn = BN_bin2bn(serial_bytes, sizeof(serial_bytes), nullptr);
        if (!bn) {
            printOpenSslError("BN_bin2bn failed");
            X509_free(cert); EVP_PKEY_free(pkey);
            return 1;
        }
        ASN1_INTEGER* ai = BN_to_ASN1_INTEGER(bn, nullptr);
        BN_free(bn);
        if (!ai) {
            printOpenSslError("BN_to_ASN1_INTEGER failed");
            X509_free(cert); EVP_PKEY_free(pkey);
            return 1;
        }
        if (X509_set_serialNumber(cert, ai) != 1) {
            printOpenSslError("X509_set_serialNumber failed");
            ASN1_INTEGER_free(ai);
            X509_free(cert); EVP_PKEY_free(pkey);
            return 1;
        }
        ASN1_INTEGER_free(ai);
    }

    // Validity period.
    if (!X509_gmtime_adj(X509_get_notBefore(cert), 0)) {
        printOpenSslError("X509_gmtime_adj(notBefore) failed");
        X509_free(cert); EVP_PKEY_free(pkey);
        return 1;
    }
    long seconds = static_cast<long>(days) * 24L * 60L * 60L;
    if (!X509_gmtime_adj(X509_get_notAfter(cert), seconds)) {
        printOpenSslError("X509_gmtime_adj(notAfter) failed");
        X509_free(cert); EVP_PKEY_free(pkey);
        return 1;
    }

    // Public key.
    if (X509_set_pubkey(cert, pkey) != 1) {
        printOpenSslError("X509_set_pubkey failed");
        X509_free(cert); EVP_PKEY_free(pkey);
        return 1;
    }

    // Subject + issuer name (self-signed, so identical).
    X509_NAME* name = X509_get_subject_name(cert);
    if (X509_NAME_add_entry_by_txt(
            name, "CN", MBSTRING_UTF8,
            reinterpret_cast<const unsigned char*>(bind.c_str()),
            -1, -1, 0) != 1) {
        printOpenSslError("X509_NAME_add_entry_by_txt(CN) failed");
        X509_free(cert); EVP_PKEY_free(pkey);
        return 1;
    }
    if (X509_set_issuer_name(cert, name) != 1) {
        printOpenSslError("X509_set_issuer_name failed");
        X509_free(cert); EVP_PKEY_free(pkey);
        return 1;
    }

    // SubjectAltName.
    {
        std::string san = "DNS:localhost,IP:127.0.0.1";
        if (bind != "localhost" && bind != "127.0.0.1") {
            san += ',';
            san += looksLikeIp(bind) ? "IP:" : "DNS:";
            san += bind;
        }
        X509_EXTENSION* ext = X509V3_EXT_conf_nid(
            nullptr, nullptr, NID_subject_alt_name, san.c_str());
        if (!ext) {
            printOpenSslError("X509V3_EXT_conf_nid(SAN) failed");
            X509_free(cert); EVP_PKEY_free(pkey);
            return 1;
        }
        if (X509_add_ext(cert, ext, -1) != 1) {
            printOpenSslError("X509_add_ext(SAN) failed");
            X509_EXTENSION_free(ext);
            X509_free(cert); EVP_PKEY_free(pkey);
            return 1;
        }
        X509_EXTENSION_free(ext);
    }

    // basicConstraints CA:FALSE — a server leaf cert, not a CA.
    {
        X509_EXTENSION* ext = X509V3_EXT_conf_nid(
            nullptr, nullptr, NID_basic_constraints, "critical,CA:FALSE");
        if (ext) {
            X509_add_ext(cert, ext, -1);
            X509_EXTENSION_free(ext);
        }
    }

    // Sign with SHA-256.
    if (X509_sign(cert, pkey, EVP_sha256()) == 0) {
        printOpenSslError("X509_sign failed");
        X509_free(cert); EVP_PKEY_free(pkey);
        return 1;
    }

    // Serialize.
    std::string cert_pem = pemEncodeCert(cert);
    std::string key_pem  = pemEncodeKey(pkey);
    X509_free(cert);
    EVP_PKEY_free(pkey);

    if (cert_pem.empty() || key_pem.empty()) {
        std::fprintf(stderr, "error: PEM encoding failed\n");
        return 1;
    }

    // Write key first (0600), then cert (0644). If either fails, the other
    // may have been written — that's acceptable; the operator will see the
    // error and retry.
    if (!writeFileWithMode(out_key, key_pem, 0600)) return 1;
    if (!writeFileWithMode(out_cert, cert_pem, 0644)) return 1;

    std::cout << "Wrote cert: " << out_cert << "\n"
              << "Wrote key:  " << out_key  << "\n";
    return 0;
}

bool execCommand(smartbotic::database::Client& client,
                 const std::string& cmd, const std::vector<std::string>& params) {
    try {
        if (cmd == "collections") {
            auto collections = client.listCollections();
            std::cout << C_BOLD << "Collections:" << C_RESET << "\n";
            for (const auto& c : collections) {
                auto info = client.getCollectionInfo(c);
                int64_t count = 0;
                if (info) count = info->documentCount;
                std::cout << "  " << C_CYAN << c << C_RESET
                          << C_DIM << " (" << count << " docs)" << C_RESET << "\n";
            }
            return true;
        }

        if (cmd == "info") {
            if (params.empty()) { printError("usage: info <collection>"); return false; }
            auto info = client.getCollectionInfo(params[0]);
            if (!info) { printError("collection not found: " + params[0]); return false; }
            std::cout << C_BOLD << params[0] << C_RESET << ":\n"
                      << "  documents: " << info->documentCount << "\n"
                      << "  size: " << info->sizeBytes << " bytes\n"
                      << "  encrypted: " << (info->encrypted ? "yes" : "no") << "\n"
                      << "  max_versions: " << info->maxVersions << "\n";
            if (info->defaultTtlSeconds > 0)
                std::cout << "  ttl: " << info->defaultTtlSeconds << "s\n";
            return true;
        }

        if (cmd == "find") {
            if (params.empty()) { printError("usage: find <collection> [--limit N] [--exists FIELD]"); return false; }
            smartbotic::database::Client::QueryOptions opts;
            opts.limit = 100;
            for (size_t i = 1; i < params.size(); ++i) {
                if (params[i] == "--limit" && i + 1 < params.size()) {
                    opts.limit = static_cast<uint32_t>(std::stoul(params[++i]));
                } else if (params[i] == "--exists" && i + 1 < params.size()) {
                    opts.filters.emplace_back(params[++i], smartbotic::database::Client::FilterOp::EXISTS, true);
                }
            }
            auto docs = client.find(params[0], opts);
            std::cout << C_DIM << "(" << docs.size() << " documents)" << C_RESET << "\n";
            for (const auto& doc : docs) {
                auto id = doc.value("_id", "");
                // Print compact summary line
                std::cout << C_GREEN << id << C_RESET;
                // Show a few key fields
                for (const auto& [k, v] : doc.items()) {
                    if (k == "_id" || k == "_created_at" || k == "_updated_at") continue;
                    auto val = v.dump();
                    if (val.size() > 60) val = val.substr(0, 57) + "...";
                    std::cout << " " << C_DIM << k << "=" << C_RESET << val;
                    // Limit to 3 fields per line
                    static int field_count = 0;
                    if (++field_count >= 3) { field_count = 0; break; }
                }
                std::cout << "\n";
            }
            return true;
        }

        if (cmd == "get") {
            if (params.size() < 2) { printError("usage: get <collection> <id>"); return false; }
            auto doc = client.get(params[0], params[1]);
            if (!doc) { printError("not found: " + params[0] + "/" + params[1]); return false; }
            printJson(*doc);
            return true;
        }

        if (cmd == "upsert") {
            if (params.size() < 3) { printError("usage: upsert <collection> <id> '<json>'"); return false; }
            auto data = json::parse(params[2], nullptr, false);
            if (!data.is_object()) { printError("invalid JSON: " + params[2]); return false; }
            client.upsert(params[0], data, params[1]);
            std::cout << C_GREEN << "ok" << C_RESET << " " << params[0] << "/" << params[1] << "\n";
            return true;
        }

        if (cmd == "remove" || cmd == "delete") {
            if (params.size() < 2) { printError("usage: remove <collection> <id>"); return false; }
            client.remove(params[0], params[1]);
            std::cout << C_GREEN << "ok" << C_RESET << " removed " << params[0] << "/" << params[1] << "\n";
            return true;
        }

        if (cmd == "count") {
            if (params.empty()) { printError("usage: count <collection>"); return false; }
            auto info = client.getCollectionInfo(params[0]);
            if (!info) { printError("collection not found: " + params[0]); return false; }
            std::cout << info->documentCount << "\n";
            return true;
        }

        if (cmd == "lock") {
            if (client.setReadOnly(true)) {
                std::cout << C_GREEN << "ok" << C_RESET << " Database locked (read-only)\n";
                return true;
            }
            printError("failed to lock database");
            return false;
        }

        if (cmd == "unlock") {
            if (client.setReadOnly(false)) {
                std::cout << C_GREEN << "ok" << C_RESET << " Database unlocked (writes accepted)\n";
                return true;
            }
            printError("failed to unlock database");
            return false;
        }

        if (cmd == "status") {
            auto s = client.getReadOnlyStatus();
            std::cout << "Read-only:          " << (s.readOnly ? (std::string(C_RED) + "YES" + C_RESET) : "no") << "\n";
            if (s.readOnly) {
                std::cout << "Reason:             " << s.reason << "\n";
            }
            std::cout << "Recovery outcome:   " << s.recoveryOutcome << "\n";
            if (!s.expectedSnapshot.empty()) {
                std::cout << "Expected snapshot:  " << s.expectedSnapshot << "\n";
            }
            if (!s.snapshotUsed.empty()) {
                std::cout << "Snapshot used:      " << s.snapshotUsed << "\n";
            }
            if (!s.failureReason.empty()) {
                std::cout << "Failure reason:     " << s.failureReason << "\n";
            }
            std::cout << "WAL replayed:       " << s.walEntriesReplayed << " entries\n";
            std::cout << "Snapshots tried:    " << s.snapshotsAttempted << "\n";
            return true;
        }

        if (cmd == "help" || cmd == "?") {
            printUsage();
            return true;
        }

        printError("unknown command: " + cmd + " (try 'help')");
        return false;

    } catch (const std::exception& e) {
        printError(e.what());
        return false;
    }
}

// Tokenize a line, respecting single/double quotes and JSON braces
std::vector<std::string> tokenize(const std::string& line) {
    std::vector<std::string> tokens;
    std::string current;
    int brace_depth = 0;
    char in_quote = 0;

    for (size_t i = 0; i < line.size(); ++i) {
        char c = line[i];

        if (in_quote) {
            current += c;
            if (c == in_quote && (i == 0 || line[i-1] != '\\')) {
                in_quote = 0;
                // Strip surrounding quotes for simple string tokens
                if (brace_depth == 0 && current.size() >= 2
                    && (current.front() == '\'' || current.front() == '"')
                    && current.front() == current.back()) {
                    current = current.substr(1, current.size() - 2);
                }
            }
            continue;
        }

        if (c == '\'' || c == '"') {
            in_quote = c;
            current += c;
            continue;
        }

        if (c == '{') { brace_depth++; current += c; continue; }
        if (c == '}') {
            brace_depth--;
            current += c;
            if (brace_depth <= 0) {
                brace_depth = 0;
                tokens.push_back(current);
                current.clear();
            }
            continue;
        }

        if (brace_depth > 0) { current += c; continue; }

        if (c == ' ' || c == '\t') {
            if (!current.empty()) {
                tokens.push_back(current);
                current.clear();
            }
            continue;
        }

        current += c;
    }
    if (!current.empty()) tokens.push_back(current);
    return tokens;
}

// v2.4.4 — offline sub-db placement audit / repair.
//
// `verify-subdbs` is read-only and safe against a running server.
// `reconcile-subdbs --apply` rewrites document placement and REQUIRES the
// service to be stopped: it holds an LMDB write txn over the whole env.
int cmdSubdbs(const std::string& command, const std::vector<std::string>& params) {
    std::string env_path, project;
    bool apply = false;
    bool stamp = false;

    for (size_t i = 0; i < params.size(); ++i) {
        const std::string& p = params[i];
        if (p == "--env" && i + 1 < params.size()) {
            env_path = params[++i];
        } else if (p == "--project" && i + 1 < params.size()) {
            project = params[++i];
        } else if (p == "--apply") {
            apply = true;
        } else if (p == "--stamp-identity") {
            stamp = true;
        } else {
            printError("unknown argument: " + p);
            return 2;
        }
    }

    if (env_path.empty()) {
        printError("--env <path> is required (e.g. "
                   "/var/lib/smartbotic-database/projects/default/env)");
        return 2;
    }
    if (command == "verify-subdbs" && apply) {
        printError("--apply is not valid for verify-subdbs; use reconcile-subdbs");
        return 2;
    }

    try {
        auto report = smartbotic::db::storage::audit(env_path, project);
        smartbotic::db::storage::print_audit(report);

        if (command == "verify-subdbs") {
            return report.misplaced.empty() ? 0 : 1;
        }

        if (!apply) {
            std::cout << "\n" << C_YELLOW << "Dry run." << C_RESET
                      << " Re-run with --apply to perform the repair.\n"
                      << C_DIM
                      << "Stop the service first, and take a backup: the repair "
                         "rewrites document placement in place.\n"
                      << C_RESET;
            return report.misplaced.empty() ? 0 : 1;
        }

        if (stamp) {
            std::cout << "\n" << C_YELLOW << "--stamp-identity:" << C_RESET
                      << " writing sentinels. This REQUIRES the server binary to be"
                         " v2.4.4 or newer;\n  earlier builds do not skip the sentinel"
                         " and will fail on scan.\n";
        }
        std::cout << "\nApplying...\n";
        auto res = smartbotic::db::storage::repair(env_path, report, stamp);
        std::cout << C_GREEN << "done." << C_RESET
                  << "  moved=" << res.moved
                  << " quarantined=" << res.quarantined
                  << " stamped=" << res.stamped << "\n";
        for (const auto& e : res.errors) {
            std::cerr << C_RED << "  warn: " << C_RESET << e << "\n";
        }
        return res.errors.empty() ? 0 : 1;
    } catch (const std::exception& e) {
        printError(e.what());
        return 1;
    }
}

} // anonymous namespace

int main(int argc, char* argv[]) {
    Args args;

    // Parse flags
    std::vector<std::string> positional;
    for (int i = 1; i < argc; ++i) {
        std::string arg = argv[i];
        if (arg == "--address" && i + 1 < argc) {
            args.address = argv[++i];
        } else if (arg == "--help" || arg == "-h") {
            printUsage();
            return 0;
        } else {
            positional.push_back(arg);
        }
    }

    if (!positional.empty()) {
        args.command = positional[0];
        args.params.assign(positional.begin() + 1, positional.end());
    }

    // Offline helpers — these don't need a running server, so dispatch
    // them before opening the gRPC channel.
    if (args.command == "generate-auth-key") {
        return cmdGenerateAuthKey();
    }
    if (args.command == "generate-tls-cert") {
        return cmdGenerateTlsCert(args.params);
    }
    if (args.command == "verify-subdbs" || args.command == "reconcile-subdbs") {
        return cmdSubdbs(args.command, args.params);
    }

    // Connect
    smartbotic::database::Client client({.address = args.address});
    client.connect();

    // Scriptable mode: single command
    if (!args.command.empty()) {
        return execCommand(client, args.command, args.params) ? 0 : 1;
    }

    // Interactive mode
    std::cout << C_BOLD << "smartbotic-db-cli" << C_RESET
              << " connected to " << C_CYAN << args.address << C_RESET << "\n"
              << C_DIM << "Type 'help' for commands, 'exit' to quit." << C_RESET << "\n";

    std::string line;
    while (true) {
        std::cout << C_YELLOW << "> " << C_RESET;
        if (!std::getline(std::cin, line)) break;

        // Trim
        auto start = line.find_first_not_of(" \t");
        if (start == std::string::npos) continue;
        line = line.substr(start);

        if (line == "exit" || line == "quit" || line == "q") break;
        if (line.empty()) continue;

        auto tokens = tokenize(line);
        if (tokens.empty()) continue;

        auto cmd = tokens[0];
        std::vector<std::string> params(tokens.begin() + 1, tokens.end());
        execCommand(client, cmd, params);
    }

    return 0;
}