postgres-mumu 0.1.2

// src/postgres.rs

use core_mumu::parser::interpreter::Interpreter;
use core_mumu::parser::types::{
    FunctionValue, IteratorHandle, IteratorKind, PluginIterator, Value,
};
use futures::StreamExt;
use indexmap::IndexMap;
use once_cell::sync::Lazy;
use regex::Regex;
use std::collections::HashMap;
use std::sync::{mpsc, Arc, Mutex};
use tokio::runtime::Runtime;
use tokio_postgres::{types::ToSql, Client, NoTls, Row};

static PG_CONN_MAP: Lazy<Mutex<HashMap<i64, Arc<Client>>>> =
    Lazy::new(|| Mutex::new(HashMap::new()));
static PG_HANDLE_SEQ: Lazy<Mutex<i64>> = Lazy::new(|| Mutex::new(1));

/* ───────────────────────────── connection ───────────────────────────── */

pub fn connect_bridge(
    _interp: &mut Interpreter,
    mut args: Vec<Value>,
) -> Result<Value, String> {
    if args.len() != 1 {
        return Err(format!(
            "pg:connect(url) expects 1 argument, got {}",
            args.len()
        ));
    }
    let url = match args.remove(0) {
        Value::SingleString(s) => s,
        other => return Err(format!("pg:connect(url) expects a string, got {:?}", other)),
    };

    // Create a dedicated runtime to establish the connection, then keep that
    // runtime alive on a background thread that drives the connection task.
    let rt = Runtime::new().map_err(|e| format!("tokio runtime error: {e}"))?;
    let (client, connection) = rt
        .block_on(tokio_postgres::connect(&url, NoTls))
        .map_err(|e| format!("postgres connect error: {e}"))?;

    // Drive the connection on a background thread.
    std::thread::spawn(move || {
        let _ = rt.block_on(connection);
    });

    let mut seq = PG_HANDLE_SEQ.lock().unwrap();
    let handle = *seq;
    *seq += 1;

    PG_CONN_MAP
        .lock()
        .unwrap()
        .insert(handle, Arc::new(client));
    Ok(Value::Long(handle))
}

/* ─────────────────── placeholders / partial application ─────────────────── */

fn is_placeholder(v: &Value) -> bool {
    match v {
        Value::Placeholder => true,
        Value::SingleString(s) if s == "_" => true,
        Value::StrArray(ss) if ss.len() == 1 && ss[0] == "_" => true,
        _ => false,
    }
}

fn query_partial_closure(store_args: Vec<Value>) -> FunctionValue {
    use FunctionValue::RustClosure;
    RustClosure(
        "pg:query-partial".to_string(),
        Arc::new(Mutex::new(
            move |interp: &mut Interpreter, new_args: Vec<Value>| {
                let mut combined = store_args.clone();
                combined.extend(new_args);
                query_bridge_partial(interp, vec![], combined)
            },
        )),
        0, // variadic-style
    )
}

pub fn query_bridge(interp: &mut Interpreter, args: Vec<Value>) -> Result<Value, String> {
    query_bridge_partial(interp, args, vec![])
}

#[derive(Debug)]
enum PlaceholderKind {
    Positional(Vec<usize>), // $1, $2, ...
    Named(Vec<String>),     // $foo, $bar (in order of first appearance)
    None,
}

fn scan_sql_placeholders(sql: &str) -> PlaceholderKind {
    let re = Regex::new(r"\$([a-zA-Z_][a-zA-Z0-9_]*|\d+)").unwrap();
    let mut named = Vec::new();
    let mut positional = Vec::new();
    for cap in re.captures_iter(sql) {
        let name = &cap[1];
        if let Ok(n) = name.parse::<usize>() {
            positional.push(n);
        } else if !named.iter().any(|s: &String| s == name) {
            named.push(name.to_string());
        }
    }
    if !named.is_empty() && !positional.is_empty() {
        PlaceholderKind::None // mixed not supported
    } else if !named.is_empty() {
        PlaceholderKind::Named(named)
    } else if !positional.is_empty() {
        PlaceholderKind::Positional(positional)
    } else {
        PlaceholderKind::None
    }
}

fn rewrite_named_sql_to_positional(sql: &str, names: &[String]) -> (String, Vec<String>) {
    let mut order = Vec::with_capacity(names.len());
    order.extend_from_slice(names);
    let map: IndexMap<_, _> = order
        .iter()
        .enumerate()
        .map(|(i, k)| (k.clone(), format!("${}", i + 1)))
        .collect();
    let re = Regex::new(r"\$([a-zA-Z_][a-zA-Z0-9_]*)").unwrap();
    let result_sql = re.replace_all(sql, |caps: &regex::Captures| {
        let key = &caps[1];
        map.get(key).cloned().unwrap_or_else(|| format!("${}", key))
    });
    (result_sql.to_string(), order.clone())
}

/* ─────────────────────── query core (partial-aware) ─────────────────────── */

fn query_bridge_partial(
    _interp: &mut Interpreter,
    mut args: Vec<Value>,
    prior_args: Vec<Value>,
) -> Result<Value, String> {
    let mut all_args = prior_args;
    all_args.append(&mut args);

    // minimal arity: conn_handle, sql
    if all_args.len() < 2 {
        return Ok(Value::Function(Box::new(query_partial_closure(
            all_args.clone(),
        ))));
    }

    let handle = match &all_args[0] {
        Value::Long(h) => *h,
        other => {
            return Err(format!(
                "pg:query: expected connection handle (Long), got {:?}",
                other
            ))
        }
    };
    let sql = match &all_args[1] {
        Value::SingleString(s) => s.clone(),
        other => return Err(format!("pg:query(sql) expects a string, got {:?}", other)),
    };

    let placeholders = scan_sql_placeholders(&sql);

    // allow partial-application holes across the first 2 or 3 args
    let mut actual_args = all_args.clone();
    if actual_args.len() < 3 {
        actual_args.push(Value::Placeholder);
    }
    let any_placeholder =
        actual_args[0..(if matches!(placeholders, PlaceholderKind::None) { 2 } else { 3 })]
            .iter()
            .any(is_placeholder);
    if any_placeholder {
        return Ok(Value::Function(Box::new(query_partial_closure(
            actual_args.clone(),
        ))));
    }

    // Build (maybe rewritten) sql + params as **Value**s (convert to ToSql in worker)
    let (final_sql, params_vals): (String, Option<Vec<Value>>) = match placeholders {
        PlaceholderKind::None => (sql.clone(), None),

        PlaceholderKind::Named(ref names) => {
            let param_arg = all_args.get(2);
            match param_arg {
                Some(Value::KeyedArray(map)) => {
                    let (rewritten_sql, order) = rewrite_named_sql_to_positional(&sql, names);
                    let mut p = Vec::with_capacity(order.len());
                    for k in order.iter() {
                        if let Some(v) = map.get(k) {
                            p.push(v.clone());
                        } else if let Some(v) = map.get(&format!("${}", k)) {
                            p.push(v.clone());
                        } else {
                            return Err(format!(
                                "pg:query: SQL expects named parameter '${}', but it was not supplied in KeyedArray",
                                k
                            ));
                        }
                    }
                    (rewritten_sql, Some(p))
                }
                Some(other) => {
                    return Err(format!(
                        "pg:query: SQL contains named placeholders; 3rd arg must be a KeyedArray (got {:?}).",
                        other
                    ));
                }
                None => {
                    return Err(
                        "pg:query: SQL contains named placeholders, but params were not supplied."
                            .into(),
                    );
                }
            }
        }

        PlaceholderKind::Positional(ref nums) => {
            let max_pos = nums.iter().copied().max().unwrap_or(0);
            let param_arg = all_args.get(2);
            match param_arg {
                Some(Value::IntArray(arr)) => {
                    if arr.len() < max_pos {
                        return Err(format!(
                            "pg:query: expects at least {} positional params, got {} (IntArray)",
                            max_pos, arr.len()
                        ));
                    }
                    (sql.clone(), Some(arr.iter().map(|i| Value::Int(*i)).collect()))
                }
                Some(Value::FloatArray(arr)) => {
                    if arr.len() < max_pos {
                        return Err(format!(
                            "pg:query: expects at least {} positional params, got {} (FloatArray)",
                            max_pos, arr.len()
                        ));
                    }
                    (sql.clone(), Some(arr.iter().map(|f| Value::Float(*f)).collect()))
                }
                Some(Value::StrArray(arr)) => {
                    if arr.len() < max_pos {
                        return Err(format!(
                            "pg:query: expects at least {} positional params, got {} (StrArray)",
                            max_pos, arr.len()
                        ));
                    }
                    (
                        sql.clone(),
                        Some(arr.iter().map(|s| Value::SingleString(s.clone())).collect()),
                    )
                }
                Some(Value::MixedArray(arr)) => {
                    if arr.len() < max_pos {
                        return Err(format!(
                            "pg:query: expects at least {} positional params, got {} (MixedArray)",
                            max_pos, arr.len()
                        ));
                    }
                    (sql.clone(), Some(arr.clone()))
                }
                Some(other) => {
                    return Err(format!(
                        "pg:query: positional placeholders → 3rd arg must be an array (got {:?}).",
                        other
                    ));
                }
                None => {
                    return Err(
                        "pg:query: positional placeholders, but params not supplied.".into(),
                    );
                }
            }
        }
    };

    // Get the client Arc
    let client = {
        let map = PG_CONN_MAP.lock().unwrap();
        map.get(&handle)
            .cloned()
            .ok_or_else(|| "pg:query: invalid connection handle".to_string())?
    };

    // Build an iterator backed by a **background worker thread** that streams rows
    // and pushes them into a channel. The iterator blocks per-item (like before).
    let (tx, rx) = mpsc::channel::<Result<Value, String>>();
    let sql_for_thread = final_sql.clone();
    let params_vals_for_thread = params_vals.clone();

    std::thread::spawn(move || {
        let rt = match Runtime::new() {
            Ok(rt) => rt,
            Err(e) => {
                let _ = tx.send(Err(format!("tokio runtime error: {e}")));
                return;
            }
        };

        // Build ToSql boxes **inside** the worker (avoids Send bounds issues)
        let mut boxed_params: Vec<Box<dyn ToSql + Sync>> = Vec::new();
        let mut byref_params: Vec<&(dyn ToSql + Sync)> = Vec::new();
        if let Some(vs) = params_vals_for_thread {
            for v in vs {
                match mu_value_to_pg_param_box(&v) {
                    Ok(b) => boxed_params.push(b),
                    Err(e) => {
                        let _ = tx.send(Err(e));
                        return;
                    }
                }
            }
            for b in &boxed_params {
                byref_params.push(&**b);
            }
        }

        // Create raw stream
        let stream_res = if byref_params.is_empty() {
            rt.block_on(client.query_raw(&sql_for_thread, std::iter::empty::<&(dyn ToSql + Sync)>()))
        } else {
            rt.block_on(client.query_raw(&sql_for_thread, byref_params))
        };

        let mut stream = match stream_res {
            Ok(s) => Box::pin(s), // Pin<Box<...>>
            Err(e) => {
                let _ = tx.send(Err(format!("postgres query error: {e}")));
                return;
            }
        };

        // Pump rows
        loop {
            match rt.block_on(stream.as_mut().next()) {
                Some(Ok(row)) => {
                    let v = row_to_value(&row);
                    if tx.send(Ok(v)).is_err() {
                        // receiver gone
                        break;
                    }
                }
                Some(Err(e)) => {
                    let _ = tx.send(Err(format!("pg:query: row error: {e}")));
                    break;
                }
                None => {
                    // EOS: close channel (receiver will see NO_MORE_DATA)
                    break;
                }
            }
        }
        // drop(tx) ends stream naturally
    });

    // Wrap receiver in a PluginIterator
    let iter: Arc<Mutex<dyn PluginIterator>> = Arc::new(Mutex::new(PgRowIter {
        rx: Mutex::new(rx),
    }));
    Ok(Value::Iterator(IteratorHandle {
        kind: IteratorKind::Plugin(iter),
    }))
}

/* ─────────────────────────── value conversions ─────────────────────────── */

fn mu_value_to_pg_param_box(val: &Value) -> Result<Box<dyn ToSql + Sync>, String> {
    match val {
        Value::Int(i) => Ok(Box::new(*i)),
        Value::Float(f) => Ok(Box::new(*f)),
        Value::Long(l) => Ok(Box::new(*l)),
        Value::Bool(b) => Ok(Box::new(*b)),
        Value::SingleString(s) => Ok(Box::new(s.clone())),
        Value::StrArray(arr) => Ok(Box::new(arr.clone())),
        Value::IntArray(arr) => Ok(Box::new(arr.clone())),
        Value::FloatArray(arr) => Ok(Box::new(arr.clone())),
        Value::BoolArray(arr) => Ok(Box::new(arr.clone())),
        Value::Placeholder => Err("Cannot bind placeholder to query param".to_string()),
        Value::KeyedArray(_) => Err("Cannot bind KeyedArray directly as a single param".to_string()),
        Value::MixedArray(_) => Err("Cannot bind MixedArray as a *single* param (use array)".to_string()),
        Value::Function(_) => Err("Cannot bind function as query param".to_string()),
        Value::Tensor(_) => Err("Cannot bind tensor as query param".to_string()),
        other => Err(format!("Cannot bind value {:?} as query param", other)),
    }
}

fn row_to_value(row: &Row) -> Value {
    let mut map = IndexMap::new();
    for col in row.columns() {
        let name_str: &str = col.name();
        let key = name_str.to_string();

        // Try common types; fall through to String/JSON as safer defaults.
        let v: Value = if let Ok(val) = row.try_get::<_, i32>(name_str) {
            Value::Int(val)
        } else if let Ok(val) = row.try_get::<_, i64>(name_str) {
            Value::Long(val)
        } else if let Ok(val) = row.try_get::<_, f64>(name_str) {
            Value::Float(val)
        } else if let Ok(val) = row.try_get::<_, bool>(name_str) {
            Value::Bool(val)
        } else if let Ok(val) = row.try_get::<_, String>(name_str) {
            Value::SingleString(val)
        } else if let Ok(val) = row.try_get::<_, Option<String>>(name_str) {
            match val {
                Some(s) => Value::SingleString(s),
                None => Value::Undefined,
            }
        } else if let Ok(val) = row.try_get::<_, serde_json::Value>(name_str) {
            Value::SingleString(val.to_string())
        } else {
            Value::Placeholder
        };

        map.insert(key, v);
    }
    Value::KeyedArray(map)
}

/* ───────────────────────── plugin iterator impl ───────────────────────── */

struct PgRowIter {
    // Wrap Receiver in a Mutex so the iterator type is Sync, satisfying
    // PluginIterator: Send + Sync.
    rx: Mutex<mpsc::Receiver<Result<Value, String>>>,
}

impl PluginIterator for PgRowIter {
    fn next_value(&mut self) -> Result<Value, String> {
        let rx = self
            .rx
            .lock()
            .map_err(|_| "Iterator internal lock poisoned".to_string())?;
        match rx.recv() {
            Ok(Ok(v)) => Ok(v),
            Ok(Err(e)) => Err(e),
            Err(_disconnected) => Err("NO_MORE_DATA".into()),
        }
    }
}