dbx/crates/dbx-core/src/db/sqlite.rs

use futures::StreamExt;
use sqlx::sqlite::{SqliteConnectOptions, SqlitePool, SqlitePoolOptions, SqliteRow};
use sqlx::{Column, Executor, Row};
use std::str::FromStr;
use std::time::{Duration, Instant};

use super::file_validator::validate_file_path;
use crate::sql::starts_with_executable_sql_keyword;
use crate::types::{ColumnInfo, DatabaseInfo, ForeignKeyInfo, IndexInfo, QueryResult, TableInfo, TriggerInfo};

pub async fn connect_path(path: &str) -> Result<SqlitePool, String> {
    let is_memory = is_memory_database_path(path);
    if !is_memory {
        validate_file_path(path, is_network_path)?;
    }

    let mut options = if is_memory {
        SqliteConnectOptions::from_str("sqlite::memory:").map_err(|e| format!("SQLite connection failed: {e}"))?
    } else {
        SqliteConnectOptions::new().filename(path).create_if_missing(false)
    };

    if is_network_path(path) {
        options = options.vfs("unix-nolock");
    }

    SqlitePoolOptions::new()
        .max_connections(if is_memory { 1 } else { 5 })
        .acquire_timeout(Duration::from_secs(10))
        .idle_timeout(Duration::from_secs(300))
        .connect_with(options)
        .await
        .map_err(|e| format!("SQLite connection failed: {e}"))
}

fn is_network_path(path: &str) -> bool {
    path.starts_with("\\\\") || path.starts_with("//") || path.contains("wsl.localhost") || path.contains("wsl$")
}

pub fn is_memory_database_path(path: &str) -> bool {
    path.trim().eq_ignore_ascii_case(":memory:")
}

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

    #[tokio::test]
    async fn connect_path_supports_memory_database_across_statements() {
        let pool = connect_path(":memory:").await.expect("connect in-memory SQLite");

        execute_query(&pool, "CREATE TABLE memory_probe (id INTEGER PRIMARY KEY, name TEXT);")
            .await
            .expect("create table");
        execute_query(&pool, "INSERT INTO memory_probe (name) VALUES ('Ada');").await.expect("insert row");
        let result = execute_query(&pool, "SELECT name FROM memory_probe WHERE id = 1;").await.expect("select row");

        assert_eq!(result.rows[0][0], serde_json::json!("Ada"));
    }
}

pub async fn list_databases(_pool: &SqlitePool) -> Result<Vec<DatabaseInfo>, String> {
    Ok(vec![DatabaseInfo { name: "main".to_string() }])
}

pub async fn list_tables(pool: &SqlitePool, _schema: &str) -> Result<Vec<TableInfo>, String> {
    let rows: Vec<SqliteRow> = sqlx::query(
        "SELECT name, type FROM sqlite_master WHERE type IN ('table', 'view') AND name NOT LIKE 'sqlite_%' ORDER BY name",
    )
    .fetch_all(pool)
    .await
    .map_err(|e| e.to_string())?;

    Ok(rows
        .iter()
        .map(|row| {
            let t: String = row.get("type");
            TableInfo {
                name: row.get::<String, _>("name"),
                table_type: if t == "view" { "VIEW".to_string() } else { "BASE TABLE".to_string() },
                comment: None,
            }
        })
        .collect())
}

pub async fn get_columns(pool: &SqlitePool, _schema: &str, table: &str) -> Result<Vec<ColumnInfo>, String> {
    let rows: Vec<SqliteRow> =
        sqlx::query(&format!("PRAGMA table_info(\"{}\")", table)).fetch_all(pool).await.map_err(|e| e.to_string())?;

    Ok(rows
        .iter()
        .map(|row| ColumnInfo {
            name: row.get::<String, _>("name"),
            data_type: row.get::<String, _>("type"),
            is_nullable: row.get::<i32, _>("notnull") == 0,
            column_default: row.get::<Option<String>, _>("dflt_value"),
            is_primary_key: row.get::<i32, _>("pk") > 0,
            extra: None,
            comment: None,
            numeric_precision: None,
            numeric_scale: None,
            character_maximum_length: None,
        })
        .collect())
}

pub async fn list_indexes(pool: &SqlitePool, _schema: &str, table: &str) -> Result<Vec<IndexInfo>, String> {
    let safe_table = table.replace('"', "\"\"");
    let idx_rows: Vec<SqliteRow> = sqlx::query(&format!("PRAGMA index_list(\"{safe_table}\")"))
        .fetch_all(pool)
        .await
        .map_err(|e| e.to_string())?;

    let mut indexes = Vec::new();
    for idx_row in &idx_rows {
        let name: String = idx_row.get("name");
        let is_unique: bool = idx_row.get::<i32, _>("unique") != 0;
        let origin: String = idx_row.get::<String, _>("origin");
        let is_primary = origin == "pk";

        let safe_name = name.replace('"', "\"\"");
        let col_rows: Vec<SqliteRow> = sqlx::query(&format!("PRAGMA index_info(\"{safe_name}\")"))
            .fetch_all(pool)
            .await
            .map_err(|e| e.to_string())?;

        let columns: Vec<String> = col_rows.iter().map(|r| r.get::<String, _>("name")).collect();

        indexes.push(IndexInfo {
            name,
            columns,
            is_unique,
            is_primary,
            filter: None,
            index_type: None,
            included_columns: None,
            comment: None,
        });
    }
    Ok(indexes)
}

pub async fn list_foreign_keys(pool: &SqlitePool, _schema: &str, table: &str) -> Result<Vec<ForeignKeyInfo>, String> {
    let rows: Vec<SqliteRow> = sqlx::query(&format!("PRAGMA foreign_key_list(\"{}\")", table))
        .fetch_all(pool)
        .await
        .map_err(|e| e.to_string())?;

    Ok(rows
        .iter()
        .map(|row| ForeignKeyInfo {
            name: format!("fk_{}", row.get::<i32, _>("id")),
            column: row.get::<String, _>("from"),
            ref_table: row.get::<String, _>("table"),
            ref_column: row.get::<String, _>("to"),
        })
        .collect())
}

pub async fn list_triggers(pool: &SqlitePool, _schema: &str, table: &str) -> Result<Vec<TriggerInfo>, String> {
    let rows: Vec<SqliteRow> =
        sqlx::query("SELECT name, sql FROM sqlite_master WHERE type = 'trigger' AND tbl_name = ? ORDER BY name")
            .bind(table)
            .fetch_all(pool)
            .await
            .map_err(|e| e.to_string())?;

    Ok(rows
        .iter()
        .map(|row| {
            let sql_text: String = row.get::<Option<String>, _>("sql").unwrap_or_default();
            let upper = sql_text.to_uppercase();
            let timing = if upper.contains("BEFORE") {
                "BEFORE"
            } else if upper.contains("AFTER") {
                "AFTER"
            } else {
                "INSTEAD OF"
            };
            let event = if upper.contains("INSERT") {
                "INSERT"
            } else if upper.contains("UPDATE") {
                "UPDATE"
            } else {
                "DELETE"
            };
            TriggerInfo { name: row.get::<String, _>("name"), event: event.to_string(), timing: timing.to_string() }
        })
        .collect())
}

pub async fn execute_query(pool: &SqlitePool, sql: &str) -> Result<QueryResult, String> {
    execute_query_with_max_rows(pool, sql, None).await
}

fn query_result_row_limit(max_rows: Option<usize>) -> usize {
    max_rows.unwrap_or(crate::query::MAX_ROWS).max(1)
}

pub async fn execute_query_with_max_rows(
    pool: &SqlitePool,
    sql: &str,
    max_rows: Option<usize>,
) -> Result<QueryResult, String> {
    let start = Instant::now();
    let row_limit = query_result_row_limit(max_rows);

    if starts_with_executable_sql_keyword(sql, &["SELECT", "PRAGMA", "EXPLAIN", "WITH"]) {
        let desc = pool.describe(sql).await.map_err(|e| e.to_string())?;
        let columns: Vec<String> = desc.columns().iter().map(|c| c.name().to_string()).collect();

        let mut stream = sqlx::query(sql).fetch(pool);
        let mut result_rows: Vec<Vec<serde_json::Value>> = Vec::new();

        while let Some(row) = stream.next().await {
            let row = row.map_err(|e| e.to_string())?;
            result_rows.push(
                (0..row.len())
                    .map(|i| {
                        row.try_get::<String, _>(i)
                            .map(serde_json::Value::String)
                            .or_else(|_| row.try_get::<i64, _>(i).map(super::safe_i64_to_json))
                            .or_else(|_| {
                                row.try_get::<f64, _>(i).map(|v| {
                                    serde_json::Number::from_f64(v)
                                        .map(serde_json::Value::Number)
                                        .unwrap_or(serde_json::Value::Null)
                                })
                            })
                            .or_else(|_| row.try_get::<bool, _>(i).map(serde_json::Value::Bool))
                            .unwrap_or(serde_json::Value::Null)
                    })
                    .collect(),
            );
            if result_rows.len() > row_limit {
                break;
            }
        }

        let truncated = result_rows.len() > row_limit;
        if truncated {
            result_rows.truncate(row_limit);
        }

        Ok(QueryResult {
            columns,
            rows: result_rows,
            affected_rows: 0,
            execution_time_ms: start.elapsed().as_millis(),
            truncated,
            session_id: None,
            has_more: false,
        })
    } else {
        let result = sqlx::query(sql).execute(pool).await.map_err(|e| e.to_string())?;

        Ok(QueryResult {
            columns: vec![],
            rows: vec![],
            affected_rows: result.rows_affected(),
            execution_time_ms: start.elapsed().as_millis(),
            truncated: false,
            session_id: None,
            has_more: false,
        })
    }
}