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

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

#[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(())
}