smartdb/rust/crates/rustdb/src/management.rs

use anyhow::Result;
use bson::{Bson, Document};
use serde::{Deserialize, Serialize};
use tokio::io::{AsyncBufReadExt, BufReader};
use tracing::{info, error};

use crate::RustDb;
use rustdb_config::{RustDbOptions, StorageType};

/// A management request from the TypeScript wrapper.
#[derive(Debug, Deserialize)]
pub struct ManagementRequest {
    pub id: String,
    pub method: String,
    #[serde(default)]
    pub params: serde_json::Value,
}

/// A management response back to the TypeScript wrapper.
#[derive(Debug, Serialize)]
pub struct ManagementResponse {
    pub id: String,
    pub success: bool,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub result: Option<serde_json::Value>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub error: Option<String>,
}

/// An unsolicited event from the server to the TypeScript wrapper.
#[derive(Debug, Serialize)]
pub struct ManagementEvent {
    pub event: String,
    pub data: serde_json::Value,
}

impl ManagementResponse {
    fn ok(id: String, result: serde_json::Value) -> Self {
        Self {
            id,
            success: true,
            result: Some(result),
            error: None,
        }
    }

    fn err(id: String, message: String) -> Self {
        Self {
            id,
            success: false,
            result: None,
            error: Some(message),
        }
    }
}

fn send_line(line: &str) {
    use std::io::Write;
    let stdout = std::io::stdout();
    let mut handle = stdout.lock();
    let _ = handle.write_all(line.as_bytes());
    let _ = handle.write_all(b"\n");
    let _ = handle.flush();
}

fn send_response(response: &ManagementResponse) {
    match serde_json::to_string(response) {
        Ok(json) => send_line(&json),
        Err(e) => error!("Failed to serialize management response: {}", e),
    }
}

fn send_event(event: &str, data: serde_json::Value) {
    let evt = ManagementEvent {
        event: event.to_string(),
        data,
    };
    match serde_json::to_string(&evt) {
        Ok(json) => send_line(&json),
        Err(e) => error!("Failed to serialize management event: {}", e),
    }
}

/// Run the management loop, reading JSON commands from stdin and writing responses to stdout.
pub async fn management_loop() -> Result<()> {
    let stdin = BufReader::new(tokio::io::stdin());
    let mut lines = stdin.lines();
    let mut db: Option<RustDb> = None;

    send_event("ready", serde_json::json!({}));

    loop {
        let line = match lines.next_line().await {
            Ok(Some(line)) => line,
            Ok(None) => {
                // stdin closed - parent process exited
                info!("Management stdin closed, shutting down");
                if let Some(ref mut d) = db {
                    let _ = d.stop().await;
                }
                break;
            }
            Err(e) => {
                error!("Error reading management stdin: {}", e);
                break;
            }
        };

        let line = line.trim().to_string();
        if line.is_empty() {
            continue;
        }

        let request: ManagementRequest = match serde_json::from_str(&line) {
            Ok(r) => r,
            Err(e) => {
                error!("Failed to parse management request: {}", e);
                send_response(&ManagementResponse::err(
                    "unknown".to_string(),
                    format!("Failed to parse request: {}", e),
                ));
                continue;
            }
        };

        let response = handle_request(&request, &mut db).await;
        send_response(&response);
    }

    Ok(())
}

async fn handle_request(
    request: &ManagementRequest,
    db: &mut Option<RustDb>,
) -> ManagementResponse {
    let id = request.id.clone();

    match request.method.as_str() {
        "start" => handle_start(&id, &request.params, db).await,
        "stop" => handle_stop(&id, db).await,
        "getStatus" => handle_get_status(&id, db),
        "getHealth" => handle_get_health(&id, db).await,
        "getMetrics" => handle_get_metrics(&id, db).await,
        "createDatabaseTenant" => handle_create_database_tenant(&id, &request.params, db).await,
        "deleteDatabaseTenant" => handle_delete_database_tenant(&id, &request.params, db).await,
        "rotateDatabaseTenantPassword" => {
            handle_rotate_database_tenant_password(&id, &request.params, db).await
        }
        "listDatabaseTenants" => handle_list_database_tenants(&id, db),
        "getDatabaseTenantDescriptor" => {
            handle_get_database_tenant_descriptor(&id, &request.params, db)
        }
        "exportDatabase" => handle_export_database(&id, &request.params, db).await,
        "importDatabase" => handle_import_database(&id, &request.params, db).await,
        "getOpLog" => handle_get_oplog(&id, &request.params, db),
        "getOpLogStats" => handle_get_oplog_stats(&id, db),
        "revertToSeq" => handle_revert_to_seq(&id, &request.params, db).await,
        "getCollections" => handle_get_collections(&id, &request.params, db).await,
        "getDocuments" => handle_get_documents(&id, &request.params, db).await,
        _ => ManagementResponse::err(id, format!("Unknown method: {}", request.method)),
    }
}

async fn handle_start(
    id: &str,
    params: &serde_json::Value,
    db: &mut Option<RustDb>,
) -> ManagementResponse {
    if db.is_some() {
        return ManagementResponse::err(id.to_string(), "Server is already running".to_string());
    }

    let config = match params.get("config") {
        Some(config) => config,
        None => return ManagementResponse::err(id.to_string(), "Missing 'config' parameter".to_string()),
    };

    let options: RustDbOptions = match serde_json::from_value(config.clone()) {
        Ok(o) => o,
        Err(e) => return ManagementResponse::err(id.to_string(), format!("Invalid config: {}", e)),
    };
    if let Err(e) = options.validate() {
        return ManagementResponse::err(id.to_string(), format!("Invalid config: {}", e));
    }

    let connection_uri = options.connection_uri();

    match RustDb::new(options).await {
        Ok(mut d) => {
            match d.start().await {
                Ok(()) => {
                    send_event("started", serde_json::json!({}));
                    *db = Some(d);
                    ManagementResponse::ok(
                        id.to_string(),
                        serde_json::json!({ "connectionUri": connection_uri }),
                    )
                }
                Err(e) => {
                    send_event("error", serde_json::json!({"message": format!("{}", e)}));
                    ManagementResponse::err(id.to_string(), format!("Failed to start: {}", e))
                }
            }
        }
        Err(e) => ManagementResponse::err(id.to_string(), format!("Failed to create server: {}", e)),
    }
}

async fn handle_stop(
    id: &str,
    db: &mut Option<RustDb>,
) -> ManagementResponse {
    match db.as_mut() {
        Some(d) => {
            match d.stop().await {
                Ok(()) => {
                    *db = None;
                    send_event("stopped", serde_json::json!({}));
                    ManagementResponse::ok(id.to_string(), serde_json::json!({}))
                }
                Err(e) => ManagementResponse::err(id.to_string(), format!("Failed to stop: {}", e)),
            }
        }
        None => ManagementResponse::ok(id.to_string(), serde_json::json!({})),
    }
}

fn handle_get_status(
    id: &str,
    db: &Option<RustDb>,
) -> ManagementResponse {
    match db.as_ref() {
        Some(_d) => ManagementResponse::ok(
            id.to_string(),
            serde_json::json!({
                "running": true,
            }),
        ),
        None => ManagementResponse::ok(
            id.to_string(),
            serde_json::json!({ "running": false }),
        ),
    }
}

async fn handle_get_health(id: &str, db: &Option<RustDb>) -> ManagementResponse {
    match db.as_ref() {
        Some(d) => {
            let ctx = d.ctx();
            let (database_count, collection_count) = database_and_collection_counts(ctx).await;
            let options = d.options();
            let storage = match &options.storage {
                StorageType::Memory => "memory",
                StorageType::File => "file",
            };
            ManagementResponse::ok(
                id.to_string(),
                serde_json::json!({
                    "running": true,
                    "storage": storage,
                    "storagePath": options.storage_path.clone().or_else(|| options.persist_path.clone()),
                    "authEnabled": ctx.auth.enabled(),
                    "authUsers": ctx.auth.user_count(),
                    "usersPathConfigured": options.auth.users_path.is_some(),
                    "databaseCount": database_count,
                    "collectionCount": collection_count,
                    "uptimeSeconds": ctx.start_time.elapsed().as_secs(),
                }),
            )
        }
        None => ManagementResponse::ok(
            id.to_string(),
            serde_json::json!({
                "running": false,
                "databaseCount": 0,
                "collectionCount": 0,
            }),
        ),
    }
}

async fn handle_get_metrics(
    id: &str,
    db: &Option<RustDb>,
) -> ManagementResponse {
    match db.as_ref() {
        Some(d) => {
            let ctx = d.ctx();
            let db_list = ctx.storage.list_databases().await.unwrap_or_default();
            let mut total_collections = 0u64;
            for db_name in &db_list {
                if let Ok(colls) = ctx.storage.list_collections(db_name).await {
                    total_collections += colls.len() as u64;
                }
            }
            let oplog_stats = ctx.oplog.stats();
            let uptime_secs = ctx.start_time.elapsed().as_secs();

            ManagementResponse::ok(
                id.to_string(),
                serde_json::json!({
                    "databases": db_list.len(),
                    "collections": total_collections,
                    "oplogEntries": oplog_stats.total_entries,
                    "oplogCurrentSeq": oplog_stats.current_seq,
                    "sessions": ctx.sessions.len(),
                    "activeTransactions": ctx.transactions.len(),
                    "authEnabled": ctx.auth.enabled(),
                    "authUsers": ctx.auth.user_count(),
                    "uptimeSeconds": uptime_secs,
                }),
            )
        }
        None => ManagementResponse::err(id.to_string(), "Server is not running".to_string()),
    }
}

async fn handle_create_database_tenant(
    id: &str,
    params: &serde_json::Value,
    db: &Option<RustDb>,
) -> ManagementResponse {
    let d = match db.as_ref() {
        Some(d) => d,
        None => return ManagementResponse::err(id.to_string(), "Server is not running".to_string()),
    };
    let ctx = d.ctx();
    if !ctx.auth.enabled() {
        return ManagementResponse::err(
            id.to_string(),
            "Authentication must be enabled to create database tenants".to_string(),
        );
    }

    let database_name = match string_param(params, "databaseName") {
        Ok(value) => value,
        Err(message) => return ManagementResponse::err(id.to_string(), message),
    };
    if let Err(message) = validate_database_name(database_name) {
        return ManagementResponse::err(id.to_string(), message);
    }
    let username = match string_param(params, "username") {
        Ok(value) => value,
        Err(message) => return ManagementResponse::err(id.to_string(), message),
    };
    if let Err(message) = validate_username(username) {
        return ManagementResponse::err(id.to_string(), message);
    }
    let password = match string_param(params, "password") {
        Ok(value) => value,
        Err(message) => return ManagementResponse::err(id.to_string(), message),
    };
    if password.is_empty() {
        return ManagementResponse::err(id.to_string(), "password must not be empty".to_string());
    }
    let roles = match roles_param(params) {
        Ok(roles) => roles,
        Err(message) => return ManagementResponse::err(id.to_string(), message),
    };

    if let Err(e) = ctx.storage.create_database(database_name).await {
        if !is_already_exists(&e.to_string()) {
            return ManagementResponse::err(
                id.to_string(),
                format!("Failed to create database: {e}"),
            );
        }
    }

    match ctx
        .auth
        .create_user(database_name, username, password, roles)
    {
        Ok(()) => {
            let users = ctx.auth.users_info(database_name, Some(username));
            match users.first() {
                Some(user) => ManagementResponse::ok(id.to_string(), tenant_descriptor_json(user)),
                None => ManagementResponse::err(
                    id.to_string(),
                    "Tenant user was created but could not be read back".to_string(),
                ),
            }
        }
        Err(e) => {
            ManagementResponse::err(id.to_string(), format!("Failed to create tenant user: {e}"))
        }
    }
}

async fn handle_delete_database_tenant(
    id: &str,
    params: &serde_json::Value,
    db: &Option<RustDb>,
) -> ManagementResponse {
    let d = match db.as_ref() {
        Some(d) => d,
        None => {
            return ManagementResponse::err(id.to_string(), "Server is not running".to_string())
        }
    };
    let ctx = d.ctx();
    let database_name = match string_param(params, "databaseName") {
        Ok(value) => value,
        Err(message) => return ManagementResponse::err(id.to_string(), message),
    };
    if let Err(message) = validate_database_name(database_name) {
        return ManagementResponse::err(id.to_string(), message);
    }
    let username = params.get("username").and_then(|v| v.as_str());
    if let Some(username) = username {
        if let Err(message) = validate_username(username) {
            return ManagementResponse::err(id.to_string(), message);
        }
    }

    if let Err(e) = ctx.storage.drop_database(database_name).await {
        return ManagementResponse::err(id.to_string(), format!("Failed to drop database: {e}"));
    }
    remove_database_indexes(ctx, database_name);

    let mut deleted_users = 0usize;
    if ctx.auth.enabled() {
        if let Some(username) = username {
            match ctx.auth.drop_user(database_name, username) {
                Ok(()) => deleted_users = 1,
                Err(rustdb_auth::AuthError::UserNotFound(_)) => deleted_users = 0,
                Err(e) => {
                    return ManagementResponse::err(
                        id.to_string(),
                        format!("Failed to drop tenant user: {e}"),
                    )
                }
            }
        } else {
            match ctx.auth.drop_users_for_database(database_name) {
                Ok(count) => deleted_users = count,
                Err(e) => {
                    return ManagementResponse::err(
                        id.to_string(),
                        format!("Failed to drop tenant users: {e}"),
                    )
                }
            }
        }
    }

    ManagementResponse::ok(
        id.to_string(),
        serde_json::json!({
            "databaseName": database_name,
            "deletedUsers": deleted_users,
            "databaseDropped": true,
        }),
    )
}

async fn handle_rotate_database_tenant_password(
    id: &str,
    params: &serde_json::Value,
    db: &Option<RustDb>,
) -> ManagementResponse {
    let d = match db.as_ref() {
        Some(d) => d,
        None => {
            return ManagementResponse::err(id.to_string(), "Server is not running".to_string())
        }
    };
    let ctx = d.ctx();
    if !ctx.auth.enabled() {
        return ManagementResponse::err(
            id.to_string(),
            "Authentication must be enabled to rotate database tenant passwords".to_string(),
        );
    }

    let username = match string_param(params, "username") {
        Ok(value) => value,
        Err(message) => return ManagementResponse::err(id.to_string(), message),
    };
    if let Err(message) = validate_username(username) {
        return ManagementResponse::err(id.to_string(), message);
    }
    let password = match string_param(params, "password") {
        Ok(value) => value,
        Err(message) => return ManagementResponse::err(id.to_string(), message),
    };
    if password.is_empty() {
        return ManagementResponse::err(id.to_string(), "password must not be empty".to_string());
    }

    let matches: Vec<_> = ctx
        .auth
        .list_users()
        .into_iter()
        .filter(|user| user.username == username)
        .collect();
    if matches.is_empty() {
        return ManagementResponse::err(
            id.to_string(),
            format!("tenant user not found: {username}"),
        );
    }
    if matches.len() > 1 {
        return ManagementResponse::err(
            id.to_string(),
            format!("tenant username is ambiguous across databases: {username}"),
        );
    }
    let user = &matches[0];
    match ctx
        .auth
        .update_user(&user.database, username, Some(password), None)
    {
        Ok(()) => {
            let users = ctx.auth.users_info(&user.database, Some(username));
            match users.first() {
                Some(user) => ManagementResponse::ok(id.to_string(), tenant_descriptor_json(user)),
                None => ManagementResponse::err(
                    id.to_string(),
                    "Tenant user was updated but could not be read back".to_string(),
                ),
            }
        }
        Err(e) => ManagementResponse::err(
            id.to_string(),
            format!("Failed to rotate tenant password: {e}"),
        ),
    }
}

fn handle_list_database_tenants(id: &str, db: &Option<RustDb>) -> ManagementResponse {
    let d = match db.as_ref() {
        Some(d) => d,
        None => {
            return ManagementResponse::err(id.to_string(), "Server is not running".to_string())
        }
    };
    let tenants: Vec<serde_json::Value> = d
        .ctx()
        .auth
        .list_users()
        .into_iter()
        .filter(|user| user.database != "admin")
        .map(|user| tenant_descriptor_json(&user))
        .collect();
    ManagementResponse::ok(id.to_string(), serde_json::json!({ "tenants": tenants }))
}

fn handle_get_database_tenant_descriptor(
    id: &str,
    params: &serde_json::Value,
    db: &Option<RustDb>,
) -> ManagementResponse {
    let d = match db.as_ref() {
        Some(d) => d,
        None => {
            return ManagementResponse::err(id.to_string(), "Server is not running".to_string())
        }
    };
    let database_name = match string_param(params, "databaseName") {
        Ok(value) => value,
        Err(message) => return ManagementResponse::err(id.to_string(), message),
    };
    let username = match string_param(params, "username") {
        Ok(value) => value,
        Err(message) => return ManagementResponse::err(id.to_string(), message),
    };
    let users = d.ctx().auth.users_info(database_name, Some(username));
    match users.first() {
        Some(user) => ManagementResponse::ok(id.to_string(), tenant_descriptor_json(user)),
        None => ManagementResponse::err(
            id.to_string(),
            format!("tenant user not found: {database_name}.{username}"),
        ),
    }
}

async fn handle_export_database(
    id: &str,
    params: &serde_json::Value,
    db: &Option<RustDb>,
) -> ManagementResponse {
    let d = match db.as_ref() {
        Some(d) => d,
        None => {
            return ManagementResponse::err(id.to_string(), "Server is not running".to_string())
        }
    };
    let ctx = d.ctx();
    let database_name = match string_param(params, "databaseName") {
        Ok(value) => value,
        Err(message) => return ManagementResponse::err(id.to_string(), message),
    };
    if let Err(message) = validate_database_name(database_name) {
        return ManagementResponse::err(id.to_string(), message);
    }
    match ctx.storage.database_exists(database_name).await {
        Ok(true) => {}
        Ok(false) => {
            return ManagementResponse::err(
                id.to_string(),
                format!("database not found: {database_name}"),
            )
        }
        Err(e) => {
            return ManagementResponse::err(
                id.to_string(),
                format!("Failed to check database: {e}"),
            )
        }
    }

    let collection_names = match ctx.storage.list_collections(database_name).await {
        Ok(collections) => collections,
        Err(e) => {
            return ManagementResponse::err(
                id.to_string(),
                format!("Failed to list collections: {e}"),
            )
        }
    };
    let mut collections = Vec::with_capacity(collection_names.len());
    for collection_name in collection_names {
        let documents = match ctx.storage.find_all(database_name, &collection_name).await {
            Ok(docs) => docs
                .into_iter()
                .map(|doc| bson_doc_to_json(&doc))
                .collect::<Vec<_>>(),
            Err(e) => {
                return ManagementResponse::err(
                    id.to_string(),
                    format!("Failed to export collection '{collection_name}': {e}"),
                )
            }
        };
        let indexes = match ctx
            .storage
            .get_indexes(database_name, &collection_name)
            .await
        {
            Ok(specs) => specs
                .into_iter()
                .map(|doc| bson_doc_to_json(&doc))
                .collect::<Vec<_>>(),
            Err(_) => Vec::new(),
        };
        collections.push(serde_json::json!({
            "name": collection_name,
            "documents": documents,
            "indexes": indexes,
        }));
    }

    ManagementResponse::ok(
        id.to_string(),
        serde_json::json!({
            "format": "smartdb.database.export.v1",
            "databaseName": database_name,
            "exportedAtMs": now_ms(),
            "collections": collections,
        }),
    )
}

async fn handle_import_database(
    id: &str,
    params: &serde_json::Value,
    db: &Option<RustDb>,
) -> ManagementResponse {
    let d = match db.as_ref() {
        Some(d) => d,
        None => {
            return ManagementResponse::err(id.to_string(), "Server is not running".to_string())
        }
    };
    let ctx = d.ctx();
    let database_name = match string_param(params, "databaseName") {
        Ok(value) => value,
        Err(message) => return ManagementResponse::err(id.to_string(), message),
    };
    if let Err(message) = validate_database_name(database_name) {
        return ManagementResponse::err(id.to_string(), message);
    }
    let source = match params.get("source") {
        Some(value) => value,
        None => {
            return ManagementResponse::err(
                id.to_string(),
                "Missing 'source' parameter".to_string(),
            )
        }
    };
    let source_collections = match source.get("collections").and_then(|value| value.as_array()) {
        Some(collections) => collections,
        None => {
            return ManagementResponse::err(
                id.to_string(),
                "source.collections must be an array".to_string(),
            )
        }
    };

    if let Err(e) = ctx.storage.drop_database(database_name).await {
        return ManagementResponse::err(
            id.to_string(),
            format!("Failed to clear database before import: {e}"),
        );
    }
    remove_database_indexes(ctx, database_name);
    if let Err(e) = ctx.storage.create_database(database_name).await {
        if !is_already_exists(&e.to_string()) {
            return ManagementResponse::err(
                id.to_string(),
                format!("Failed to create database: {e}"),
            );
        }
    }

    let mut imported_collections = 0usize;
    let mut imported_documents = 0usize;
    for collection in source_collections {
        let collection_name = match collection.get("name").and_then(|value| value.as_str()) {
            Some(value) => value,
            None => {
                return ManagementResponse::err(
                    id.to_string(),
                    "source collection is missing a string 'name'".to_string(),
                )
            }
        };
        if let Err(message) = validate_collection_name(collection_name) {
            return ManagementResponse::err(id.to_string(), message);
        }
        if let Err(e) = ctx
            .storage
            .create_collection(database_name, collection_name)
            .await
        {
            if !is_already_exists(&e.to_string()) {
                return ManagementResponse::err(
                    id.to_string(),
                    format!("Failed to create collection '{collection_name}': {e}"),
                );
            }
        }

        if let Some(documents) = collection
            .get("documents")
            .and_then(|value| value.as_array())
        {
            for document_value in documents {
                let document = match json_to_bson_doc(document_value) {
                    Ok(document) => document,
                    Err(message) => {
                        return ManagementResponse::err(
                            id.to_string(),
                            format!("Invalid document in '{collection_name}': {message}"),
                        )
                    }
                };
                if let Err(e) = ctx
                    .storage
                    .insert_one(database_name, collection_name, document)
                    .await
                {
                    return ManagementResponse::err(
                        id.to_string(),
                        format!("Failed to import document into '{collection_name}': {e}"),
                    );
                }
                imported_documents += 1;
            }
        }

        if let Some(indexes) = collection.get("indexes").and_then(|value| value.as_array()) {
            for index_value in indexes {
                let index_doc = match json_to_bson_doc(index_value) {
                    Ok(document) => document,
                    Err(message) => {
                        return ManagementResponse::err(
                            id.to_string(),
                            format!("Invalid index in '{collection_name}': {message}"),
                        )
                    }
                };
                let name = index_doc.get_str("name").unwrap_or("_id_").to_string();
                if let Err(e) = ctx
                    .storage
                    .save_index(database_name, collection_name, &name, index_doc)
                    .await
                {
                    return ManagementResponse::err(
                        id.to_string(),
                        format!("Failed to import index '{name}' into '{collection_name}': {e}"),
                    );
                }
            }
        }

        imported_collections += 1;
    }

    ManagementResponse::ok(
        id.to_string(),
        serde_json::json!({
            "databaseName": database_name,
            "collections": imported_collections,
            "documents": imported_documents,
        }),
    )
}

fn handle_get_oplog(
    id: &str,
    params: &serde_json::Value,
    db: &Option<RustDb>,
) -> ManagementResponse {
    let d = match db.as_ref() {
        Some(d) => d,
        None => {
            return ManagementResponse::err(id.to_string(), "Server is not running".to_string())
        }
    };

    let ctx = d.ctx();
    let since_seq = params.get("sinceSeq").and_then(|v| v.as_u64()).unwrap_or(1);
    let limit = params.get("limit").and_then(|v| v.as_u64()).unwrap_or(100) as usize;
    let filter_db = params.get("db").and_then(|v| v.as_str());
    let filter_coll = params.get("collection").and_then(|v| v.as_str());

    let mut entries = ctx.oplog.entries_since(since_seq);

    // Apply filters.
    if let Some(fdb) = filter_db {
        entries.retain(|e| e.db == fdb);
    }
    if let Some(fcoll) = filter_coll {
        entries.retain(|e| e.collection == fcoll);
    }

    let total = entries.len();
    entries.truncate(limit);

    // Serialize entries to JSON.
    let entries_json: Vec<serde_json::Value> = entries
        .iter()
        .map(|e| {
            let doc_json = e.document.as_ref().map(|d| bson_doc_to_json(d));
            let prev_json = e.previous_document.as_ref().map(|d| bson_doc_to_json(d));
            serde_json::json!({
                "seq": e.seq,
                "timestampMs": e.timestamp_ms,
                "op": match e.op {
                    rustdb_storage::OpType::Insert => "insert",
                    rustdb_storage::OpType::Update => "update",
                    rustdb_storage::OpType::Delete => "delete",
                },
                "db": e.db,
                "collection": e.collection,
                "documentId": e.document_id,
                "document": doc_json,
                "previousDocument": prev_json,
            })
        })
        .collect();

    ManagementResponse::ok(
        id.to_string(),
        serde_json::json!({
            "entries": entries_json,
            "currentSeq": ctx.oplog.current_seq(),
            "totalEntries": total,
        }),
    )
}

fn handle_get_oplog_stats(
    id: &str,
    db: &Option<RustDb>,
) -> ManagementResponse {
    let d = match db.as_ref() {
        Some(d) => d,
        None => return ManagementResponse::err(id.to_string(), "Server is not running".to_string()),
    };

    let stats = d.ctx().oplog.stats();
    ManagementResponse::ok(
        id.to_string(),
        serde_json::json!({
            "currentSeq": stats.current_seq,
            "totalEntries": stats.total_entries,
            "oldestSeq": stats.oldest_seq,
            "entriesByOp": {
                "insert": stats.inserts,
                "update": stats.updates,
                "delete": stats.deletes,
            },
        }),
    )
}

async fn handle_revert_to_seq(
    id: &str,
    params: &serde_json::Value,
    db: &Option<RustDb>,
) -> ManagementResponse {
    let d = match db.as_ref() {
        Some(d) => d,
        None => return ManagementResponse::err(id.to_string(), "Server is not running".to_string()),
    };

    let target_seq = match params.get("seq").and_then(|v| v.as_u64()) {
        Some(s) => s,
        None => return ManagementResponse::err(id.to_string(), "Missing 'seq' parameter".to_string()),
    };
    let dry_run = params.get("dryRun").and_then(|v| v.as_bool()).unwrap_or(false);

    let ctx = d.ctx();
    let current = ctx.oplog.current_seq();

    if target_seq > current {
        return ManagementResponse::err(
            id.to_string(),
            format!("Target seq {} is beyond current seq {}", target_seq, current),
        );
    }

    // Collect entries to revert (from target+1 to current), sorted descending for reverse processing.
    let mut entries_to_revert = ctx.oplog.entries_range(target_seq + 1, current);
    entries_to_revert.reverse();

    if dry_run {
        let entries_json: Vec<serde_json::Value> = entries_to_revert
            .iter()
            .map(|e| {
                serde_json::json!({
                    "seq": e.seq,
                    "op": match e.op {
                        rustdb_storage::OpType::Insert => "insert",
                        rustdb_storage::OpType::Update => "update",
                        rustdb_storage::OpType::Delete => "delete",
                    },
                    "db": e.db,
                    "collection": e.collection,
                    "documentId": e.document_id,
                })
            })
            .collect();

        return ManagementResponse::ok(
            id.to_string(),
            serde_json::json!({
                "dryRun": true,
                "reverted": entries_to_revert.len(),
                "entries": entries_json,
            }),
        );
    }

    // Execute revert: process each entry in reverse, using storage directly.
    let mut reverted = 0u64;
    let mut errors: Vec<String> = Vec::new();

    for entry in &entries_to_revert {
        let result = match entry.op {
            rustdb_storage::OpType::Insert => {
                // Undo insert -> delete the document.
                ctx.storage.delete_by_id(&entry.db, &entry.collection, &entry.document_id).await
            }
            rustdb_storage::OpType::Update => {
                // Undo update -> restore the previous document.
                if let Some(ref prev_doc) = entry.previous_document {
                    ctx.storage
                        .update_by_id(&entry.db, &entry.collection, &entry.document_id, prev_doc.clone())
                        .await
                } else {
                    errors.push(format!("seq {}: update entry missing previous_document", entry.seq));
                    continue;
                }
            }
            rustdb_storage::OpType::Delete => {
                // Undo delete -> re-insert the previous document.
                if let Some(ref prev_doc) = entry.previous_document {
                    ctx.storage
                        .insert_one(&entry.db, &entry.collection, prev_doc.clone())
                        .await
                        .map(|_| ())
                } else {
                    errors.push(format!("seq {}: delete entry missing previous_document", entry.seq));
                    continue;
                }
            }
        };

        match result {
            Ok(()) => reverted += 1,
            Err(e) => errors.push(format!("seq {}: {}", entry.seq, e)),
        }
    }

    // Truncate the oplog to the target sequence.
    ctx.oplog.truncate_after(target_seq);

    let mut response = serde_json::json!({
        "dryRun": false,
        "reverted": reverted,
        "targetSeq": target_seq,
    });

    if !errors.is_empty() {
        response["errors"] = serde_json::json!(errors);
    }

    ManagementResponse::ok(id.to_string(), response)
}

async fn handle_get_collections(
    id: &str,
    params: &serde_json::Value,
    db: &Option<RustDb>,
) -> ManagementResponse {
    let d = match db.as_ref() {
        Some(d) => d,
        None => return ManagementResponse::err(id.to_string(), "Server is not running".to_string()),
    };

    let ctx = d.ctx();
    let filter_db = params.get("db").and_then(|v| v.as_str());

    let databases = match ctx.storage.list_databases().await {
        Ok(dbs) => dbs,
        Err(e) => return ManagementResponse::err(id.to_string(), format!("Failed to list databases: {}", e)),
    };

    let mut collections: Vec<serde_json::Value> = Vec::new();

    for db_name in &databases {
        if let Some(fdb) = filter_db {
            if db_name != fdb {
                continue;
            }
        }

        if let Ok(colls) = ctx.storage.list_collections(db_name).await {
            for coll_name in colls {
                let count = ctx.storage.count(db_name, &coll_name).await.unwrap_or(0);
                collections.push(serde_json::json!({
                    "db": db_name,
                    "name": coll_name,
                    "count": count,
                }));
            }
        }
    }

    ManagementResponse::ok(
        id.to_string(),
        serde_json::json!({ "collections": collections }),
    )
}

async fn handle_get_documents(
    id: &str,
    params: &serde_json::Value,
    db: &Option<RustDb>,
) -> ManagementResponse {
    let d = match db.as_ref() {
        Some(d) => d,
        None => return ManagementResponse::err(id.to_string(), "Server is not running".to_string()),
    };

    let db_name = match params.get("db").and_then(|v| v.as_str()) {
        Some(s) => s,
        None => return ManagementResponse::err(id.to_string(), "Missing 'db' parameter".to_string()),
    };
    let coll_name = match params.get("collection").and_then(|v| v.as_str()) {
        Some(s) => s,
        None => return ManagementResponse::err(id.to_string(), "Missing 'collection' parameter".to_string()),
    };
    let limit = params.get("limit").and_then(|v| v.as_u64()).unwrap_or(50) as usize;
    let skip = params.get("skip").and_then(|v| v.as_u64()).unwrap_or(0) as usize;

    let ctx = d.ctx();

    let all_docs = match ctx.storage.find_all(db_name, coll_name).await {
        Ok(docs) => docs,
        Err(e) => return ManagementResponse::err(id.to_string(), format!("Failed to find documents: {}", e)),
    };

    let total = all_docs.len();
    let docs: Vec<serde_json::Value> = all_docs
        .into_iter()
        .skip(skip)
        .take(limit)
        .map(|d| bson_doc_to_json(&d))
        .collect();

    ManagementResponse::ok(
        id.to_string(),
        serde_json::json!({
            "documents": docs,
            "total": total,
        }),
    )
}

async fn database_and_collection_counts(ctx: &rustdb_commands::CommandContext) -> (usize, u64) {
    let databases = ctx.storage.list_databases().await.unwrap_or_default();
    let mut collections = 0u64;
    for database in &databases {
        if let Ok(database_collections) = ctx.storage.list_collections(database).await {
            collections += database_collections.len() as u64;
        }
    }
    (databases.len(), collections)
}

fn remove_database_indexes(ctx: &rustdb_commands::CommandContext, database_name: &str) {
    let prefix = format!("{}.", database_name);
    let keys_to_remove: Vec<String> = ctx
        .indexes
        .iter()
        .filter(|entry| entry.key().starts_with(&prefix))
        .map(|entry| entry.key().clone())
        .collect();
    for key in keys_to_remove {
        ctx.indexes.remove(&key);
    }
}

fn tenant_descriptor_json(user: &rustdb_auth::AuthenticatedUser) -> serde_json::Value {
    serde_json::json!({
        "databaseName": user.database.clone(),
        "username": user.username.clone(),
        "roles": user.roles.clone(),
        "authSource": user.database.clone(),
    })
}

fn string_param<'a>(params: &'a serde_json::Value, key: &str) -> Result<&'a str, String> {
    params
        .get(key)
        .and_then(|value| value.as_str())
        .ok_or_else(|| format!("Missing '{key}' parameter"))
}

fn roles_param(params: &serde_json::Value) -> Result<Vec<String>, String> {
    let Some(value) = params.get("roles") else {
        return Ok(vec!["readWrite".to_string(), "dbAdmin".to_string()]);
    };
    let roles = value
        .as_array()
        .ok_or_else(|| "roles must be an array of strings".to_string())?;
    let mut result = Vec::with_capacity(roles.len());
    for role in roles {
        let Some(role_name) = role.as_str() else {
            return Err("roles must be an array of strings".to_string());
        };
        if role_name.is_empty() {
            return Err("roles must not contain empty role names".to_string());
        }
        result.push(role_name.to_string());
    }
    Ok(result)
}

fn validate_database_name(name: &str) -> Result<(), String> {
    if name.is_empty() {
        return Err("databaseName must not be empty".to_string());
    }
    if name == "."
        || name == ".."
        || name.contains('/')
        || name.contains('\\')
        || name.contains('\0')
    {
        return Err(format!(
            "databaseName contains invalid path characters: {name}"
        ));
    }
    Ok(())
}

fn validate_collection_name(name: &str) -> Result<(), String> {
    if name.is_empty() {
        return Err("collection name must not be empty".to_string());
    }
    if name == "."
        || name == ".."
        || name.contains('/')
        || name.contains('\\')
        || name.contains('\0')
    {
        return Err(format!(
            "collection name contains invalid path characters: {name}"
        ));
    }
    Ok(())
}

fn validate_username(username: &str) -> Result<(), String> {
    if username.is_empty() {
        return Err("username must not be empty".to_string());
    }
    if username.contains('\0') {
        return Err("username must not contain NUL bytes".to_string());
    }
    Ok(())
}

fn is_already_exists(message: &str) -> bool {
    message.contains("AlreadyExists") || message.contains("already exists")
}

fn json_to_bson_doc(value: &serde_json::Value) -> Result<Document, String> {
    let bson_value: Bson = serde_json::from_value(value.clone()).map_err(|e| e.to_string())?;
    match bson_value {
        Bson::Document(document) => Ok(document),
        _ => Err("expected BSON document".to_string()),
    }
}

fn now_ms() -> u64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap_or_default()
        .as_millis() as u64
}

/// Convert a BSON Document to a serde_json::Value.
fn bson_doc_to_json(doc: &bson::Document) -> serde_json::Value {
    // Use bson's built-in relaxed extended JSON serialization.
    let bson_val = bson::Bson::Document(doc.clone());
    bson_val.into_relaxed_extjson()
}