zam 0.8.1

//! Database management for zam
//!
//! This module provides SQLite-based storage for command history with support for:
//! - Multi-host history tracking
//! - Session management
//! - Token/password storage for retrieval
//! - Import from shell history files

use crate::error::Result;
use crate::types::{CommandId, HostId, SessionId};
use chrono::{DateTime, Utc};
use rusqlite::{Connection, OptionalExtension, params};
use std::path::Path;
use uuid::Uuid;

/// Represents a host in the database
#[derive(Debug, Clone)]
pub struct Host {
    pub id: HostId,
    pub hostname: String,
    pub created_at: DateTime<Utc>,
}

/// Represents a shell session
#[derive(Debug, Clone)]
pub struct Session {
    pub id: SessionId,
    pub host_id: HostId,
    pub hostname: String,
    pub started_at: DateTime<Utc>,
    pub ended_at: Option<DateTime<Utc>>,
}

/// Represents a command entry in the database
#[derive(Debug, Clone, serde::Serialize)]
pub struct CommandEntry {
    pub id: CommandId,
    pub session_id: SessionId,
    pub command: String,
    pub timestamp: DateTime<Utc>,
    pub directory: String,
    pub redacted: bool,
    pub exit_code: Option<i32>,
}

/// Represents a redacted token that can be retrieved
#[derive(Debug, Clone)]
pub struct Token {
    pub id: i64,
    pub command_id: CommandId,
    pub token_type: String, // e.g., "password", "api_key", "token"
    pub placeholder: String,
    pub original_value: String,
    pub created_at: DateTime<Utc>,
}

/// Represents a shell alias
#[derive(Debug, Clone, serde::Serialize)]
pub struct Alias {
    pub alias: String,
    pub command: String,
    pub description: String,
    pub date_created: DateTime<Utc>,
    pub date_updated: DateTime<Utc>,
}

/// Represents a secret key loaded into a session (value is NOT stored)
#[derive(Debug, Clone)]
pub struct SessionSecret {
    pub id: i64,
    pub session_id: SessionId,
    pub key_name: String,
    pub source: String,
    pub loaded_at: DateTime<Utc>,
}

/// Statistics about the database
#[derive(Debug, Clone, Default)]
pub struct DatabaseStats {
    pub total_commands: usize,
    pub total_sessions: usize,
    pub total_hosts: usize,
    pub redacted_commands: usize,
    pub stored_tokens: usize,
    pub oldest_entry: Option<DateTime<Utc>>,
    pub newest_entry: Option<DateTime<Utc>>,
}

/// Main database manager
pub struct Database {
    conn: Connection,
    current_host_id: HostId,
    current_session_id: Option<SessionId>,
}

impl Database {
    /// Create a new database connection and initialize schema
    #[must_use = "Database connection must be used"]
    pub fn new(db_path: &Path) -> Result<Self> {
        // Create parent directory if it doesn't exist
        if let Some(parent) = db_path.parent() {
            std::fs::create_dir_all(parent)?;
        }

        let conn = Connection::open(db_path)?;

        // Enable WAL journal mode for better concurrent performance
        conn.execute_batch("PRAGMA journal_mode=WAL; PRAGMA synchronous=NORMAL;")?;

        // Enable foreign keys
        conn.execute("PRAGMA foreign_keys = ON", [])?;

        let mut db = Self {
            conn,
            current_host_id: HostId::new(0),
            current_session_id: None,
        };

        db.initialize_schema()?;
        db.ensure_current_host()?;

        Ok(db)
    }

    /// Initialize database schema
    fn initialize_schema(&self) -> Result<()> {
        // Hosts table
        self.conn.execute(
            "CREATE TABLE IF NOT EXISTS hosts (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                hostname TEXT NOT NULL UNIQUE,
                created_at TEXT NOT NULL
            )",
            [],
        )?;

        // Sessions table
        self.conn.execute(
            "CREATE TABLE IF NOT EXISTS sessions (
                id TEXT PRIMARY KEY,
                host_id INTEGER NOT NULL,
                started_at TEXT NOT NULL,
                ended_at TEXT,
                FOREIGN KEY (host_id) REFERENCES hosts(id) ON DELETE CASCADE
            )",
            [],
        )?;

        // Commands table
        self.conn.execute(
            "CREATE TABLE IF NOT EXISTS commands (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                session_id TEXT NOT NULL,
                command TEXT NOT NULL,
                timestamp TEXT NOT NULL,
                directory TEXT NOT NULL,
                redacted INTEGER NOT NULL DEFAULT 0,
                exit_code INTEGER,
                FOREIGN KEY (session_id) REFERENCES sessions(id) ON DELETE CASCADE
            )",
            [],
        )?;

        // Tokens table - stores redacted values for retrieval
        self.conn.execute(
            "CREATE TABLE IF NOT EXISTS tokens (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                command_id INTEGER NOT NULL,
                token_type TEXT NOT NULL,
                placeholder TEXT NOT NULL,
                original_value TEXT NOT NULL,
                created_at TEXT NOT NULL,
                FOREIGN KEY (command_id) REFERENCES commands(id) ON DELETE CASCADE
            )",
            [],
        )?;

        // Aliases table
        self.conn.execute(
            "CREATE TABLE IF NOT EXISTS aliases (
                alias TEXT PRIMARY KEY,
                command TEXT NOT NULL,
                description TEXT NOT NULL,
                date_created TEXT NOT NULL,
                date_updated TEXT NOT NULL
            )",
            [],
        )?;

        // Create indices for common queries
        self.conn.execute(
            "CREATE INDEX IF NOT EXISTS idx_commands_timestamp ON commands(timestamp DESC)",
            [],
        )?;

        self.conn.execute(
            "CREATE INDEX IF NOT EXISTS idx_commands_session ON commands(session_id)",
            [],
        )?;

        self.conn.execute(
            "CREATE INDEX IF NOT EXISTS idx_commands_directory ON commands(directory)",
            [],
        )?;

        self.conn.execute(
            "CREATE INDEX IF NOT EXISTS idx_tokens_command ON tokens(command_id)",
            [],
        )?;

        self.conn.execute(
            "CREATE INDEX IF NOT EXISTS idx_sessions_host ON sessions(host_id)",
            [],
        )?;

        self.conn.execute(
            "CREATE INDEX IF NOT EXISTS idx_commands_command ON commands(command)",
            [],
        )?;

        // Session secrets table - tracks keys loaded from external sources (values NOT stored)
        self.conn.execute(
            "CREATE TABLE IF NOT EXISTS session_secrets (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                session_id TEXT NOT NULL,
                key_name TEXT NOT NULL,
                source TEXT NOT NULL,
                loaded_at TEXT NOT NULL,
                FOREIGN KEY (session_id) REFERENCES sessions(id) ON DELETE CASCADE
            )",
            [],
        )?;

        self.conn.execute(
            "CREATE INDEX IF NOT EXISTS idx_session_secrets_session ON session_secrets(session_id)",
            [],
        )?;

        // Preferences table - key/value store for TUI settings
        self.conn.execute(
            "CREATE TABLE IF NOT EXISTS preferences (
                key TEXT PRIMARY KEY,
                value TEXT NOT NULL
            )",
            [],
        )?;

        Ok(())
    }

    /// Ensure the current host exists in the database
    fn ensure_current_host(&mut self) -> Result<()> {
        let hostname = hostname::get()
            .map(|h| h.to_string_lossy().to_string())
            .unwrap_or_else(|_| "unknown".to_string());

        // Try to find existing host
        let host_id: Option<i64> = self
            .conn
            .query_row(
                "SELECT id FROM hosts WHERE hostname = ?1",
                params![hostname],
                |row| row.get(0),
            )
            .optional()?;

        self.current_host_id = if let Some(id) = host_id {
            HostId::new(id)
        } else {
            // Insert new host
            let now = Utc::now().to_rfc3339();
            self.conn.execute(
                "INSERT INTO hosts (hostname, created_at) VALUES (?1, ?2)",
                params![hostname, now],
            )?;
            HostId::new(self.conn.last_insert_rowid())
        };

        Ok(())
    }

    /// Start a new session
    pub fn start_session(&mut self) -> Result<String> {
        let session_id = Uuid::new_v4().to_string();
        let now = Utc::now().to_rfc3339();

        self.conn.execute(
            "INSERT INTO sessions (id, host_id, started_at) VALUES (?1, ?2, ?3)",
            params![session_id, self.current_host_id.as_i64(), now],
        )?;

        self.current_session_id = Some(SessionId::new(session_id.clone()));
        Ok(session_id)
    }

    /// End the current session
    pub fn end_session(&mut self, session_id: &str) -> Result<()> {
        let now = Utc::now().to_rfc3339();
        self.conn.execute(
            "UPDATE sessions SET ended_at = ?1 WHERE id = ?2",
            params![now, session_id],
        )?;

        if self.current_session_id.as_deref() == Some(session_id) {
            self.current_session_id = None;
        }

        Ok(())
    }

    /// Get or create a session for the current shell
    pub fn ensure_session(&mut self) -> Result<String> {
        if let Some(ref session_id) = self.current_session_id {
            Ok(session_id.as_str().to_string())
        } else {
            self.start_session()
        }
    }

    /// Resume an existing session by ID, or create one with that ID if it doesn't exist.
    /// Useful for static sessions (e.g. Claude Code) where multiple `zam log` invocations
    /// should share the same session.
    pub fn resume_session(&mut self, session_id: &str) -> Result<()> {
        let exists: bool = self.conn.query_row(
            "SELECT EXISTS(SELECT 1 FROM sessions WHERE id = ?1)",
            params![session_id],
            |row| row.get(0),
        )?;

        if !exists {
            let now = Utc::now().to_rfc3339();
            self.conn.execute(
                "INSERT INTO sessions (id, host_id, started_at) VALUES (?1, ?2, ?3)",
                params![session_id, self.current_host_id.as_i64(), now],
            )?;
        }

        self.current_session_id = Some(SessionId::new(session_id.to_string()));
        Ok(())
    }

    /// Add a command to the database
    pub fn add_command(
        &mut self,
        command: &str,
        directory: &str,
        timestamp: DateTime<Utc>,
        redacted: bool,
        exit_code: Option<i32>,
    ) -> Result<i64> {
        let session_id = self.ensure_session()?;
        let timestamp_str = timestamp.to_rfc3339();

        self.conn.execute(
            "INSERT INTO commands (session_id, command, timestamp, directory, redacted, exit_code)
             VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
            params![
                session_id,
                command,
                timestamp_str,
                directory,
                redacted as i32,
                exit_code
            ],
        )?;

        Ok(self.conn.last_insert_rowid())
    }

    /// Store a redacted token for later retrieval
    pub fn store_token(
        &self,
        command_id: i64,
        token_type: &str,
        placeholder: &str,
        original_value: &str,
    ) -> Result<i64> {
        let now = Utc::now().to_rfc3339();

        self.conn.execute(
            "INSERT INTO tokens (command_id, token_type, placeholder, original_value, created_at)
             VALUES (?1, ?2, ?3, ?4, ?5)",
            params![command_id, token_type, placeholder, original_value, now],
        )?;

        Ok(self.conn.last_insert_rowid())
    }

    /// Get tokens for a specific command
    #[must_use = "Token query results should be used"]
    pub fn get_tokens_for_command(&self, command_id: CommandId) -> Result<Vec<Token>> {
        let mut stmt = self.conn.prepare(
            "SELECT id, command_id, token_type, placeholder, original_value, created_at
             FROM tokens WHERE command_id = ?1",
        )?;

        let tokens = stmt
            .query_map(params![command_id.as_i64()], |row| {
                Ok(Token {
                    id: row.get(0)?,
                    command_id: CommandId::new(row.get(1)?),
                    token_type: row.get(2)?,
                    placeholder: row.get(3)?,
                    original_value: row.get(4)?,
                    created_at: row
                        .get::<_, String>(5)?
                        .parse()
                        .unwrap_or_else(|_| Utc::now()),
                })
            })?
            .collect::<rusqlite::Result<Vec<_>>>()?;

        Ok(tokens)
    }

    /// Get tokens by session
    pub fn get_tokens_by_session(&self, session_id: &str) -> Result<Vec<Token>> {
        let mut stmt = self.conn.prepare(
            "SELECT t.id, t.command_id, t.token_type, t.placeholder, t.original_value, t.created_at
             FROM tokens t
             JOIN commands c ON t.command_id = c.id
             WHERE c.session_id = ?1
             ORDER BY t.created_at DESC",
        )?;

        let tokens = stmt
            .query_map(params![session_id], |row| {
                Ok(Token {
                    id: row.get(0)?,
                    command_id: row.get(1)?,
                    token_type: row.get(2)?,
                    placeholder: row.get(3)?,
                    original_value: row.get(4)?,
                    created_at: row
                        .get::<_, String>(5)?
                        .parse()
                        .unwrap_or_else(|_| Utc::now()),
                })
            })?
            .collect::<rusqlite::Result<Vec<_>>>()?;

        Ok(tokens)
    }

    /// Get tokens by directory
    pub fn get_tokens_by_directory(&self, directory: &str) -> Result<Vec<Token>> {
        let mut stmt = self.conn.prepare(
            "SELECT t.id, t.command_id, t.token_type, t.placeholder, t.original_value, t.created_at
             FROM tokens t
             JOIN commands c ON t.command_id = c.id
             WHERE c.directory = ?1
             ORDER BY t.created_at DESC",
        )?;

        let tokens = stmt
            .query_map(params![directory], |row| {
                Ok(Token {
                    id: row.get(0)?,
                    command_id: row.get(1)?,
                    token_type: row.get(2)?,
                    placeholder: row.get(3)?,
                    original_value: row.get(4)?,
                    created_at: row
                        .get::<_, String>(5)?
                        .parse()
                        .unwrap_or_else(|_| Utc::now()),
                })
            })?
            .collect::<rusqlite::Result<Vec<_>>>()?;

        Ok(tokens)
    }

    /// Search commands
    #[must_use = "Search results should be used"]
    pub fn search_commands(
        &self,
        query: &str,
        directory_filter: Option<&str>,
        host_filter: Option<&str>,
        limit: Option<usize>,
    ) -> Result<Vec<CommandEntry>> {
        let mut sql = String::from(
            "SELECT c.id, c.session_id, c.command, c.timestamp, c.directory, c.redacted, c.exit_code
             FROM commands c
             JOIN sessions s ON c.session_id = s.id
             JOIN hosts h ON s.host_id = h.id
             WHERE c.command LIKE ?1",
        );

        let mut params: Vec<Box<dyn rusqlite::ToSql>> = vec![Box::new(format!("%{}%", query))];

        if let Some(dir) = directory_filter {
            sql.push_str(" AND c.directory LIKE ?");
            params.push(Box::new(format!("%{}%", dir)));
        }

        if let Some(host) = host_filter {
            sql.push_str(" AND h.hostname = ?");
            params.push(Box::new(host.to_string()));
        }

        sql.push_str(" ORDER BY c.timestamp DESC");

        if let Some(lim) = limit {
            sql.push_str(" LIMIT ?");
            params.push(Box::new(lim as i64));
        }

        let mut stmt = self.conn.prepare(&sql)?;
        let param_refs: Vec<&dyn rusqlite::ToSql> = params.iter().map(|b| b.as_ref()).collect();

        let commands = stmt
            .query_map(param_refs.as_slice(), |row| {
                Ok(CommandEntry {
                    id: row.get(0)?,
                    session_id: row.get(1)?,
                    command: row.get(2)?,
                    timestamp: row
                        .get::<_, String>(3)?
                        .parse()
                        .unwrap_or_else(|_| Utc::now()),
                    directory: row.get(4)?,
                    redacted: row.get::<_, i32>(5)? != 0,
                    exit_code: row.get(6)?,
                })
            })?
            .collect::<rusqlite::Result<Vec<_>>>()?;

        Ok(commands)
    }

    /// Get recent commands
    #[must_use = "Query results should be used"]
    pub fn get_recent_commands(&self, limit: usize) -> Result<Vec<CommandEntry>> {
        let mut stmt = self.conn.prepare(
            "SELECT id, session_id, command, timestamp, directory, redacted, exit_code
             FROM commands
             ORDER BY timestamp DESC
             LIMIT ?1",
        )?;

        let commands = stmt
            .query_map(params![limit as i64], |row| {
                Ok(CommandEntry {
                    id: row.get(0)?,
                    session_id: row.get(1)?,
                    command: row.get(2)?,
                    timestamp: row
                        .get::<_, String>(3)?
                        .parse()
                        .unwrap_or_else(|_| Utc::now()),
                    directory: row.get(4)?,
                    redacted: row.get::<_, i32>(5)? != 0,
                    exit_code: row.get(6)?,
                })
            })?
            .collect::<rusqlite::Result<Vec<_>>>()?;

        Ok(commands)
    }

    /// Get all commands (for export/migration)
    #[must_use = "Query results should be used"]
    pub fn get_all_commands(&self) -> Result<Vec<CommandEntry>> {
        let mut stmt = self.conn.prepare(
            "SELECT id, session_id, command, timestamp, directory, redacted, exit_code
             FROM commands
             ORDER BY timestamp ASC",
        )?;

        let commands = stmt
            .query_map([], |row| {
                Ok(CommandEntry {
                    id: row.get(0)?,
                    session_id: row.get(1)?,
                    command: row.get(2)?,
                    timestamp: row
                        .get::<_, String>(3)?
                        .parse()
                        .unwrap_or_else(|_| Utc::now()),
                    directory: row.get(4)?,
                    redacted: row.get::<_, i32>(5)? != 0,
                    exit_code: row.get(6)?,
                })
            })?
            .collect::<rusqlite::Result<Vec<_>>>()?;

        Ok(commands)
    }

    /// Get commands excluding imported, with pagination (most recent first)
    pub fn get_commands_paginated(&self, offset: usize, limit: usize) -> Result<Vec<CommandEntry>> {
        let mut stmt = self.conn.prepare(
            "SELECT id, session_id, command, timestamp, directory, redacted, exit_code
             FROM commands
             WHERE directory != '<imported>'
             ORDER BY timestamp DESC
             LIMIT ?1 OFFSET ?2",
        )?;

        let commands = stmt
            .query_map(rusqlite::params![limit as i64, offset as i64], |row| {
                Ok(CommandEntry {
                    id: row.get(0)?,
                    session_id: row.get(1)?,
                    command: row.get(2)?,
                    timestamp: row
                        .get::<_, String>(3)?
                        .parse()
                        .unwrap_or_else(|_| Utc::now()),
                    directory: row.get(4)?,
                    redacted: row.get::<_, i32>(5)? != 0,
                    exit_code: row.get(6)?,
                })
            })?
            .collect::<rusqlite::Result<Vec<_>>>()?;

        Ok(commands)
    }

    /// Get unique commands per directory, paginated (most recent first).
    /// Groups by (command, directory), keeping the latest timestamp and highest id.
    pub fn get_unique_commands_paginated(
        &self,
        offset: usize,
        limit: usize,
    ) -> Result<Vec<CommandEntry>> {
        self.get_unique_commands_filtered(offset, limit, None)
    }

    /// Get unique commands with optional filter, paginated
    pub fn get_unique_commands_filtered(
        &self,
        offset: usize,
        limit: usize,
        filter: Option<&str>,
    ) -> Result<Vec<CommandEntry>> {
        let (where_clause, params): (&str, Vec<Box<dyn rusqlite::types::ToSql>>) = match filter {
            Some(f) if !f.is_empty() => (
                "WHERE directory != '<imported>' AND command LIKE ?3",
                vec![
                    Box::new(limit as i64),
                    Box::new(offset as i64),
                    Box::new(format!("%{}%", f)),
                ],
            ),
            _ => (
                "WHERE directory != '<imported>'",
                vec![Box::new(limit as i64), Box::new(offset as i64)],
            ),
        };

        let sql = format!(
            "SELECT id, session_id, command, timestamp, directory, redacted, exit_code
             FROM commands
             WHERE id IN (
                 SELECT MAX(id) FROM commands {where_clause} GROUP BY command
             )
             ORDER BY timestamp DESC
             LIMIT ?1 OFFSET ?2"
        );

        let mut stmt = self.conn.prepare(&sql)?;
        let params_ref: Vec<&dyn rusqlite::types::ToSql> =
            params.iter().map(|p| p.as_ref()).collect();

        let commands = stmt
            .query_map(params_ref.as_slice(), |row| {
                Ok(CommandEntry {
                    id: row.get(0)?,
                    session_id: row.get(1)?,
                    command: row.get(2)?,
                    timestamp: row
                        .get::<_, String>(3)?
                        .parse()
                        .unwrap_or_else(|_| Utc::now()),
                    directory: row.get(4)?,
                    redacted: row.get::<_, i32>(5)? != 0,
                    exit_code: row.get(6)?,
                })
            })?
            .collect::<rusqlite::Result<Vec<_>>>()?;

        Ok(commands)
    }

    /// Count unique (command, directory) pairs excluding imported
    pub fn count_unique_commands(&self) -> Result<usize> {
        self.count_unique_commands_filtered(None)
    }

    /// Count unique commands with optional filter
    pub fn count_unique_commands_filtered(&self, filter: Option<&str>) -> Result<usize> {
        let (where_extra, params): (&str, Vec<Box<dyn rusqlite::types::ToSql>>) = match filter {
            Some(f) if !f.is_empty() => {
                (" AND command LIKE ?1", vec![Box::new(format!("%{}%", f))])
            }
            _ => ("", vec![]),
        };

        let sql = format!(
            "SELECT COUNT(*) FROM (
                SELECT 1 FROM commands
                WHERE directory != '<imported>'{where_extra}
                GROUP BY command
            )"
        );

        let params_ref: Vec<&dyn rusqlite::types::ToSql> =
            params.iter().map(|p| p.as_ref()).collect();
        let count: i64 = self
            .conn
            .query_row(&sql, params_ref.as_slice(), |row| row.get(0))?;
        Ok(count as usize)
    }

    /// Get unique commands for a specific directory (no duplicates, most recent first)
    pub fn get_commands_for_directory(&self, directory: &str) -> Result<Vec<CommandEntry>> {
        let mut stmt = self.conn.prepare(
            "SELECT MAX(id), session_id, command, MAX(timestamp) as ts, directory, redacted, exit_code
             FROM commands
             WHERE directory = ?1
             GROUP BY command
             ORDER BY ts DESC",
        )?;

        let commands = stmt
            .query_map(rusqlite::params![directory], |row| {
                Ok(CommandEntry {
                    id: row.get(0)?,
                    session_id: row.get(1)?,
                    command: row.get(2)?,
                    timestamp: row
                        .get::<_, String>(3)?
                        .parse()
                        .unwrap_or_else(|_| Utc::now()),
                    directory: row.get(4)?,
                    redacted: row.get::<_, i32>(5)? != 0,
                    exit_code: row.get(6)?,
                })
            })?
            .collect::<rusqlite::Result<Vec<_>>>()?;

        Ok(commands)
    }

    /// Get all commands for a specific session
    /// Count commands for a batch of session IDs (single query)
    pub fn count_commands_for_sessions(&self, session_ids: &[&str]) -> Result<Vec<usize>> {
        if session_ids.is_empty() {
            return Ok(Vec::new());
        }
        let placeholders: Vec<String> = (1..=session_ids.len()).map(|i| format!("?{i}")).collect();
        let sql = format!(
            "SELECT session_id, COUNT(*) FROM commands WHERE session_id IN ({}) GROUP BY session_id",
            placeholders.join(",")
        );
        let mut stmt = self.conn.prepare(&sql)?;
        let params: Vec<&dyn rusqlite::types::ToSql> = session_ids
            .iter()
            .map(|s| s as &dyn rusqlite::types::ToSql)
            .collect();
        let mut counts: std::collections::HashMap<String, usize> = std::collections::HashMap::new();
        let rows = stmt.query_map(params.as_slice(), |row| {
            Ok((row.get::<_, String>(0)?, row.get::<_, i64>(1)? as usize))
        })?;
        for r in rows {
            let (sid, cnt) = r?;
            counts.insert(sid, cnt);
        }
        Ok(session_ids
            .iter()
            .map(|sid| *counts.get(*sid).unwrap_or(&0))
            .collect())
    }

    /// Get all commands for a specific session
    pub fn get_commands_for_session(&self, session_id: &str) -> Result<Vec<CommandEntry>> {
        let mut stmt = self.conn.prepare(
            "SELECT id, session_id, command, timestamp, directory, redacted, exit_code
             FROM commands
             WHERE session_id = ?1
             ORDER BY timestamp DESC",
        )?;

        let commands = stmt
            .query_map(rusqlite::params![session_id], |row| {
                Ok(CommandEntry {
                    id: row.get(0)?,
                    session_id: row.get(1)?,
                    command: row.get(2)?,
                    timestamp: row
                        .get::<_, String>(3)?
                        .parse()
                        .unwrap_or_else(|_| Utc::now()),
                    directory: row.get(4)?,
                    redacted: row.get::<_, i32>(5)? != 0,
                    exit_code: row.get(6)?,
                })
            })?
            .collect::<rusqlite::Result<Vec<_>>>()?;

        Ok(commands)
    }

    /// Get the most frequently used unique commands globally
    pub fn get_frequent_commands(&self, limit: usize) -> Result<Vec<(String, usize)>> {
        let mut stmt = self.conn.prepare(
            "SELECT command, COUNT(*) as cnt
             FROM commands
             WHERE directory != '<imported>'
             GROUP BY command
             ORDER BY cnt DESC
             LIMIT ?1",
        )?;

        let results = stmt
            .query_map(rusqlite::params![limit as i64], |row| {
                Ok((row.get::<_, String>(0)?, row.get::<_, i64>(1)? as usize))
            })?
            .collect::<rusqlite::Result<Vec<_>>>()?;

        Ok(results)
    }

    /// Count non-imported commands
    pub fn count_commands(&self) -> Result<usize> {
        let count: i64 = self.conn.query_row(
            "SELECT COUNT(*) FROM commands WHERE directory != '<imported>'",
            [],
            |row| row.get(0),
        )?;
        Ok(count as usize)
    }

    /// Get database statistics
    pub fn get_stats(&self) -> Result<DatabaseStats> {
        let total_commands: i64 =
            self.conn
                .query_row("SELECT COUNT(*) FROM commands", [], |row| row.get(0))?;

        let total_sessions: i64 =
            self.conn
                .query_row("SELECT COUNT(*) FROM sessions", [], |row| row.get(0))?;

        let total_hosts: i64 = self
            .conn
            .query_row("SELECT COUNT(*) FROM hosts", [], |row| row.get(0))?;

        let redacted_commands: i64 = self.conn.query_row(
            "SELECT COUNT(*) FROM commands WHERE redacted = 1",
            [],
            |row| row.get(0),
        )?;

        let stored_tokens: i64 = self
            .conn
            .query_row("SELECT COUNT(*) FROM tokens", [], |row| row.get(0))?;

        let oldest_entry: Option<String> = self
            .conn
            .query_row(
                "SELECT timestamp FROM commands ORDER BY timestamp ASC LIMIT 1",
                [],
                |row| row.get(0),
            )
            .optional()?;

        let newest_entry: Option<String> = self
            .conn
            .query_row(
                "SELECT timestamp FROM commands ORDER BY timestamp DESC LIMIT 1",
                [],
                |row| row.get(0),
            )
            .optional()?;

        Ok(DatabaseStats {
            total_commands: total_commands as usize,
            total_sessions: total_sessions as usize,
            total_hosts: total_hosts as usize,
            redacted_commands: redacted_commands as usize,
            stored_tokens: stored_tokens as usize,
            oldest_entry: oldest_entry.and_then(|s| s.parse().ok()),
            newest_entry: newest_entry.and_then(|s| s.parse().ok()),
        })
    }

    /// Get all hosts
    pub fn get_hosts(&self) -> Result<Vec<Host>> {
        let mut stmt = self
            .conn
            .prepare("SELECT id, hostname, created_at FROM hosts ORDER BY hostname")?;

        let hosts = stmt
            .query_map([], |row| {
                Ok(Host {
                    id: row.get(0)?,
                    hostname: row.get(1)?,
                    created_at: row
                        .get::<_, String>(2)?
                        .parse()
                        .unwrap_or_else(|_| Utc::now()),
                })
            })?
            .collect::<rusqlite::Result<Vec<_>>>()?;

        Ok(hosts)
    }

    /// Get sessions for a host
    pub fn get_sessions_for_host(&self, host_id: HostId) -> Result<Vec<Session>> {
        let mut stmt = self.conn.prepare(
            "SELECT s.id, s.host_id, COALESCE(h.hostname, '?'), s.started_at, s.ended_at
             FROM sessions s
             LEFT JOIN hosts h ON s.host_id = h.id
             WHERE s.host_id = ?1
             ORDER BY s.started_at DESC",
        )?;

        let sessions = stmt
            .query_map(params![host_id.as_i64()], |row| {
                Ok(Session {
                    id: SessionId::new(row.get(0)?),
                    host_id: HostId::new(row.get(1)?),
                    hostname: row.get(2)?,
                    started_at: row
                        .get::<_, String>(3)?
                        .parse()
                        .unwrap_or_else(|_| Utc::now()),
                    ended_at: row
                        .get::<_, Option<String>>(4)?
                        .and_then(|s| s.parse().ok()),
                })
            })?
            .collect::<rusqlite::Result<Vec<_>>>()?;

        Ok(sessions)
    }

    /// Import from bash history
    pub fn import_from_bash_history(&mut self, bash_history_path: &Path) -> Result<usize> {
        let content = std::fs::read_to_string(bash_history_path)?;
        let mut imported_count = 0;
        let now = Utc::now();

        for line in content.lines() {
            let line = line.trim();
            if line.is_empty() || line.starts_with('#') {
                continue;
            }

            self.add_command(line, "<imported>", now, false, None)?;
            imported_count += 1;
        }

        Ok(imported_count)
    }

    /// Import from zsh history
    pub fn import_from_zsh_history(&mut self, zsh_history_path: &Path) -> Result<usize> {
        let content = std::fs::read_to_string(zsh_history_path)?;
        let mut imported_count = 0;

        // Zsh format: ": 1609786800:0;command"
        let re = regex::Regex::new(r"^: (\d+):\d+;(.*)").unwrap();

        for line in content.lines() {
            if let Some(caps) = re.captures(line) {
                let timestamp_str = caps.get(1).unwrap().as_str();
                let command = caps.get(2).unwrap().as_str();

                if let Ok(timestamp_secs) = timestamp_str.parse::<i64>()
                    && let Some(datetime) = DateTime::from_timestamp(timestamp_secs, 0)
                {
                    self.add_command(command, "<imported>", datetime, false, None)?;
                    imported_count += 1;
                }
            }
        }

        Ok(imported_count)
    }

    /// Merge another database into this one
    pub fn merge_from_database(&mut self, other_db_path: &Path) -> Result<usize> {
        let other_conn = Connection::open(other_db_path)?;
        let mut imported_count = 0;

        // Get all commands from the other database
        let mut stmt = other_conn.prepare(
            "SELECT c.command, c.timestamp, c.directory, c.redacted, c.exit_code,
                    s.started_at, h.hostname
             FROM commands c
             JOIN sessions s ON c.session_id = s.id
             JOIN hosts h ON s.host_id = h.id
             ORDER BY c.timestamp ASC",
        )?;

        let commands: Vec<_> = stmt
            .query_map([], |row| {
                Ok((
                    row.get::<_, String>(0)?,
                    row.get::<_, String>(1)?,
                    row.get::<_, String>(2)?,
                    row.get::<_, i32>(3)? != 0,
                    row.get::<_, Option<i32>>(4)?,
                ))
            })?
            .collect::<rusqlite::Result<Vec<_>>>()?;

        for (command, timestamp_str, directory, redacted, exit_code) in commands {
            if let Ok(timestamp) = timestamp_str.parse() {
                self.add_command(&command, &directory, timestamp, redacted, exit_code)?;
                imported_count += 1;
            }
        }

        Ok(imported_count)
    }

    /// Add a new alias
    pub fn add_alias(&self, alias: &str, command: &str, description: &str) -> Result<()> {
        let now = Utc::now().to_rfc3339();
        self.conn.execute(
            "INSERT INTO aliases (alias, command, description, date_created, date_updated)
             VALUES (?1, ?2, ?3, ?4, ?5)",
            params![alias, command, description, now, now],
        )?;
        Ok(())
    }

    /// Update an existing alias
    pub fn update_alias(
        &self,
        alias: &str,
        command: &str,
        description: Option<&str>,
    ) -> Result<()> {
        let now = Utc::now().to_rfc3339();
        if let Some(desc) = description {
            self.conn.execute(
                "UPDATE aliases SET command = ?1, description = ?2, date_updated = ?3
                 WHERE alias = ?4",
                params![command, desc, now, alias],
            )?;
        } else {
            self.conn.execute(
                "UPDATE aliases SET command = ?1, date_updated = ?2 WHERE alias = ?3",
                params![command, now, alias],
            )?;
        }
        Ok(())
    }

    /// Remove an alias
    pub fn remove_alias(&self, alias: &str) -> Result<()> {
        self.conn
            .execute("DELETE FROM aliases WHERE alias = ?1", params![alias])?;
        Ok(())
    }

    /// List all aliases ordered by name
    #[must_use = "Alias list should be used"]
    pub fn list_aliases(&self) -> Result<Vec<Alias>> {
        let mut stmt = self.conn.prepare(
            "SELECT alias, command, description, date_created, date_updated
             FROM aliases ORDER BY alias ASC",
        )?;

        let aliases = stmt
            .query_map([], |row| {
                Ok(Alias {
                    alias: row.get(0)?,
                    command: row.get(1)?,
                    description: row.get(2)?,
                    date_created: row
                        .get::<_, String>(3)?
                        .parse()
                        .unwrap_or_else(|_| Utc::now()),
                    date_updated: row
                        .get::<_, String>(4)?
                        .parse()
                        .unwrap_or_else(|_| Utc::now()),
                })
            })?
            .collect::<rusqlite::Result<Vec<_>>>()?;

        Ok(aliases)
    }

    /// Upsert aliases from shell environment (sync)
    /// Returns the number of aliases upserted
    pub fn sync_aliases(&self, aliases: &[(String, String)]) -> Result<usize> {
        let now = Utc::now().to_rfc3339();
        let mut count = 0;

        for (name, command) in aliases {
            self.conn.execute(
                "INSERT INTO aliases (alias, command, description, date_created, date_updated)
                 VALUES (?1, ?2, '', ?3, ?4)
                 ON CONFLICT(alias) DO UPDATE SET command = ?2, date_updated = ?4",
                params![name, command, now, now],
            )?;
            count += 1;
        }

        Ok(count)
    }

    /// Clear all data (for testing)
    pub fn clear(&self) -> Result<()> {
        self.conn.execute("DELETE FROM tokens", [])?;
        self.conn.execute("DELETE FROM commands", [])?;
        self.conn.execute("DELETE FROM sessions", [])?;
        self.conn.execute("DELETE FROM hosts", [])?;
        Ok(())
    }

    /// Delete a specific command by ID
    pub fn delete_command(&self, id: CommandId) -> Result<()> {
        self.conn
            .execute("DELETE FROM commands WHERE id = ?1", [id.0])?;
        Ok(())
    }

    /// Get all sessions across all hosts
    pub fn get_all_sessions(&self) -> Result<Vec<Session>> {
        let mut stmt = self.conn.prepare(
            "SELECT s.id, s.host_id, COALESCE(h.hostname, '?'), s.started_at, s.ended_at
             FROM sessions s
             LEFT JOIN hosts h ON s.host_id = h.id
             ORDER BY s.started_at DESC",
        )?;

        let sessions = stmt
            .query_map([], |row| {
                Ok(Session {
                    id: SessionId::new(row.get(0)?),
                    host_id: HostId::new(row.get(1)?),
                    hostname: row.get(2)?,
                    started_at: row
                        .get::<_, String>(3)?
                        .parse()
                        .unwrap_or_else(|_| Utc::now()),
                    ended_at: row
                        .get::<_, Option<String>>(4)?
                        .and_then(|s| s.parse().ok()),
                })
            })?
            .collect::<rusqlite::Result<Vec<_>>>()?;

        Ok(sessions)
    }

    /// Count total sessions
    pub fn count_sessions(&self) -> Result<usize> {
        self.count_sessions_filtered(None)
    }

    /// Count sessions with optional filter on session id or hostname
    pub fn count_sessions_filtered(&self, filter: Option<&str>) -> Result<usize> {
        let (where_clause, params): (&str, Vec<Box<dyn rusqlite::types::ToSql>>) = match filter {
            Some(f) if !f.is_empty() => (
                "WHERE s.id LIKE ?1 OR h.hostname LIKE ?1",
                vec![Box::new(format!("%{}%", f))],
            ),
            _ => ("", vec![]),
        };

        let sql = format!(
            "SELECT COUNT(*) FROM sessions s
             LEFT JOIN hosts h ON s.host_id = h.id
             {where_clause}"
        );

        let params_ref: Vec<&dyn rusqlite::types::ToSql> =
            params.iter().map(|p| p.as_ref()).collect();
        let count: i64 = self
            .conn
            .query_row(&sql, params_ref.as_slice(), |row| row.get(0))?;
        Ok(count as usize)
    }

    /// Get sessions with pagination
    pub fn get_sessions_paginated(&self, offset: usize, limit: usize) -> Result<Vec<Session>> {
        self.get_sessions_filtered(offset, limit, None)
    }

    /// Get sessions with optional filter, paginated
    pub fn get_sessions_filtered(
        &self,
        offset: usize,
        limit: usize,
        filter: Option<&str>,
    ) -> Result<Vec<Session>> {
        let (where_clause, params): (&str, Vec<Box<dyn rusqlite::types::ToSql>>) = match filter {
            Some(f) if !f.is_empty() => (
                "WHERE s.id LIKE ?3 OR h.hostname LIKE ?3",
                vec![
                    Box::new(limit as i64),
                    Box::new(offset as i64),
                    Box::new(format!("%{}%", f)),
                ],
            ),
            _ => ("", vec![Box::new(limit as i64), Box::new(offset as i64)]),
        };

        let sql = format!(
            "SELECT s.id, s.host_id, COALESCE(h.hostname, '?'), s.started_at, s.ended_at
             FROM sessions s
             LEFT JOIN hosts h ON s.host_id = h.id
             {where_clause}
             ORDER BY s.started_at DESC
             LIMIT ?1 OFFSET ?2"
        );

        let mut stmt = self.conn.prepare(&sql)?;
        let params_ref: Vec<&dyn rusqlite::types::ToSql> =
            params.iter().map(|p| p.as_ref()).collect();

        let sessions = stmt
            .query_map(params_ref.as_slice(), |row| {
                Ok(Session {
                    id: SessionId::new(row.get(0)?),
                    host_id: HostId::new(row.get(1)?),
                    hostname: row.get(2)?,
                    started_at: row
                        .get::<_, String>(3)?
                        .parse()
                        .unwrap_or_else(|_| Utc::now()),
                    ended_at: row
                        .get::<_, Option<String>>(4)?
                        .and_then(|s| s.parse().ok()),
                })
            })?
            .collect::<rusqlite::Result<Vec<_>>>()?;

        Ok(sessions)
    }

    /// Get all tokens
    pub fn get_all_tokens(&self) -> Result<Vec<Token>> {
        let mut stmt = self.conn.prepare(
            "SELECT id, command_id, token_type, placeholder, original_value, created_at
             FROM tokens
             ORDER BY created_at DESC",
        )?;

        let tokens = stmt
            .query_map([], |row| {
                Ok(Token {
                    id: row.get(0)?,
                    command_id: CommandId::new(row.get(1)?),
                    token_type: row.get(2)?,
                    placeholder: row.get(3)?,
                    original_value: row.get(4)?,
                    created_at: row
                        .get::<_, String>(5)?
                        .parse()
                        .unwrap_or_else(|_| Utc::now()),
                })
            })?
            .collect::<rusqlite::Result<Vec<_>>>()?;

        Ok(tokens)
    }

    /// Delete a host and cascade to sessions/commands/tokens
    pub fn delete_host(&self, id: HostId) -> Result<()> {
        self.conn
            .execute("DELETE FROM hosts WHERE id = ?1", [id.as_i64()])?;
        Ok(())
    }

    /// Delete a session and cascade to commands/tokens
    pub fn delete_session(&self, id: &str) -> Result<()> {
        self.conn
            .execute("DELETE FROM sessions WHERE id = ?1", [id])?;
        Ok(())
    }

    /// Delete a specific token
    pub fn delete_token(&self, id: i64) -> Result<()> {
        self.conn
            .execute("DELETE FROM tokens WHERE id = ?1", [id])?;
        Ok(())
    }

    /// Store a session secret key name (value is NOT stored)
    pub fn store_session_secret(
        &self,
        session_id: &str,
        key_name: &str,
        source: &str,
    ) -> Result<()> {
        let now = Utc::now().to_rfc3339();
        self.conn.execute(
            "INSERT INTO session_secrets (session_id, key_name, source, loaded_at)
             VALUES (?1, ?2, ?3, ?4)",
            params![session_id, key_name, source, now],
        )?;
        Ok(())
    }

    /// Get all secret key names for a session
    pub fn get_session_secrets(&self, session_id: &str) -> Result<Vec<SessionSecret>> {
        let mut stmt = self.conn.prepare(
            "SELECT id, session_id, key_name, source, loaded_at
             FROM session_secrets
             WHERE session_id = ?1
             ORDER BY loaded_at",
        )?;

        let secrets = stmt
            .query_map(params![session_id], |row| {
                Ok(SessionSecret {
                    id: row.get(0)?,
                    session_id: SessionId::new(row.get(1)?),
                    key_name: row.get(2)?,
                    source: row.get(3)?,
                    loaded_at: row
                        .get::<_, String>(4)?
                        .parse()
                        .unwrap_or_else(|_| Utc::now()),
                })
            })?
            .collect::<rusqlite::Result<Vec<_>>>()?;

        Ok(secrets)
    }

    /// Clear session secrets and return the key names (for unset output)
    pub fn clear_session_secrets(&self, session_id: &str) -> Result<Vec<String>> {
        let keys = self.get_session_secrets(session_id)?;
        let key_names: Vec<String> = keys.into_iter().map(|s| s.key_name).collect();

        self.conn.execute(
            "DELETE FROM session_secrets WHERE session_id = ?1",
            params![session_id],
        )?;

        Ok(key_names)
    }

    /// Get a preference value by key
    pub fn get_preference(&self, key: &str) -> Result<Option<String>> {
        let val = self
            .conn
            .query_row(
                "SELECT value FROM preferences WHERE key = ?1",
                rusqlite::params![key],
                |row| row.get(0),
            )
            .optional()?;
        Ok(val)
    }

    /// Set a preference value
    pub fn set_preference(&self, key: &str, value: &str) -> Result<()> {
        self.conn.execute(
            "INSERT INTO preferences (key, value) VALUES (?1, ?2)
             ON CONFLICT(key) DO UPDATE SET value = excluded.value",
            rusqlite::params![key, value],
        )?;
        Ok(())
    }

    /// Get a boolean preference (defaults to false if missing)
    pub fn get_bool_preference(&self, key: &str) -> Result<bool> {
        Ok(self
            .get_preference(key)?
            .map(|v| v == "true")
            .unwrap_or(false))
    }

    /// Run VACUUM to reclaim unused space and defragment the database file
    pub fn vacuum(&self) -> Result<()> {
        self.conn.execute_batch("VACUUM")?;
        Ok(())
    }

    /// Delete the oldest commands beyond `max_entries`, keeping the most recent ones.
    /// Returns the number of deleted rows.
    pub fn prune_old_commands(&self, max_entries: usize) -> Result<usize> {
        let deleted = self.conn.execute(
            "DELETE FROM commands WHERE id NOT IN (
                SELECT id FROM commands ORDER BY timestamp DESC LIMIT ?1
            )",
            params![max_entries as i64],
        )?;
        Ok(deleted)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::NamedTempFile;

    #[test]
    fn test_database_creation() {
        let temp_file = NamedTempFile::new().unwrap();
        let db = Database::new(temp_file.path()).unwrap();
        let stats = db.get_stats().unwrap();
        assert_eq!(stats.total_commands, 0);
    }

    #[test]
    fn test_add_command() {
        let temp_file = NamedTempFile::new().unwrap();
        let mut db = Database::new(temp_file.path()).unwrap();

        let cmd_id = db
            .add_command("ls -la", "/home/user", Utc::now(), false, Some(0))
            .unwrap();
        assert!(cmd_id > 0);

        let stats = db.get_stats().unwrap();
        assert_eq!(stats.total_commands, 1);
    }

    #[test]
    fn test_token_storage() {
        let temp_file = NamedTempFile::new().unwrap();
        let mut db = Database::new(temp_file.path()).unwrap();

        let cmd_id = db
            .add_command("echo password123", "/home", Utc::now(), true, None)
            .unwrap();

        db.store_token(cmd_id, "password", "<redacted>", "password123")
            .unwrap();

        let tokens = db.get_tokens_for_command(CommandId::new(cmd_id)).unwrap();
        assert_eq!(tokens.len(), 1);
        assert_eq!(tokens[0].original_value, "password123");
    }

    #[test]
    fn test_alias_crud() {
        let temp_file = NamedTempFile::new().unwrap();
        let db = Database::new(temp_file.path()).unwrap();

        // Add
        db.add_alias("ll", "ls -la", "long listing").unwrap();
        let aliases = db.list_aliases().unwrap();
        assert_eq!(aliases.len(), 1);
        assert_eq!(aliases[0].alias, "ll");
        assert_eq!(aliases[0].command, "ls -la");
        assert_eq!(aliases[0].description, "long listing");

        // Update
        db.update_alias("ll", "ls -lah", Some("long listing with human sizes"))
            .unwrap();
        let aliases = db.list_aliases().unwrap();
        assert_eq!(aliases[0].command, "ls -lah");
        assert_eq!(aliases[0].description, "long listing with human sizes");

        // Remove
        db.remove_alias("ll").unwrap();
        let aliases = db.list_aliases().unwrap();
        assert!(aliases.is_empty());
    }

    #[test]
    fn test_alias_sync() {
        let temp_file = NamedTempFile::new().unwrap();
        let db = Database::new(temp_file.path()).unwrap();

        let aliases = vec![
            ("ll".to_string(), "ls -la".to_string()),
            ("gs".to_string(), "git status".to_string()),
        ];

        let count = db.sync_aliases(&aliases).unwrap();
        assert_eq!(count, 2);

        // Sync again with updated command — should upsert
        let updated = vec![("ll".to_string(), "ls -lah".to_string())];
        db.sync_aliases(&updated).unwrap();

        let all = db.list_aliases().unwrap();
        assert_eq!(all.len(), 2);
        let ll = all.iter().find(|a| a.alias == "ll").unwrap();
        assert_eq!(ll.command, "ls -lah");
    }
}