stoolap 0.4.0

// Copyright 2025 Stoolap Contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Stoolap CLI - Interactive SQL database command-line interface
//!

use std::fs::File;
use std::io::{self, BufRead, BufReader, IsTerminal};
use std::time::Instant;

use clap::Parser;
use comfy_table::{presets::UTF8_FULL_CONDENSED, Cell, ContentArrangement, Table};
use rustyline::error::ReadlineError;
use rustyline::history::DefaultHistory;
use rustyline::{Config, DefaultEditor, EditMode, Editor};

use stoolap::api::{Database, Transaction as ApiTransaction};
use stoolap::common::version::{MAJOR, MINOR, PATCH};
use stoolap::Value;

/// Version string constant
const VERSION: &str = concat!(
    env!("CARGO_PKG_VERSION_MAJOR"),
    ".",
    env!("CARGO_PKG_VERSION_MINOR"),
    ".",
    env!("CARGO_PKG_VERSION_PATCH")
);

/// Stoolap SQL Database CLI
#[derive(Parser, Debug)]
#[command(name = "stoolap")]
#[command(author = "Stoolap Contributors")]
#[command(version = VERSION)]
#[command(about = "High-performance embedded SQL database with MVCC")]
#[command(
    long_about = "Stoolap is a high-performance embedded SQL database with MVCC transactions.\n\
This CLI provides an interactive interface to execute SQL queries and manage your database.\n\n\
PERSISTENCE DSN PARAMETERS:\n\
  file:///path/to/db?param=value&param2=value2\n\n\
  sync_mode=none|normal|full   Fsync mode (default: normal)\n\
    none: no fsync, durable at checkpoint only\n\
    normal: fsync every 1 second, DDL fsyncs immediately\n\
    full: fsync on every write\n\
  checkpoint_interval=SECS    Checkpoint interval in seconds (default: 60)\n\
  compact_threshold=COUNT     Sub-target volumes per table before merging (default: 4)\n\
  compression=on|off          WAL + volume LZ4 compression (default: on)\n\
  wal_compression=on|off      WAL compression only (default: on)\n\
  volume_compression=on|off   Volume LZ4 compression only (default: on)\n\
  keep_snapshots=COUNT        Backup snapshots per table (default: 3)\n\
  checkpoint_on_close=on|off  Checkpoint on shutdown (default: on)\n\
  wal_max_size=BYTES          Max WAL file size before rotation (default: 67108864)\n\
  wal_buffer_size=BYTES       WAL buffer size (default: 65536)\n\
  wal_flush_trigger=BYTES     Buffer size to trigger flush (default: 32768)\n\
  commit_batch_size=COUNT     Commits to batch before sync (default: 100)\n\
  sync_interval_ms=MS         Min time between syncs in normal mode (default: 1000)\n\
  compression_threshold=BYTES Min size to compress (default: 64)\n\
  target_volume_rows=ROWS     Target rows per cold volume (default: 1048576)\n\n\
EXAMPLES:\n\
  stoolap -d memory://                                    In-memory database\n\
  stoolap -d file:///tmp/mydb                             Persistent database\n\
  stoolap -d file:///tmp/mydb?sync_mode=full               Maximum durability\n\
  stoolap -d file:///tmp/mydb?sync_mode=none&compression=off  Max performance\n\
  stoolap -d file:///tmp/mydb --profile durable           Use durable preset\n\
  stoolap -d file:///tmp/mydb --sync full --compression off\n\n\
BACKUP & RESTORE:\n\
  stoolap -d file:///tmp/mydb --snapshot                  Create backup snapshot\n\
  stoolap -d file:///tmp/mydb --restore                   Restore from latest snapshot\n\
  stoolap -d file:///tmp/mydb --restore 20260315-100000   Restore specific snapshot\n\
  stoolap -d file:///tmp/mydb --reset-volumes --restore   Recovery from corrupted volumes"
)]
struct Args {
    /// Database path (file://<path> or memory://)
    #[arg(short = 'd', long = "db", default_value = "memory://")]
    db_path: String,

    /// Output results in JSON format
    #[arg(short = 'j', long = "json", default_value = "false")]
    json_output: bool,

    /// Suppress connection messages
    #[arg(short = 'q', long = "quiet", default_value = "false")]
    quiet: bool,

    /// Maximum number of rows to display (0 for unlimited)
    #[arg(short = 'l', long = "limit", default_value = "40")]
    limit: usize,

    /// Execute a single SQL statement and exit
    #[arg(short = 'e', long = "execute")]
    execute: Option<String>,

    /// Execute SQL statements from a file
    #[arg(short = 'f', long = "file")]
    file: Option<String>,

    /// WAL sync mode for durability (none, normal, full)
    /// - none: No fsync, data durable at checkpoint only
    /// - normal: Fsync every 1 second, DDL fsyncs immediately (default)
    /// - full: Fsync on every write
    #[arg(short = 's', long = "sync", value_name = "MODE")]
    sync_mode: Option<String>,

    /// Persistence profile preset (fast, normal, durable)
    /// - fast: Optimized for performance, less durable
    /// - normal: Balanced performance and durability (default)
    /// - durable: Maximum durability, slower performance
    #[arg(short = 'p', long = "profile", value_name = "PROFILE")]
    persistence_profile: Option<String>,

    /// Checkpoint interval in seconds (default: 60)
    #[arg(long = "checkpoint-interval", value_name = "SECONDS")]
    checkpoint_interval: Option<u32>,

    /// Sub-target volumes per table before merging (default: 4)
    #[arg(long = "compact-threshold", value_name = "COUNT")]
    compact_threshold: Option<u32>,

    /// Maximum WAL file size in MB before rotation (default: 64)
    #[arg(long = "wal-max-size", value_name = "MB")]
    wal_max_size: Option<u32>,

    /// Enable or disable LZ4 compression for WAL and volumes (default: on)
    #[arg(long = "compression", value_name = "on|off")]
    compression: Option<String>,

    /// Number of backup snapshots to keep per table (default: 3)
    #[arg(long = "keep-snapshots", value_name = "COUNT")]
    keep_snapshots: Option<u32>,

    /// Disable checkpoint on close (for crash simulation in tests)
    #[arg(long = "no-checkpoint-on-close")]
    no_checkpoint_on_close: bool,

    /// Restore database from backup snapshot and exit.
    /// Deletes corrupted volumes/ if needed, loads snapshot data, recreates indexes.
    /// Optionally specify a timestamp: --restore "20260315-100000.000"
    #[arg(long = "restore", value_name = "TIMESTAMP", num_args = 0..=1, default_missing_value = "")]
    restore: Option<String>,

    /// Create a backup snapshot and exit.
    #[arg(long = "snapshot")]
    snapshot: bool,

    /// Force remove volumes/ directory before opening (for recovery from corrupted volumes).
    /// Use with --restore to force snapshot-based recovery.
    #[arg(long = "reset-volumes")]
    reset_volumes: bool,

    /// Query timeout in milliseconds (0 for no timeout, default: 0)
    /// Long-running queries will be cancelled after this time.
    #[arg(short = 't', long = "timeout", value_name = "MS", default_value = "0")]
    timeout_ms: u64,
}

/// CLI state for interactive mode
struct Cli {
    db: Database,
    tx: Option<ApiTransaction>,
    in_transaction: bool,
    json_output: bool,
    limit: usize,
    #[allow(dead_code)]
    quiet: bool,
    timeout_ms: u64,
    editor: Editor<(), DefaultHistory>,
    current_query: String,
    in_multi_line: bool,
}

impl Cli {
    fn new(
        db: Database,
        json_output: bool,
        limit: usize,
        quiet: bool,
        timeout_ms: u64,
    ) -> io::Result<Self> {
        let config = Config::builder()
            .history_ignore_space(true)
            .edit_mode(EditMode::Emacs)
            .build();

        let mut editor =
            DefaultEditor::with_config(config).map_err(|e| io::Error::other(e.to_string()))?;

        // Load history from file
        if let Some(home) = dirs::home_dir() {
            let history_file = home.join(".stoolap_history");
            let _ = editor.load_history(&history_file);
        }

        Ok(Self {
            db,
            tx: None,
            in_transaction: false,
            json_output,
            limit,
            quiet,
            timeout_ms,
            editor,
            current_query: String::new(),
            in_multi_line: false,
        })
    }

    fn get_prompt(&self) -> &'static str {
        if self.in_multi_line {
            if self.in_transaction {
                "\x1b[1;33m[TXN]->\x1b[0m "
            } else {
                "\x1b[1;36m->\x1b[0m "
            }
        } else if self.in_transaction {
            "\x1b[1;33m[TXN]>\x1b[0m "
        } else {
            "\x1b[1;36m>\x1b[0m "
        }
    }

    fn run(&mut self) -> io::Result<()> {
        println!("Stoolap v{}.{}.{}", MAJOR, MINOR, PATCH);
        println!("Enter SQL commands, 'help' for assistance, or 'exit' to quit.");
        println!("Use Up/Down arrows for history, Ctrl+R to search history.");
        if self.json_output {
            println!("JSON output mode enabled.");
        }
        println!();

        loop {
            let prompt = self.get_prompt();
            match self.editor.readline(prompt) {
                Ok(line) => {
                    let line = line.trim();

                    // Handle empty line
                    if !self.in_multi_line && line.is_empty() {
                        continue;
                    }

                    // Handle special commands (only when not in multi-line mode)
                    if !self.in_multi_line {
                        match line.to_lowercase().as_str() {
                            "exit" | "quit" | "\\q" => {
                                if self.in_transaction {
                                    eprintln!("\x1b[1;33mWarning: Exiting with active transaction. Rolling back...\x1b[0m");
                                    let _ = self.rollback_transaction();
                                }
                                break;
                            }
                            "help" | "\\h" | "\\?" => {
                                self.print_help();
                                continue;
                            }
                            _ => {}
                        }
                    }

                    // Check for transaction control statements
                    let upper_line = line.to_uppercase();
                    if upper_line == "BEGIN"
                        || upper_line == "COMMIT"
                        || upper_line == "ROLLBACK"
                        || upper_line.starts_with("BEGIN ")
                    {
                        let _ = self.editor.add_history_entry(line);
                        let start = Instant::now();
                        if let Err(e) = self.execute_query(line) {
                            eprintln!("\x1b[1;31mError:\x1b[0m {}", e);
                        } else {
                            println!("\x1b[1;32mQuery executed in {:?}\x1b[0m", start.elapsed());
                        }
                        continue;
                    }

                    // Add line to current query
                    if !self.current_query.is_empty() {
                        self.current_query.push('\n');
                    }
                    self.current_query.push_str(line);

                    // Check if query ends with semicolon
                    let full_query = self.current_query.trim().to_string();
                    if full_query.ends_with(';') {
                        // Add to history
                        let history_entry = full_query.replace('\n', "\\n");
                        let _ = self.editor.add_history_entry(&history_entry);

                        self.in_multi_line = false;

                        // Split and execute statements
                        let statements = split_sql_statements(&full_query);
                        for stmt in statements {
                            let stmt = stmt.trim();
                            if stmt.is_empty() {
                                continue;
                            }

                            let start = Instant::now();
                            if let Err(e) = self.execute_query(stmt) {
                                eprintln!("\x1b[1;31mError:\x1b[0m {}", e);
                            } else {
                                println!(
                                    "\x1b[1;32mQuery executed in {:?}\x1b[0m",
                                    start.elapsed()
                                );
                            }
                        }

                        self.current_query.clear();
                    } else {
                        self.in_multi_line = true;
                    }
                }
                Err(ReadlineError::Interrupted) | Err(ReadlineError::Eof) => {
                    if self.in_transaction {
                        eprintln!("\n\x1b[1;33mWarning: Exiting with active transaction. Rolling back...\x1b[0m");
                        let _ = self.rollback_transaction();
                    }
                    break;
                }
                Err(e) => {
                    eprintln!("Error: {}", e);
                    break;
                }
            }
        }

        // Save history
        if let Some(home) = dirs::home_dir() {
            let history_file = home.join(".stoolap_history");
            let _ = self.editor.save_history(&history_file);
        }

        Ok(())
    }

    fn execute_query(&mut self, query: &str) -> Result<(), String> {
        let upper_query = query.to_uppercase();
        let upper_query = upper_query.trim();

        // Handle special commands
        match upper_query {
            "HELP" | "\\H" | "\\?" => {
                self.print_help();
                return Ok(());
            }
            _ => {}
        }

        // Handle transaction commands
        if upper_query.starts_with("BEGIN") {
            return self.begin_transaction();
        } else if upper_query == "COMMIT" {
            return self.commit_transaction();
        } else if upper_query == "ROLLBACK" {
            return self.rollback_transaction();
        }

        // Check if it's a query that returns rows
        // Include statements with RETURNING clause (INSERT/UPDATE/DELETE RETURNING)
        if upper_query.starts_with("SELECT")
            || upper_query.starts_with("WITH")
            || upper_query.starts_with("SHOW")
            || upper_query.starts_with("DESCRIBE")
            || upper_query.starts_with("DESC ")
            || upper_query.starts_with("EXPLAIN")
            || upper_query.starts_with("VACUUM")
            || (upper_query.starts_with("PRAGMA") && !upper_query.contains('='))
            || upper_query.contains(" RETURNING ")
            || upper_query.ends_with(" RETURNING")
        {
            self.execute_read_query(query)
        } else {
            self.execute_write_query(query)
        }
    }

    fn begin_transaction(&mut self) -> Result<(), String> {
        if self.in_transaction {
            return Err("already in a transaction".to_string());
        }

        let tx = self.db.begin().map_err(|e| e.to_string())?;
        self.tx = Some(tx);
        self.in_transaction = true;

        println!("\x1b[1;32mTransaction started\x1b[0m");
        Ok(())
    }

    fn commit_transaction(&mut self) -> Result<(), String> {
        if !self.in_transaction {
            return Err("not in a transaction".to_string());
        }

        if let Some(mut tx) = self.tx.take() {
            tx.commit().map_err(|e| e.to_string())?;
        }

        self.in_transaction = false;
        println!("\x1b[1;32mTransaction committed\x1b[0m");
        Ok(())
    }

    fn rollback_transaction(&mut self) -> Result<(), String> {
        if !self.in_transaction {
            return Err("not in a transaction".to_string());
        }

        if let Some(mut tx) = self.tx.take() {
            tx.rollback().map_err(|e| e.to_string())?;
        }

        self.in_transaction = false;
        println!("\x1b[1;33mTransaction rolled back\x1b[0m");
        Ok(())
    }

    fn execute_read_query(&mut self, query: &str) -> Result<(), String> {
        let rows_result = if self.in_transaction {
            if let Some(ref mut tx) = self.tx {
                tx.query(query, ()).map_err(|e| e.to_string())?
            } else {
                return Err("Transaction not available".to_string());
            }
        } else if self.timeout_ms > 0 {
            self.db
                .query_with_timeout(query, (), self.timeout_ms)
                .map_err(|e| e.to_string())?
        } else {
            self.db.query(query, ()).map_err(|e| e.to_string())?
        };

        let columns: Vec<String> = rows_result.columns().to_vec();

        // Collect all rows
        let mut all_rows: Vec<Vec<Value>> = Vec::new();
        for row_result in rows_result {
            let row = row_result.map_err(|e| e.to_string())?;
            let mut values = Vec::new();
            for i in 0..row.len() {
                values.push(row.get_value(i).cloned().unwrap_or(Value::null_unknown()));
            }
            all_rows.push(values);
        }

        let row_count = all_rows.len();

        if self.json_output {
            self.output_json(&columns, &all_rows, row_count)?;
        } else {
            self.output_table(&columns, &all_rows, row_count)?;
        }

        Ok(())
    }

    fn execute_write_query(&mut self, query: &str) -> Result<(), String> {
        let rows_affected = if self.in_transaction {
            if let Some(ref mut tx) = self.tx {
                tx.execute(query, ()).map_err(|e| e.to_string())?
            } else {
                return Err("Transaction not available".to_string());
            }
        } else if self.timeout_ms > 0 {
            self.db
                .execute_with_timeout(query, (), self.timeout_ms)
                .map_err(|e| e.to_string())?
        } else {
            self.db.execute(query, ()).map_err(|e| e.to_string())?
        };

        if self.json_output {
            println!(r#"{{"rows_affected":{}}}"#, rows_affected);
        } else {
            let row_text = if rows_affected == 1 { "row" } else { "rows" };
            println!("\x1b[1;32m{} {} affected\x1b[0m", rows_affected, row_text);
        }

        Ok(())
    }

    fn output_json(
        &self,
        columns: &[String],
        rows: &[Vec<Value>],
        row_count: usize,
    ) -> Result<(), String> {
        let json_rows: Vec<Vec<serde_json::Value>> = rows
            .iter()
            .map(|row| row.iter().map(value_to_json).collect())
            .collect();

        let result = serde_json::json!({
            "columns": columns,
            "rows": json_rows,
            "count": row_count
        });

        println!(
            "{}",
            serde_json::to_string(&result).map_err(|e| e.to_string())?
        );
        Ok(())
    }

    fn output_explain_plan(&self, rows: &[Vec<Value>], row_count: usize) -> Result<(), String> {
        println!();
        for row in rows {
            if let Some(Value::Text(line)) = row.first() {
                println!("{}", colorize_plan_line(line));
            }
        }
        println!();
        let line_text = if row_count == 1 { "line" } else { "lines" };
        println!("\x1b[2m({} {} in plan)\x1b[0m", row_count, line_text);
        Ok(())
    }

    fn output_table(
        &self,
        columns: &[String],
        rows: &[Vec<Value>],
        row_count: usize,
    ) -> Result<(), String> {
        // Render EXPLAIN output without table borders, with color highlighting.
        // Verify first row starts with a known statement keyword to avoid false positives
        // from queries that alias a column as "plan".
        if columns.len() == 1 && columns[0] == "plan" && is_explain_output(rows) {
            return self.output_explain_plan(rows, row_count);
        }

        let mut table = Table::new();
        table
            .load_preset(UTF8_FULL_CONDENSED)
            .set_content_arrangement(ContentArrangement::Dynamic);

        // Add header
        table.set_header(columns.iter().map(Cell::new));

        // Smart truncation with limit
        if self.limit > 0 && row_count > self.limit {
            let top_rows = self.limit / 2;
            let bottom_rows = self.limit - top_rows;

            // Add top rows
            for row in rows.iter().take(top_rows) {
                table.add_row(row.iter().map(|v| Cell::new(format_value(v))));
            }

            // Add truncation indicator
            let hidden_rows = row_count - self.limit;
            let mut truncation_row: Vec<Cell> = Vec::new();
            let message = format!("... ({} more rows) ...", hidden_rows);
            for (i, _) in columns.iter().enumerate() {
                if i == columns.len() / 2 {
                    truncation_row.push(Cell::new(&message));
                } else {
                    truncation_row.push(Cell::new(""));
                }
            }
            table.add_row(truncation_row);

            // Add bottom rows
            let start_idx = row_count.saturating_sub(bottom_rows).max(top_rows);
            for row in rows.iter().skip(start_idx) {
                table.add_row(row.iter().map(|v| Cell::new(format_value(v))));
            }
        } else {
            // Add all rows
            for row in rows {
                table.add_row(row.iter().map(|v| Cell::new(format_value(v))));
            }
        }

        println!("{table}");

        // Print summary
        let row_text = if row_count == 1 { "row" } else { "rows" };
        if self.limit > 0 && row_count > self.limit {
            println!(
                "\x1b[1;32m{} {} in set (showing {})\x1b[0m",
                row_count, row_text, self.limit
            );
        } else {
            println!("\x1b[1;32m{} {} in set\x1b[0m", row_count, row_text);
        }

        Ok(())
    }

    fn print_help(&self) {
        println!("\x1b[1mStoolap SQL CLI Commands:\x1b[0m");
        println!();
        println!("  \x1b[1;33mSQL Commands:\x1b[0m");
        println!("    SELECT ...             Execute a SELECT query");
        println!("    INSERT ...             Insert data into a table");
        println!("    UPDATE ...             Update data in a table");
        println!("    DELETE ...             Delete data from a table");
        println!("    CREATE TABLE ...       Create a new table");
        println!("    CREATE INDEX ...       Create an index on a column");
        println!("    ALTER TABLE ...        Modify table schema");
        println!("    SHOW TABLES            List all tables");
        println!("    SHOW CREATE TABLE ...  Show CREATE TABLE statement for a table");
        println!("    SHOW INDEXES FROM ...  Show indexes for a table");
        println!();
        println!("  \x1b[1;33mTransaction Commands:\x1b[0m");
        println!("    BEGIN                  Start a new transaction");
        println!("    COMMIT                 Commit the current transaction");
        println!("    ROLLBACK               Rollback the current transaction");
        println!();
        println!("  \x1b[1;33mPRAGMA Commands:\x1b[0m");
        println!("    PRAGMA CHECKPOINT      Seal hot data to volumes, compact, truncate WAL");
        println!("    PRAGMA SNAPSHOT        Create backup snapshot (.bin files)");
        println!("    PRAGMA RESTORE         Restore from latest backup snapshot");
        println!(
            "    PRAGMA key=value       Set configuration (sync_mode, checkpoint_interval, ...)"
        );
        println!("    PRAGMA key             Show current configuration value");
        println!();
        println!("  \x1b[1;33mSpecial Commands:\x1b[0m");
        println!("    exit, quit, \\q         Exit the CLI");
        println!("    help, \\h, \\?          Show this help message");
        println!();
        println!("  \x1b[1;33mKeyboard Shortcuts:\x1b[0m");
        println!("    Up/Down arrow keys     Navigate command history");
        println!("    Ctrl+R                 Search command history");
        println!("    Ctrl+A                 Move cursor to beginning of line");
        println!("    Ctrl+E                 Move cursor to end of line");
        println!("    Ctrl+W                 Delete word before cursor");
        println!("    Ctrl+U                 Delete from cursor to beginning of line");
        println!("    Ctrl+K                 Delete from cursor to end of line");
        println!("    Ctrl+L                 Clear screen");
        println!();
    }
}

/// Build DSN with query parameters from CLI args
fn build_dsn(args: &Args) -> String {
    let mut dsn = args.db_path.clone();

    // Only add params for file:// databases
    if !dsn.starts_with("file://") {
        return dsn;
    }

    let mut params = Vec::new();

    // Handle persistence profile presets
    if let Some(ref profile) = args.persistence_profile {
        match profile.to_lowercase().as_str() {
            "fast" => {
                params.push("sync=none".to_string());
            }
            "durable" => {
                params.push("sync=full".to_string());
            }
            "normal" => {
                // Default, no params needed
            }
            _ => {
                eprintln!(
                    "Warning: Unknown profile '{}', using 'normal'. Valid: fast, normal, durable",
                    profile
                );
            }
        }
    }

    // Individual sync mode overrides profile
    if let Some(ref sync) = args.sync_mode {
        // Remove any existing sync param from profile
        params.retain(|p| !p.starts_with("sync="));
        match sync.to_lowercase().as_str() {
            "none" | "off" => params.push("sync=none".to_string()),
            "normal" => params.push("sync=normal".to_string()),
            "full" => params.push("sync=full".to_string()),
            _ => {
                eprintln!(
                    "Warning: Unknown sync mode '{}', using 'normal'. Valid: none, normal, full",
                    sync
                );
            }
        }
    }

    // Checkpoint interval
    if let Some(interval) = args.checkpoint_interval {
        params.push(format!("checkpoint_interval={}", interval));
    }

    // Compact threshold
    if let Some(count) = args.compact_threshold {
        params.push(format!("compact_threshold={}", count));
    }

    // WAL max size (convert MB to bytes)
    if let Some(mb) = args.wal_max_size {
        params.push(format!("wal_max_size={}", mb as u64 * 1024 * 1024));
    }

    // Compression
    if let Some(ref comp) = args.compression {
        match comp.to_lowercase().as_str() {
            "on" | "true" | "1" | "yes" => params.push("compression=on".to_string()),
            "off" | "false" | "0" | "no" => params.push("compression=off".to_string()),
            _ => {
                eprintln!(
                    "Warning: Unknown compression value '{}', using 'on'. Valid: on, off",
                    comp
                );
            }
        }
    }

    // Keep snapshots
    if let Some(count) = args.keep_snapshots {
        params.push(format!("keep_snapshots={}", count));
    }

    // Checkpoint on close
    if args.no_checkpoint_on_close {
        params.push("checkpoint_on_close=off".to_string());
    }

    // Append params to DSN
    if !params.is_empty() {
        let separator = if dsn.contains('?') { "&" } else { "?" };
        dsn.push_str(separator);
        dsn.push_str(&params.join("&"));
    }

    dsn
}

/// Print persistence configuration info
fn print_persistence_info(args: &Args) {
    let sync_mode = if let Some(ref sync) = args.sync_mode {
        sync.to_lowercase()
    } else if let Some(ref profile) = args.persistence_profile {
        match profile.to_lowercase().as_str() {
            "fast" => "none".to_string(),
            "durable" => "full".to_string(),
            _ => "normal".to_string(),
        }
    } else {
        "normal".to_string()
    };

    let sync_desc = match sync_mode.as_str() {
        "none" | "off" => "none (fastest, less durable)",
        "full" => "full (slowest, most durable)",
        _ => "normal (balanced)",
    };

    println!("Persistence: WAL sync mode = {}", sync_desc);

    if let Some(interval) = args.checkpoint_interval {
        println!("Persistence: Checkpoint interval = {}s", interval);
    }

    if let Some(count) = args.compact_threshold {
        println!("Persistence: Compact threshold = {}", count);
    }

    if let Some(count) = args.keep_snapshots {
        println!("Persistence: Keep snapshots = {}", count);
    }

    if let Some(mb) = args.wal_max_size {
        println!("Persistence: WAL max size = {}MB", mb);
    }

    if let Some(ref comp) = args.compression {
        println!("Persistence: Compression = {}", comp);
    }

    if args.no_checkpoint_on_close {
        println!("Persistence: Checkpoint on close = off");
    }
}

fn main() {
    let args = Args::parse();

    // Build the DSN with optional query parameters
    let db_path = build_dsn(&args);

    // ── Filesystem-level operations (before Database::open) ────────
    // These work directly on disk so they can recover from broken state.

    let db_dir = if let Some(path) = db_path.strip_prefix("file://") {
        Some(std::path::PathBuf::from(
            path.split('?').next().unwrap_or(path),
        ))
    } else {
        None
    };

    // Handle --reset-volumes BEFORE --restore: delete volumes/ (and wal/)
    // so that Database::open() in the restore path doesn't choke on
    // corrupted on-disk state. This makes `--reset-volumes --restore`
    // work as documented.
    if args.reset_volumes {
        if let Some(ref dir) = db_dir {
            let vol_dir = dir.join("volumes");
            if vol_dir.exists() {
                match std::fs::remove_dir_all(&vol_dir) {
                    Ok(()) => {
                        if !args.quiet {
                            eprintln!("Removed volumes/ directory for recovery");
                        }
                    }
                    Err(e) => {
                        eprintln!("Error removing volumes/ directory: {}", e);
                        std::process::exit(1);
                    }
                }
            }
            let wal_dir = dir.join("wal");
            if wal_dir.exists() {
                match std::fs::remove_dir_all(&wal_dir) {
                    Ok(()) => {
                        if !args.quiet {
                            eprintln!("Removed wal/ directory for recovery");
                        }
                    }
                    Err(e) => {
                        eprintln!("Error removing wal/ directory: {}", e);
                        std::process::exit(1);
                    }
                }
            }
            // Also remove the stale lock file so open() doesn't fail
            let lock_file = dir.join("db.lock");
            if lock_file.exists() {
                let _ = std::fs::remove_file(&lock_file);
            }
        } else if !args.quiet {
            eprintln!("Warning: --reset-volumes only applies to file:// databases");
        }
    }

    // Handle --restore: filesystem-level recovery from backup snapshots.
    // If --reset-volumes was used above, volumes/ and wal/ are already gone.
    if args.restore.is_some() {
        let dir = match &db_dir {
            Some(d) => d,
            None => {
                eprintln!("Error: --restore requires a file:// database");
                std::process::exit(1);
            }
        };

        let snapshot_dir = dir.join("snapshots");
        if !snapshot_dir.exists() {
            eprintln!("Error: No snapshots/ directory found in {:?}", dir);
            eprintln!(
                "Create a backup first with: stoolap -d {} --snapshot",
                db_path
            );
            std::process::exit(1);
        }

        // Find snapshot files to validate
        let mut table_count = 0u32;
        if let Ok(entries) = std::fs::read_dir(&snapshot_dir) {
            for entry in entries.flatten() {
                if entry.file_type().map(|ft| ft.is_dir()).unwrap_or(false) {
                    let table_name = entry.file_name().to_string_lossy().to_string();
                    let has_bin = std::fs::read_dir(entry.path())
                        .ok()
                        .map(|e| {
                            e.flatten().any(|f| {
                                f.path()
                                    .extension()
                                    .map(|ext| ext == "bin")
                                    .unwrap_or(false)
                            })
                        })
                        .unwrap_or(false);
                    if has_bin {
                        table_count += 1;
                        if !args.quiet {
                            eprintln!("[restore] Found snapshot for table '{}'", table_name);
                        }
                    }
                }
            }
        }

        if table_count == 0 {
            eprintln!("Error: No snapshot .bin files found in {:?}", snapshot_dir);
            std::process::exit(1);
        }

        // Timestamp restore: open the database normally and use the engine's
        // restore_snapshot() API, which loads the specific timestamp without
        // deleting other backup files. Other snapshots are preserved.
        if let Some(ref ts) = args.restore {
            if !ts.is_empty() {
                if !args.quiet {
                    eprintln!("[restore] Restoring to timestamp: {}", ts);
                }

                let db = match Database::open(&db_path) {
                    Ok(db) => db,
                    Err(e) => {
                        eprintln!("Error opening database: {}", e);
                        std::process::exit(1);
                    }
                };

                match db.restore_snapshot(Some(ts)) {
                    Ok(msg) => {
                        println!("{}", msg);
                    }
                    Err(e) => {
                        eprintln!("Error restoring snapshot: {}", e);
                        db.close().unwrap_or(());
                        std::process::exit(1);
                    }
                }

                db.close().unwrap_or(());
                return;
            }
        }

        // Crash recovery (no timestamp): find the newest complete manifest
        // and use the timestamped restore path for a coherent snapshot set.
        let snap_dir = dir.join("snapshots");

        // Scan manifests newest-first, pick the latest whose table snapshots all exist.
        let restore_ts = {
            let mut manifests: Vec<String> = std::fs::read_dir(&snap_dir)
                .ok()
                .map(|entries| {
                    entries
                        .filter_map(|e| e.ok())
                        .filter_map(|e| {
                            let name = e.file_name().to_string_lossy().to_string();
                            name.strip_prefix("manifest-")
                                .and_then(|s| s.strip_suffix(".json"))
                                .map(|ts| ts.to_string())
                        })
                        .collect()
                })
                .unwrap_or_default();
            manifests.sort();
            manifests.reverse(); // newest first

            let mut chosen: Option<String> = None;
            for ts in &manifests {
                // Read manifest to get table list
                let manifest_path = snap_dir.join(format!("manifest-{}.json", ts));
                let Ok(content) = std::fs::read_to_string(&manifest_path) else {
                    continue;
                };
                let table_names: Vec<String> = if let Some(start) = content.find("\"tables\":[") {
                    let after = &content[start + 10..];
                    // Parse JSON string values properly (handle escaped quotes/backslashes)
                    let mut names = Vec::new();
                    let mut i = 0;
                    let bytes = after.as_bytes();
                    while i < bytes.len() {
                        if bytes[i] == b']' {
                            break;
                        }
                        if bytes[i] == b'"' {
                            i += 1;
                            let mut name = String::new();
                            while i < bytes.len() && bytes[i] != b'"' {
                                if bytes[i] == b'\\' && i + 1 < bytes.len() {
                                    i += 1; // skip backslash, take next char
                                }
                                name.push(bytes[i] as char);
                                i += 1;
                            }
                            if !name.is_empty() {
                                names.push(name);
                            }
                            i += 1; // skip closing quote
                        } else {
                            i += 1;
                        }
                    }
                    names
                } else {
                    Vec::new()
                };

                if table_names.is_empty() {
                    continue;
                }

                // Verify all table snapshots AND DDL blob exist for this timestamp
                let ddl_file = snap_dir.join(format!("ddl-{}.bin", ts));
                let all_exist = ddl_file.exists()
                    && table_names.iter().all(|table| {
                        let snapshot_file =
                            snap_dir.join(table).join(format!("snapshot-{}.bin", ts));
                        snapshot_file.exists()
                    });

                if all_exist {
                    if !args.quiet {
                        eprintln!(
                            "[restore] Found complete manifest for timestamp {} ({} tables)",
                            ts,
                            table_names.len()
                        );
                    }
                    chosen = Some(ts.clone());
                    break;
                } else if !args.quiet {
                    eprintln!(
                        "[restore] Manifest {} has missing table snapshots, trying older...",
                        ts
                    );
                }
            }
            chosen
        };

        if let Some(ts) = restore_ts {
            // Use the timestamped restore path for a coherent snapshot set.
            // This reuses the engine's restore_from_snapshot which handles
            // data + DDL matching correctly.
            if !args.quiet {
                eprintln!("[restore] Restoring from timestamp {}...", ts);
            }

            let db = match Database::open(&db_path) {
                Ok(db) => db,
                Err(e) => {
                    eprintln!("Error opening database: {}", e);
                    std::process::exit(1);
                }
            };

            match db.restore_snapshot(Some(&ts)) {
                Ok(msg) => println!("{}", msg),
                Err(e) => {
                    eprintln!("Error restoring snapshot: {}", e);
                    db.close().unwrap_or(());
                    std::process::exit(1);
                }
            }
            db.close().unwrap_or(());
        } else {
            // No manifest found — legacy fallback: open DB and let load_snapshots
            // pick the newest per table. Warn about potential DDL mismatch.
            if !args.quiet {
                eprintln!("[restore] No snapshot manifest found, using legacy restore path.");
                eprintln!("[restore] Warning: DDL (indexes, views) may not match restored data.");
            }

            // Remove volumes/ and wal/ for clean startup
            let vol_dir = dir.join("volumes");
            if vol_dir.exists() {
                if let Err(e) = std::fs::remove_dir_all(&vol_dir) {
                    eprintln!("Error removing volumes/: {}", e);
                    std::process::exit(1);
                }
            }
            let wal_dir = dir.join("wal");
            if wal_dir.exists() {
                if let Err(e) = std::fs::remove_dir_all(&wal_dir) {
                    eprintln!("Error removing wal/: {}", e);
                    std::process::exit(1);
                }
            }

            let db = match Database::open(&db_path) {
                Ok(db) => db,
                Err(e) => {
                    eprintln!("Error opening database after restore: {}", e);
                    std::process::exit(1);
                }
            };

            // Verify data loaded
            match db.query("SHOW TABLES", ()) {
                Ok(rows) => {
                    let tables: Vec<String> = rows
                        .filter_map(|r| r.ok())
                        .filter_map(|r| r.get::<String>(0).ok())
                        .collect();
                    println!(
                        "Restore complete (legacy). {} tables recovered: {}",
                        tables.len(),
                        tables.join(", ")
                    );
                }
                Err(_) => {
                    println!(
                        "Restore complete (legacy). {} tables recovered.",
                        table_count
                    );
                }
            }
            db.close().unwrap_or(());
        }
        return;
    }

    // --reset-volumes already handled above (before --restore).

    // Open the database
    let db = match Database::open(&db_path) {
        Ok(db) => db,
        Err(e) => {
            eprintln!("Error opening database: {}", e);
            std::process::exit(1);
        }
    };

    if !args.quiet {
        println!("Connected to database: {}", db_path);
        // Show persistence info for file databases
        if db_path.starts_with("file://") {
            print_persistence_info(&args);
        }
    }

    // Handle --snapshot: open DB, create backup, close, exit
    if args.snapshot {
        match db.query("PRAGMA SNAPSHOT", ()) {
            Ok(mut rows) => {
                if let Some(Ok(row)) = rows.next() {
                    if let Ok(msg) = row.get::<String>(0) {
                        println!("{}", msg);
                    }
                }
            }
            Err(e) => {
                eprintln!("Error creating snapshot: {}", e);
                std::process::exit(1);
            }
        }
        db.close().unwrap_or(());
        return;
    }

    // Handle execute flag - run single query and exit
    if let Some(ref sql) = args.execute {
        if let Err(e) = execute_query_with_options(
            &db,
            sql,
            args.json_output,
            args.quiet,
            args.limit,
            args.timeout_ms,
        ) {
            eprintln!("Error: {}", e);
            std::process::exit(1);
        }
        return;
    }

    // Handle file flag - execute SQL from file
    if let Some(ref filename) = args.file {
        let result = execute_from_file(
            &db,
            filename,
            args.json_output,
            args.quiet,
            args.limit,
            args.timeout_ms,
        );
        drop(db); // ensure close_engine() runs before exit
        if let Err(e) = result {
            eprintln!("Error: {}", e);
            std::process::exit(1);
        }
        return;
    }

    // Check if we're getting input from a pipe
    let is_pipe = !std::io::stdin().is_terminal();

    if is_pipe {
        let result = execute_piped_input(
            &db,
            args.json_output,
            args.quiet,
            args.limit,
            args.timeout_ms,
        );
        drop(db); // ensure close_engine() runs before exit
        if let Err(e) = result {
            eprintln!("Error: {}", e);
            std::process::exit(1);
        }
        return;
    }

    // Interactive mode
    let mut cli = match Cli::new(
        db,
        args.json_output,
        args.limit,
        args.quiet,
        args.timeout_ms,
    ) {
        Ok(cli) => cli,
        Err(e) => {
            eprintln!("Error initializing CLI: {}", e);
            std::process::exit(1);
        }
    };

    if let Err(e) = cli.run() {
        // Drop cli (and db) before exiting so close_engine() runs
        drop(cli);
        eprintln!("Error: {}", e);
        std::process::exit(1);
    }
}

fn execute_from_file(
    db: &Database,
    filename: &str,
    json_output: bool,
    quiet: bool,
    row_limit: usize,
    timeout_ms: u64,
) -> Result<(), String> {
    let file =
        File::open(filename).map_err(|e| format!("Error opening file {}: {}", filename, e))?;
    let reader = BufReader::new(file);
    let mut current_statement = String::new();

    for line_result in reader.lines() {
        let line = line_result.map_err(|e| format!("Error reading file: {}", e))?;

        // Skip shell-style comment lines
        let trimmed = line.trim();
        if trimmed.starts_with('#') {
            continue;
        }

        // Skip SQL-style comment lines
        if trimmed.starts_with("--") || (trimmed.starts_with("/*") && trimmed.ends_with("*/")) {
            continue;
        }

        // If blank line and we have a statement, execute it
        if trimmed.is_empty() && !current_statement.is_empty() {
            let q = current_statement.trim().to_string();
            current_statement.clear();

            if !q.is_empty() {
                let statements = split_sql_statements(&q);
                for stmt in statements {
                    let stmt = stmt.trim();
                    if stmt.is_empty() {
                        continue;
                    }

                    if let Err(e) = execute_query_with_options(
                        db,
                        stmt,
                        json_output,
                        quiet,
                        row_limit,
                        timeout_ms,
                    ) {
                        eprintln!("Error: {}", e);
                    }
                }
            }
        } else {
            current_statement.push_str(&line);
            current_statement.push('\n');
        }
    }

    // Execute any remaining statement
    if !current_statement.is_empty() {
        let q = current_statement.trim().to_string();
        if !q.is_empty() {
            let statements = split_sql_statements(&q);
            for stmt in statements {
                let stmt = stmt.trim();
                if stmt.is_empty() {
                    continue;
                }

                if let Err(e) =
                    execute_query_with_options(db, stmt, json_output, quiet, row_limit, timeout_ms)
                {
                    eprintln!("Error: {}", e);
                }
            }
        }
    }

    Ok(())
}

fn execute_piped_input(
    db: &Database,
    json_output: bool,
    quiet: bool,
    row_limit: usize,
    timeout_ms: u64,
) -> Result<(), String> {
    let stdin = io::stdin();
    let reader = stdin.lock();
    let mut current_statement = String::new();

    for line_result in reader.lines() {
        let line = line_result.map_err(|e| format!("Error reading input: {}", e))?;

        // Skip shell-style comment lines
        let trimmed = line.trim();
        if trimmed.starts_with('#') {
            continue;
        }

        // Skip SQL-style comment lines
        if trimmed.starts_with("--") || (trimmed.starts_with("/*") && trimmed.ends_with("*/")) {
            continue;
        }

        // If blank line and we have a statement, execute it
        if trimmed.is_empty() && !current_statement.is_empty() {
            let q = current_statement.trim().to_string();
            current_statement.clear();

            if !q.is_empty() {
                let statements = split_sql_statements(&q);
                for stmt in statements {
                    let stmt = stmt.trim();
                    if stmt.is_empty() {
                        continue;
                    }

                    let start = Instant::now();
                    if let Err(e) = execute_query_with_options(
                        db,
                        stmt,
                        json_output,
                        quiet,
                        row_limit,
                        timeout_ms,
                    ) {
                        eprintln!("Error: {}", e);
                    } else if !json_output && !quiet {
                        println!("Query executed in {:?}", start.elapsed());
                    }
                }
            }
        } else {
            current_statement.push_str(&line);
            current_statement.push('\n');
        }
    }

    // Execute any remaining statement
    if !current_statement.is_empty() {
        let q = current_statement.trim().to_string();
        if !q.is_empty() {
            let statements = split_sql_statements(&q);
            for stmt in statements {
                let stmt = stmt.trim();
                if stmt.is_empty() {
                    continue;
                }

                let start = Instant::now();
                if let Err(e) =
                    execute_query_with_options(db, stmt, json_output, quiet, row_limit, timeout_ms)
                {
                    eprintln!("Error: {}", e);
                } else if !json_output && !quiet {
                    println!("Query executed in {:?}", start.elapsed());
                }
            }
        }
    }

    Ok(())
}

fn execute_query_with_options(
    db: &Database,
    query: &str,
    json_output: bool,
    quiet: bool,
    row_limit: usize,
    timeout_ms: u64,
) -> Result<(), String> {
    let upper_query = query.to_uppercase();
    let upper_query = upper_query.trim();

    // Handle special commands
    match upper_query {
        "HELP" | "\\H" | "\\?" => {
            print_help_main();
            return Ok(());
        }
        "EXIT" | "QUIT" | "\\Q" => {
            return Err("exit requested".to_string());
        }
        _ => {}
    }

    // Handle parameter syntax: "SELECT ... -- PARAMS: val1, val2"
    let (sql, params) = parse_params(query);

    // Check if it's a query that returns rows
    // Include statements with RETURNING clause (INSERT/UPDATE/DELETE RETURNING)
    if upper_query.starts_with("SELECT")
        || upper_query.starts_with("WITH")
        || upper_query.starts_with("SHOW")
        || upper_query.starts_with("DESCRIBE")
        || upper_query.starts_with("DESC ")
        || upper_query.starts_with("EXPLAIN")
        || upper_query.starts_with("VACUUM")
        || (upper_query.starts_with("PRAGMA") && !upper_query.contains('='))
        || upper_query.contains(" RETURNING ")
        || upper_query.ends_with(" RETURNING")
    {
        let rows_result = if timeout_ms > 0 {
            db.query_with_timeout(&sql, params, timeout_ms)
                .map_err(|e| e.to_string())?
        } else {
            db.query(&sql, params).map_err(|e| e.to_string())?
        };

        let columns: Vec<String> = rows_result.columns().to_vec();

        // Collect all rows
        let mut all_rows: Vec<Vec<Value>> = Vec::new();
        for row_result in rows_result {
            let row = row_result.map_err(|e| e.to_string())?;
            let mut values = Vec::new();
            for i in 0..row.len() {
                values.push(row.get_value(i).cloned().unwrap_or(Value::null_unknown()));
            }
            all_rows.push(values);
        }

        let row_count = all_rows.len();

        if json_output {
            output_json(&columns, &all_rows, row_count)?;
        } else {
            output_table(&columns, &all_rows, row_count, row_limit, quiet)?;
        }
    } else {
        // Execute a non-query statement
        let rows_affected = if timeout_ms > 0 {
            db.execute_with_timeout(&sql, params, timeout_ms)
                .map_err(|e| e.to_string())?
        } else {
            db.execute(&sql, params).map_err(|e| e.to_string())?
        };

        if json_output {
            println!(r#"{{"rows_affected":{}}}"#, rows_affected);
        } else if !quiet {
            println!("{} rows affected", rows_affected);
        }
    }

    Ok(())
}

fn parse_params(query: &str) -> (String, Vec<Value>) {
    let parts: Vec<&str> = query.split(" -- PARAMS: ").collect();

    if parts.len() <= 1 {
        // Remove trailing semicolon
        let sql = parts[0].trim().trim_end_matches(';').to_string();
        return (sql, Vec::new());
    }

    let sql = parts[0].trim().trim_end_matches(';').to_string();
    let param_string = parts[1].trim();

    let mut params = Vec::new();
    for val in param_string.split(',') {
        params.push(convert_param_value(val.trim()));
    }

    (sql, params)
}

fn convert_param_value(value: &str) -> Value {
    // Try to convert to integer
    if let Ok(i) = value.parse::<i64>() {
        return Value::Integer(i);
    }

    // Try to convert to float (only if contains '.')
    if value.contains('.') {
        if let Ok(f) = value.parse::<f64>() {
            return Value::Float(f);
        }
    }

    // Try to convert to boolean
    if value == "true" {
        return Value::Boolean(true);
    }
    if value == "false" {
        return Value::Boolean(false);
    }

    // Check for null
    if value.eq_ignore_ascii_case("null") {
        return Value::null_unknown();
    }

    // Default to string
    Value::text(value)
}

fn output_json(columns: &[String], rows: &[Vec<Value>], row_count: usize) -> Result<(), String> {
    let json_rows: Vec<Vec<serde_json::Value>> = rows
        .iter()
        .map(|row| row.iter().map(value_to_json).collect())
        .collect();

    let result = serde_json::json!({
        "columns": columns,
        "rows": json_rows,
        "count": row_count
    });

    println!(
        "{}",
        serde_json::to_string(&result).map_err(|e| e.to_string())?
    );
    Ok(())
}

fn output_table(
    columns: &[String],
    rows: &[Vec<Value>],
    row_count: usize,
    row_limit: usize,
    quiet: bool,
) -> Result<(), String> {
    // Print the column names
    for (i, column) in columns.iter().enumerate() {
        if i > 0 {
            print!(" | ");
        }
        print!("{}", column);
    }
    println!();

    // Print a separator
    for (i, _) in columns.iter().enumerate() {
        if i > 0 {
            print!("-+-");
        }
        print!("----");
    }
    println!();

    // Display rows with smart truncation
    if row_limit == 0 || row_count <= row_limit {
        for row in rows {
            for (i, value) in row.iter().enumerate() {
                if i > 0 {
                    print!(" | ");
                }
                print!("{}", format_value(value));
            }
            println!();
        }

        if !quiet {
            println!("{} rows in set", row_count);
        }
    } else {
        // Smart truncation
        let top_rows = row_limit / 2;
        let bottom_rows = row_limit - top_rows;

        // Show top rows
        for row in rows.iter().take(top_rows) {
            for (i, value) in row.iter().enumerate() {
                if i > 0 {
                    print!(" | ");
                }
                print!("{}", format_value(value));
            }
            println!();
        }

        // Show truncation indicator
        let hidden_rows = row_count - row_limit;
        println!();
        println!("    \x1b[2m... ({} more rows) ...\x1b[0m", hidden_rows);
        println!();

        // Show bottom rows
        let start_idx = row_count.saturating_sub(bottom_rows).max(top_rows);
        for row in rows.iter().skip(start_idx) {
            for (i, value) in row.iter().enumerate() {
                if i > 0 {
                    print!(" | ");
                }
                print!("{}", format_value(value));
            }
            println!();
        }

        if !quiet {
            println!("{} rows in set (showing {})", row_count, row_limit);
        }
    }

    Ok(())
}

// ============================================================================
// EXPLAIN Plan Colorization
// ============================================================================

// ANSI color codes for EXPLAIN plan output (256-color, high contrast)
const C_HEADER: &str = "\x1b[1;4m"; // Bold + underline for headers
const C_DIM: &str = "\x1b[38;5;245m"; // Medium gray for arrows
const C_RESET: &str = "\x1b[0m";
const C_INDEX: &str = "\x1b[1;38;5;46m"; // Bright green for index access
const C_SEQ: &str = "\x1b[1;38;5;208m"; // Orange for seq scan warning
const C_JOIN: &str = "\x1b[1;38;5;75m"; // Light blue for join operators
const C_STATS: &str = "\x1b[1;38;5;196m"; // Bright pure red for timing/row stats
const C_LABEL: &str = "\x1b[38;5;117m"; // Sky blue for labels
const C_SUBIDX: &str = "\x1b[38;5;114m"; // Medium green for sub-index details
const C_CTE: &str = "\x1b[38;5;183m"; // Lavender for CTE/subquery scans

/// Check if rows look like EXPLAIN plan output (not a regular query with a "plan" column)
fn is_explain_output(rows: &[Vec<Value>]) -> bool {
    if let Some(first_row) = rows.first() {
        if let Some(Value::Text(line)) = first_row.first() {
            let trimmed = line.trim();
            return trimmed.starts_with("SELECT")
                || trimmed.starts_with("INSERT")
                || trimmed.starts_with("UPDATE")
                || trimmed.starts_with("DELETE")
                || trimmed.starts_with("WITH ");
        }
    }
    false
}

/// Colorize a single EXPLAIN plan line with ANSI escape codes
fn colorize_plan_line(line: &str) -> String {
    let trimmed = line.trim_start();
    let indent = &line[..line.len() - trimmed.len()];

    let mut result = String::with_capacity(line.len() + 80);
    result.push_str(indent);

    if trimmed.starts_with("-> ") {
        colorize_node(trimmed, &mut result);
    } else if trimmed.starts_with("SELECT")
        || trimmed.starts_with("INSERT")
        || trimmed.starts_with("UPDATE")
        || trimmed.starts_with("DELETE")
    {
        colorize_header(trimmed, &mut result);
    } else if trimmed.starts_with("WITH ") || trimmed.starts_with("Statement:") {
        result.push_str(C_HEADER);
        result.push_str(trimmed);
        result.push_str(C_RESET);
    } else {
        colorize_detail(trimmed, &mut result);
    }

    result
}

/// Split a line into the main content and trailing stats parenthetical.
/// Stats always start with "(actual " or "(cost=".
fn split_stats(line: &str) -> (&str, Option<&str>) {
    for marker in &["(actual ", "(cost="] {
        if let Some(pos) = line.rfind(marker) {
            return (line[..pos].trim_end(), Some(&line[pos..]));
        }
    }
    (line, None)
}

/// Colorize a node line starting with "-> "
fn colorize_node(line: &str, result: &mut String) {
    let after_arrow = &line[3..]; // skip "-> "

    // Dim arrow
    result.push_str(C_DIM);
    result.push_str("->");
    result.push_str(C_RESET);
    result.push(' ');

    // Indexed access patterns (green)
    let indexed = [
        "Index Scan",
        "PK Lookup",
        "Multi-Index Scan",
        "Composite Index Scan",
        "Index Nested Loop",
    ];
    // Sequential scan patterns (orange warning)
    let sequential = ["Seq Scan", "Parallel Seq Scan"];
    // Join operator patterns (blue)
    let joins = ["Hash Join", "Merge Join", "Nested Loop ("];

    let color = if indexed.iter().any(|p| after_arrow.starts_with(p)) {
        C_INDEX
    } else if sequential.iter().any(|p| after_arrow.starts_with(p)) {
        C_SEQ
    } else if joins.iter().any(|p| after_arrow.starts_with(p)) {
        C_JOIN
    } else if after_arrow.starts_with("Subquery Scan") || after_arrow.starts_with("CTE Scan") {
        C_CTE
    } else {
        // Sub-index detail lines like "idx_cat on category: = Electronics"
        result.push_str(C_SUBIDX);
        result.push_str(after_arrow);
        result.push_str(C_RESET);
        return;
    };

    let (scan_part, stats_part) = split_stats(after_arrow);
    result.push_str(color);
    result.push_str(scan_part);
    result.push_str(C_RESET);
    if let Some(stats) = stats_part {
        result.push(' ');
        result.push_str(C_STATS);
        result.push_str(stats);
        result.push_str(C_RESET);
    }
}

/// Colorize a statement header line (SELECT, INSERT, UPDATE, DELETE)
fn colorize_header(line: &str, result: &mut String) {
    let (scan_part, stats_part) = split_stats(line);
    result.push_str(C_HEADER);
    result.push_str(scan_part);
    result.push_str(C_RESET);
    if let Some(stats) = stats_part {
        result.push(' ');
        result.push_str(C_STATS);
        result.push_str(stats);
        result.push_str(C_RESET);
    }
}

/// Colorize a detail line (labels like Columns:, Filter:, Group By:, etc.)
fn colorize_detail(line: &str, result: &mut String) {
    let labels = [
        "Columns:",
        "Filter:",
        "Index Cond:",
        "Join Cond:",
        "Group By:",
        "Order By:",
        "Having:",
        "Limit:",
        "Offset:",
        "Using:",
        "Alias:",
        "Values:",
        "Set:",
        "Source:",
    ];

    for label in &labels {
        if let Some(rest) = line.strip_prefix(label) {
            result.push_str(C_LABEL);
            result.push_str(label);
            result.push_str(C_RESET);
            result.push_str(rest);
            return;
        }
    }

    // No recognized label — render as-is
    result.push_str(line);
}

fn format_value(value: &Value) -> String {
    match value {
        Value::Null(_) => "NULL".to_string(),
        Value::Integer(i) => i.to_string(),
        Value::Float(f) => format_float(*f),
        Value::Text(s) => s.to_string(),
        Value::Boolean(b) => if *b { "true" } else { "false" }.to_string(),
        Value::Timestamp(ts) => format_timestamp(ts),
        Value::Extension(data) if data.first() == Some(&(stoolap::DataType::Json as u8)) => {
            std::str::from_utf8(&data[1..]).unwrap_or("").to_string()
        }
        Value::Extension(data) if data.first() == Some(&(stoolap::DataType::Vector as u8)) => {
            stoolap::core::value::format_vector_bytes(&data[1..])
        }
        Value::Extension(_) => "<extension>".to_string(),
    }
}

/// Format a timestamp preserving sub-second precision.
/// Omits fractional part when zero, trims trailing zeros otherwise.
fn format_timestamp(ts: &chrono::DateTime<chrono::Utc>) -> String {
    let nanos = ts.timestamp_subsec_nanos();
    if nanos == 0 {
        return ts.format("%Y-%m-%dT%H:%M:%SZ").to_string();
    }
    // Format with 9-digit nanoseconds, then trim trailing zeros
    let s = ts.format("%Y-%m-%dT%H:%M:%S.%9fZ").to_string();
    // Find the 'Z' and trim zeros before it
    if let Some(z_pos) = s.rfind('Z') {
        let trimmed = s[..z_pos].trim_end_matches('0');
        format!("{}Z", trimmed)
    } else {
        s
    }
}

/// Format a float value consistently, using scientific notation for extreme values
fn format_float(v: f64) -> String {
    // Handle special cases
    if v.is_nan() {
        return "NaN".to_string();
    }
    if v.is_infinite() {
        return if v.is_sign_positive() {
            "Infinity"
        } else {
            "-Infinity"
        }
        .to_string();
    }

    let abs_v = v.abs();

    // Use scientific notation for very large or very small numbers
    if abs_v != 0.0 && !(1e-4..1e15).contains(&abs_v) {
        let s = format!("{:e}", v);
        // Clean up trailing zeros in mantissa
        if let Some(e_pos) = s.find('e') {
            let (mantissa, exp) = s.split_at(e_pos);
            let clean_mantissa = if mantissa.contains('.') {
                mantissa
                    .trim_end_matches('0')
                    .trim_end_matches('.')
                    .to_string()
            } else {
                mantissa.to_string()
            };
            return format!("{}{}", clean_mantissa, exp);
        }
        return s;
    }

    // Integer-like float
    if v == v.trunc() {
        format!("{:.1}", v)
    } else {
        // Normal range - show reasonable precision
        format!("{:.6}", v)
            .trim_end_matches('0')
            .trim_end_matches('.')
            .to_string()
    }
}

fn value_to_json(value: &Value) -> serde_json::Value {
    match value {
        Value::Null(_) => serde_json::Value::Null,
        Value::Integer(i) => serde_json::json!(i),
        Value::Float(f) => serde_json::json!(f),
        Value::Text(s) => serde_json::json!(s.as_str()),
        Value::Boolean(b) => serde_json::json!(b),
        Value::Timestamp(ts) => serde_json::json!(format_timestamp(ts)),
        Value::Extension(data) if data.first() == Some(&(stoolap::DataType::Json as u8)) => {
            serde_json::json!(std::str::from_utf8(&data[1..]).unwrap_or(""))
        }
        Value::Extension(data) if data.first() == Some(&(stoolap::DataType::Vector as u8)) => {
            serde_json::json!(stoolap::core::value::format_vector_bytes(&data[1..]))
        }
        Value::Extension(_) => serde_json::Value::Null,
    }
}

/// Split SQL statements by semicolons, handling quotes and comments
fn split_sql_statements(input: &str) -> Vec<String> {
    let mut statements = Vec::new();
    let mut current_statement = String::new();

    let mut in_single_quotes = false;
    let mut in_double_quotes = false;
    let mut in_line_comment = false;
    let mut in_block_comment = false;

    let chars: Vec<char> = input.chars().collect();
    let mut i = 0;

    while i < chars.len() {
        let char = chars[i];

        // Handle end of line comment
        if in_line_comment {
            if char == '\n' {
                in_line_comment = false;
                current_statement.push(char);
            }
            i += 1;
            continue;
        }

        // Handle start of line comment
        // Only treat -- as comment if followed by whitespace/newline/EOF
        // Otherwise, it's double negation (--val or --5)
        if !in_single_quotes
            && !in_double_quotes
            && !in_block_comment
            && char == '-'
            && i + 1 < chars.len()
            && chars[i + 1] == '-'
        {
            // Check what follows the second dash
            let after_second_dash = if i + 2 < chars.len() {
                chars[i + 2]
            } else {
                '\0' // End of input
            };
            // Only treat as comment if followed by whitespace, newline, or EOF
            if after_second_dash == '\0'
                || after_second_dash == ' '
                || after_second_dash == '\t'
                || after_second_dash == '\n'
                || after_second_dash == '\r'
            {
                in_line_comment = true;
                i += 2;
                continue;
            }
            // Otherwise, fall through and treat as two minus operators
        }

        // Handle end of block comment
        if in_block_comment {
            if char == '*' && i + 1 < chars.len() && chars[i + 1] == '/' {
                in_block_comment = false;
                i += 2;
                continue;
            }
            i += 1;
            continue;
        }

        // Handle start of block comment
        if !in_single_quotes
            && !in_double_quotes
            && char == '/'
            && i + 1 < chars.len()
            && chars[i + 1] == '*'
        {
            in_block_comment = true;
            i += 2;
            continue;
        }

        // Handle quotes
        if !in_block_comment && !in_line_comment {
            if char == '\'' && (i == 0 || chars[i - 1] != '\\') {
                in_single_quotes = !in_single_quotes;
            } else if char == '"' && (i == 0 || chars[i - 1] != '\\') {
                in_double_quotes = !in_double_quotes;
            }
        }

        // If we find a semicolon outside of quotes and comments, it's a statement delimiter
        if char == ';'
            && !in_single_quotes
            && !in_double_quotes
            && !in_block_comment
            && !in_line_comment
        {
            statements.push(current_statement.clone());
            current_statement.clear();
        } else {
            current_statement.push(char);
        }

        i += 1;
    }

    // Add any remaining statement
    if !current_statement.is_empty() {
        statements.push(current_statement);
    }

    statements
}

fn print_help_main() {
    println!("Stoolap SQL CLI");
    println!();
    println!("  SQL Commands:");
    println!("    SELECT ...             Execute a SELECT query");
    println!("    INSERT ...             Insert data into a table");
    println!("    UPDATE ...             Update data in a table");
    println!("    DELETE ...             Delete data from a table");
    println!("    CREATE TABLE ...       Create a new table");
    println!("    CREATE INDEX ...       Create an index on a column");
    println!("    SHOW TABLES            List all tables");
    println!("    SHOW CREATE TABLE ...  Show CREATE TABLE statement for a table");
    println!("    SHOW INDEXES FROM ...  Show indexes for a table");
    println!();
    println!("  Transaction Commands:");
    println!("    BEGIN                  Start a new transaction");
    println!("    COMMIT                 Commit the current transaction");
    println!("    ROLLBACK               Rollback the current transaction");
    println!();
    println!("  Special Commands:");
    println!("    help, \\h, \\?          Show this help message");
    println!();
}