flowcode-core 0.4.3-alpha

//! flowcode-core/src/duckdb_session.rs
//! Manages DuckDB connections and SQL query execution for FlowCode.

use std::cell::RefCell;
use std::rc::Rc;

use crate::ast::ArgValue;
use crate::types::{TypedValue, ValueKind};
use crate::error::FCError;

// Use DuckDB's re-exported Arrow types to ensure version compatibility
use duckdb::Connection;
use duckdb::types::Value;

/// A session for interacting with a DuckDB database.
#[cfg(feature = "duckdb")]
pub struct DuckSession {
    conn: Rc<RefCell<Connection>>,
}

#[cfg(feature = "duckdb")]
impl DuckSession {
    /// Creates a new DuckDB session with an in-memory database for testing.
    pub fn new() -> Result<Self, FCError> {
        // Use in-memory database for testing to avoid file locking issues
        let conn = Connection::open_in_memory()
            .map_err(|e| FCError::BackendError(format!("Connection error: {}", e.to_string())))?;
        
        Ok(DuckSession {
            conn: Rc::new(RefCell::new(conn)),
        })
    }

    /// Executes an SQL query and returns the result as a TypedValue.
    pub fn execute_sql(&self, query: &str) -> Result<TypedValue, FCError> {
        println!("Debug - Full SQL Query to Execute: {}", query);
        self.execute_with_params(query, &[])
    }
    
    /// Executes a query with parameters and returns the result as a TypedValue.
    pub fn execute_with_params(
        &self,
        query: &str,
        params: &[ArgValue],
    ) -> Result<TypedValue, FCError> {
        use crate::duckdb_interchange::record_batch_to_typed_value;
        
        let conn = self.conn.borrow();
        let mut stmt = conn.prepare(query)
            .map_err(|e| FCError::BackendError(format!("Query error: {}", e.to_string())))?;

        // Convert parameters to DuckDB-compatible values
        let duck_params: Vec<Value> = params
            .iter()
            .map(convert_arg_to_duckdb_value)
            .collect::<Result<Vec<_>, _>>()?;

        // Create parameter references
        let params: Vec<&dyn duckdb::ToSql> = duck_params
            .iter()
            .map(|p| p as &dyn duckdb::ToSql)
            .collect();

        // Execute query and get Arrow results iterator
        let mut arrow_results = stmt.query_arrow(&params[..])
            .map_err(|e| FCError::BackendError(format!("Query error: {}", e.to_string())))?;
        
        // Collect batches from Arrow results
        let mut batches: Vec<duckdb::arrow::record_batch::RecordBatch> = Vec::new();
        while let Some(batch) = arrow_results.next() {
            batches.push(batch);
        }
        
        // If no batches, return empty table
        if batches.is_empty() {
            return Ok(TypedValue::new(
                ValueKind::Table,
                ArgValue::Table(Vec::new())
            ));
        }
        
        // For simplicity, take the first batch or concatenate in future
        let combined_batch = batches[0].clone();
        
        // Convert RecordBatch to TypedValue using interchange layer
        record_batch_to_typed_value(combined_batch)
            .map_err(|e| FCError::BackendError(format!("Table conversion error: {}", e.to_string())))
    }
    
    /// Registers a CSV file as a table in DuckDB for querying.
    pub fn register_csv(&self, table_name: &str, file_path: &str) -> Result<(), FCError> {
        let create_view_sql = format!(
            "CREATE OR REPLACE VIEW {} AS SELECT * FROM read_csv_auto('{}')",
            table_name, file_path
        );
        let conn = self.conn.borrow();
        conn.execute(&create_view_sql, [])
            .map_err(|e| FCError::BackendError(format!("Failed to register CSV as table: {}", e.to_string())))?;
        println!("Debug - Registered CSV: {} as table: {}", file_path, table_name);
        Ok(())
    }
    
    /// Executes a query that returns a single value.
    pub fn query_scalar(&self, query: &str) -> Result<TypedValue, FCError> {
        let conn = self.conn.borrow();
        let mut stmt = conn.prepare(query)
            .map_err(|e| FCError::BackendError(format!("Query error: {}", e.to_string())))?;
        
        // Use query_row to get a single row
        let mut rows = stmt.query([])
            .map_err(|e| FCError::BackendError(format!("Query error: {}", e.to_string())))?;
        
        if let Ok(Some(row)) = rows.next() {
            let value: Value = row.get(0)
                .map_err(|e| FCError::BackendError(format!("Fetch error: {}", e.to_string())))?;
            Ok(duckdb_value_to_typed_value(&value))
        } else {
            // No result, return NULL
            Ok(TypedValue::new(ValueKind::Null, ArgValue::Null))
        }
    }
    
    /// Executes a query on a TypedValue Table and returns the result as a new Table.
    pub fn query_table(
        &self,
        _table: &TypedValue,
        query: &str,
    ) -> Result<TypedValue, FCError> {
        // Placeholder: Currently just executes the query without using the input table
        // In a real implementation, we would bind the table data to a temporary view
        self.execute_sql(query)
    }
    
    /// Convenience wrapper: LOAD CSV via DuckDB's read_csv_auto into a table.
    pub fn load_csv_auto(&self, table_name: &str, file_path: &str) -> Result<(), FCError> {
        let create_table_sql = format!(
            "CREATE OR REPLACE TABLE {} AS SELECT * FROM read_csv_auto('{}', HEADER=TRUE)",
            table_name, file_path
        );
        let conn = self.conn.borrow();
        conn.execute(&create_table_sql, [])
            .map_err(|e| FCError::BackendError(format!("Failed to load CSV into table: {}", e.to_string())))?;
        println!("Debug - Loaded CSV: {} into table: {}", file_path, table_name);
        Ok(())
    }
    
    /// Fetches column names for a given table using PRAGMA table_info.
    pub fn get_table_columns(&self, table_name: &str) -> Result<Vec<String>, FCError> {
        let pragma_sql = format!("PRAGMA table_info('{}')", table_name);
        let conn = self.conn.borrow();
        let mut stmt = conn.prepare(&pragma_sql)
            .map_err(|e| FCError::BackendError(format!("Query error: {}", e.to_string())))?;
        let mut rows = stmt.query([])
            .map_err(|e| FCError::BackendError(format!("Query error: {}", e.to_string())))?;
        let mut columns = Vec::new();
        while let Ok(Some(row)) = rows.next() {
            let name: String = row.get(1)
                .map_err(|e| FCError::BackendError(format!("Fetch error: {}", e.to_string())))?;
            columns.push(name);
        }
        Ok(columns)
    }

    fn run_query(&self, sql: &str) -> Result<(), FCError> {
        let result = self.conn.borrow().execute(sql, [])
            .map_err(|e| FCError::BackendError(format!("Backend error: {}", e.to_string())))?;
        println!("Result: DuckDB executed: {} with query: {} (affected rows: {})", sql.split_whitespace().next().unwrap_or("unknown"), sql, result);
        Ok(())
    }
}

#[cfg(feature = "duckdb")]
impl crate::executor::backend::Backend for DuckSession {
    fn execute(&self, command: &crate::parser::ParsedCommand) -> Result<String, Box<dyn std::error::Error>> {
        use crate::duckdb_verbs::map_verb_to_sql;

        // Convert command arguments to TypedValue for mapping
        let typed_args: Vec<TypedValue> = command.args.iter().map(|arg| {
            match arg {
                crate::ast::ArgValue::Number(n) => TypedValue::new(crate::types::ValueKind::Number, crate::ast::ArgValue::Number(*n)),
                crate::ast::ArgValue::String(s) => TypedValue::new(crate::types::ValueKind::String, crate::ast::ArgValue::String(s.clone())),
                crate::ast::ArgValue::Bool(b) => TypedValue::new(crate::types::ValueKind::Bool, crate::ast::ArgValue::Bool(*b)),
                crate::ast::ArgValue::Null => TypedValue::new(crate::types::ValueKind::Null, crate::ast::ArgValue::Null),
                crate::ast::ArgValue::Table(tbl) => TypedValue::new(crate::types::ValueKind::Table, crate::ast::ArgValue::Table(tbl.clone())),
            }
        }).collect();

        // Map the command name and arguments to an SQL query, passing self as the session
        let sql_query = map_verb_to_sql(&command.name, &typed_args, self)
            .map_err(|e| Box::new(std::io::Error::new(std::io::ErrorKind::Other, e.to_string())) as Box<dyn std::error::Error>)?;

        // Debug: Print the full SQL query before execution
        println!("Debug - Full SQL Query to Execute: {}", sql_query);

        // Execute the SQL query using DuckDB with mutable access
        self.run_query(&sql_query)
            .map_err(|e| Box::new(std::io::Error::new(std::io::ErrorKind::Other, e.to_string())) as Box<dyn std::error::Error>)?;

        // For now, return a simple success message with the result count
        // In a full implementation, we would convert the query result to a string representation
        Ok(format!("DuckDB executed: {} with query: {} (affected rows: 0)", command.name, sql_query))
    }
}

/// Helper function to convert ArgValue to DuckDB Value
fn convert_arg_to_duckdb_value(arg: &ArgValue) -> Result<Value, FCError> {
    match arg {
        ArgValue::Number(n) => Ok(Value::Double(*n)),
        ArgValue::String(s) => Ok(Value::Text(s.clone())),
        ArgValue::Bool(b) => Ok(Value::Boolean(*b)),
        ArgValue::Null => Ok(Value::Null),
        _ => Err(FCError::BackendError(
            "Unsupported parameter type".to_string()
        )),
    }
}

/// Helper function to convert DuckDB Value to TypedValue
fn duckdb_value_to_typed_value(value: &Value) -> TypedValue {
    match value {
        Value::Null => TypedValue::new(ValueKind::Null, ArgValue::Null),
        Value::Boolean(b) => TypedValue::new(ValueKind::Bool, ArgValue::Bool(*b)),
        Value::TinyInt(i) => TypedValue::new(ValueKind::Number, ArgValue::Number(*i as f64)),
        Value::SmallInt(i) => TypedValue::new(ValueKind::Number, ArgValue::Number(*i as f64)),
        Value::Int(i) => TypedValue::new(ValueKind::Number, ArgValue::Number(*i as f64)),
        Value::BigInt(i) => TypedValue::new(ValueKind::Number, ArgValue::Number(*i as f64)),
        Value::HugeInt(i) => TypedValue::new(ValueKind::Number, ArgValue::Number(*i as f64)),
        Value::Float(f) => TypedValue::new(ValueKind::Number, ArgValue::Number(*f as f64)),
        Value::Double(d) => TypedValue::new(ValueKind::Number, ArgValue::Number(*d)),
        Value::Text(s) => {
            let s = s.clone();
            TypedValue::new(ValueKind::String, ArgValue::String(s))
        },
        _ => TypedValue::new(ValueKind::Null, ArgValue::Null),
    }
}

#[cfg(feature = "duckdb")]
unsafe impl Send for DuckSession {}
#[cfg(feature = "duckdb")]
unsafe impl Sync for DuckSession {}