markbase/markbase-core/src/cli/metadata/db.rs

use anyhow::Context;
use clap::Subcommand;
use rusqlite::Connection;

#[derive(Subcommand)]
pub enum DbCommand {
    #[command(name = "db-create")]
    Create {
        #[arg(short, long)]
        user: String,
    },
    #[command(name = "db-status")]
    Status {
        #[arg(short, long)]
        user: String,
    },
    #[command(name = "db-backup")]
    Backup {
        #[arg(short, long)]
        user: String,
        #[arg(short, long)]
        output: String,
    },
    #[command(name = "db-restore")]
    Restore {
        #[arg(short, long)]
        user: String,
        #[arg(short, long)]
        input: String,
    },
}

pub fn handle_db_command(cmd: DbCommand) -> anyhow::Result<()> {
    match cmd {
        DbCommand::Create { user } => {
            let db_path = filetree::FileTree::user_db_path(&user);

            if std::path::Path::new(&db_path).exists() {
                println!("Database already exists: {}", db_path);
                println!("Use 'db status' to check database info");
                return Ok(());
            }

            println!("Creating database for user: {}", user);

            let conn =
                filetree::FileTree::init_user_db(&user).context("Failed to initialize database")?;

            println!("✓ Database created: {}", db_path);
            println!(
                "✓ Tables initialized: file_nodes, file_registry, file_locations, tree_registry"
            );

            conn.close()
                .map_err(|e| anyhow::anyhow!("Failed to close database: {:?}", e))?;
        }
        DbCommand::Status { user } => {
            let db_path = filetree::FileTree::user_db_path(&user);

            if !std::path::Path::new(&db_path).exists() {
                return Err(anyhow::anyhow!("Database not found: {}", db_path));
            }

            let conn = Connection::open(&db_path).context("Failed to open database")?;

            let file_size = std::fs::metadata(&db_path)?.len();
            let file_size_mb = file_size as f64 / 1024.0 / 1024.0;

            let node_count: i64 = conn
                .query_row("SELECT COUNT(*) FROM file_nodes", [], |row| row.get(0))
                .context("Failed to count nodes")?;

            let file_count: i64 = conn
                .query_row("SELECT COUNT(*) FROM file_registry", [], |row| row.get(0))
                .context("Failed to count files")?;

            let tree_types: Vec<String> = {
                let mut stmt = conn.prepare("SELECT tree_type FROM tree_registry")?;
                let rows = stmt.query_map([], |row| row.get(0))?;
                rows.collect::<Result<Vec<_>, _>>()?
            };

            println!("=== Database Status ===");
            println!("User: {}", user);
            println!("Path: {}", db_path);
            println!("Size: {:.2} MB", file_size_mb);
            println!("Nodes: {}", node_count);
            println!("Files: {}", file_count);
            println!("Tree Types: {:?}", tree_types);

            conn.close()
                .map_err(|e| anyhow::anyhow!("Failed to close database: {:?}", e))?;
        }
        DbCommand::Backup { user, output } => {
            let db_path = filetree::FileTree::user_db_path(&user);

            if !std::path::Path::new(&db_path).exists() {
                return Err(anyhow::anyhow!("Database not found: {}", db_path));
            }

            println!("Backing up database for user: {} to {}", user, output);

            std::fs::copy(&db_path, &output).context("Failed to backup database")?;

            println!("✓ Database backed up to: {}", output);
            println!("✓ Backup size: {} bytes", std::fs::metadata(&output)?.len());
        }
        DbCommand::Restore { user, input } => {
            if !std::path::Path::new(&input).exists() {
                return Err(anyhow::anyhow!("Backup file not found: {}", input));
            }

            let db_path = filetree::FileTree::user_db_path(&user);

            if std::path::Path::new(&db_path).exists() {
                let backup_path = format!("{}.bak", db_path);
                println!("Warning: Database exists, creating backup: {}", backup_path);
                std::fs::copy(&db_path, &backup_path)
                    .context("Failed to create backup before restore")?;
            }

            println!("Restoring database for user: {} from {}", user, input);

            std::fs::copy(&input, &db_path).context("Failed to restore database")?;

            println!("✓ Database restored from: {}", input);
            println!(
                "✓ Database size: {} bytes",
                std::fs::metadata(&db_path)?.len()
            );
        }
    }
    Ok(())
}