rpo 0.1.0-beta.4

use std::fs::File;
use std::path::{Path, PathBuf};

use polars::prelude::{DataFrame, ParquetWriter};

use super::FrameSink;
use crate::RpoError;

#[derive(Clone, Copy, Debug)]
/// What to do with an existing table.
pub enum DuckDbWriteMode {
    /// Drop and recreate it.
    Replace,
    /// Add rows to it.
    Append,
}

/// Streams blame snapshots into a DuckDB table.
///
/// Requires the `sink-duckdb` feature. DuckDB is statically bundled, so
/// no system library is needed.
pub struct DuckDbSink {
    db_path: PathBuf,
    table: String,
    mode: DuckDbWriteMode,
    conn: Option<::duckdb::Connection>,
    first_write_done: bool,
    snapshot_seq: u64,
}

impl DuckDbSink {
    /// Open (or create) `db_path` and write snapshots into `table`.
    ///
    /// # Errors
    ///
    /// [`RpoError::Sink`] if `table` is not a plain identifier matching
    /// `^[A-Za-z_][A-Za-z0-9_]*$` — table names cannot be parameterized
    /// in SQL, so they are validated rather than escaped.
    pub fn new(
        db_path: impl AsRef<Path>,
        table: &str,
        mode: DuckDbWriteMode,
    ) -> Result<Self, RpoError> {
        if !is_valid_table_name(table) {
            return Err(RpoError::Sink(format!("invalid table name: {table:?}")));
        }
        let conn = ::duckdb::Connection::open(db_path.as_ref())
            .map_err(|e| RpoError::Sink(e.to_string()))?;
        Ok(Self {
            db_path: db_path.as_ref().to_path_buf(),
            table: table.to_string(),
            mode,
            conn: Some(conn),
            first_write_done: false,
            snapshot_seq: 0,
        })
    }

    fn conn(&mut self) -> Result<&mut ::duckdb::Connection, RpoError> {
        self.conn
            .as_mut()
            .ok_or_else(|| RpoError::Sink("sink already finished".into()))
    }
}

impl FrameSink for DuckDbSink {
    fn write_snapshot(&mut self, mut frame: DataFrame) -> Result<(), RpoError> {
        // The duckdb 1.1 crate's `polars` feature only exposes
        // *result*-side polars conversion (Statement::query_polars). It does
        // not expose `register_polars_dataframe` / `unregister`, and even its
        // arrow-FFI table function (ArrowVTab) would require a different
        // arrow crate version than polars 0.53 ships internally. The
        // pragmatic, type-safe bridge is to spill each snapshot to a
        // temporary Parquet file and let DuckDB read it back via the
        // built-in `read_parquet` table function (parquet support is
        // statically linked via the `bundled` feature).
        let table = self.table.clone();
        let mode = self.mode;
        let first = !self.first_write_done;
        self.snapshot_seq += 1;
        let seq = self.snapshot_seq;

        let tmp_dir = std::env::temp_dir();
        let tmp_path = tmp_dir.join(format!(
            "rpo_duckdb_sink_{}_{}.parquet",
            std::process::id(),
            seq,
        ));

        struct TempFileGuard<'a>(&'a std::path::Path);
        impl Drop for TempFileGuard<'_> {
            fn drop(&mut self) {
                let _ = std::fs::remove_file(self.0);
            }
        }
        let _guard = TempFileGuard(&tmp_path);

        {
            let file = File::create(&tmp_path)?;
            ParquetWriter::new(file).finish(&mut frame)?;
        }

        let conn = self.conn()?;
        let parquet_literal = sql_string_literal(&tmp_path.to_string_lossy());

        let sql = match (mode, first) {
            (DuckDbWriteMode::Replace, true) => format!(
                "CREATE OR REPLACE TABLE {table} AS SELECT * FROM read_parquet({parquet_literal})"
            ),
            (DuckDbWriteMode::Append, true) => {
                // Create-if-not-exists with the right schema, then insert.
                conn.execute_batch(&format!(
                    "CREATE TABLE IF NOT EXISTS {table} AS SELECT * FROM read_parquet({parquet_literal}) WHERE 1=0"
                ))
                .map_err(|e| RpoError::Sink(e.to_string()))?;
                format!("INSERT INTO {table} SELECT * FROM read_parquet({parquet_literal})")
            }
            (_, false) => {
                format!("INSERT INTO {table} SELECT * FROM read_parquet({parquet_literal})")
            }
        };

        conn.execute_batch(&sql)
            .map_err(|e| RpoError::Sink(e.to_string()))?;

        self.first_write_done = true;
        Ok(())
    }

    fn finish(&mut self) -> Result<(), RpoError> {
        // Drop the connection on finish to flush WAL.
        self.conn.take();
        let _ = &self.db_path;
        Ok(())
    }
}

/// Quote a string for safe interpolation as a SQL string literal.
/// DuckDB uses single-quoted strings with `''` to escape an embedded quote.
fn sql_string_literal(s: &str) -> String {
    let mut out = String::with_capacity(s.len() + 2);
    out.push('\'');
    for ch in s.chars() {
        if ch == '\'' {
            out.push('\'');
        }
        out.push(ch);
    }
    out.push('\'');
    out
}

/// Validate a DuckDB table identifier against `^[A-Za-z_][A-Za-z0-9_]*$`.
/// Used to defend against SQL injection in interpolated DDL/DML.
fn is_valid_table_name(s: &str) -> bool {
    let mut chars = s.chars();
    let Some(first) = chars.next() else {
        return false;
    };
    if !(first.is_ascii_alphabetic() || first == '_') {
        return false;
    }
    chars.all(|c| c.is_ascii_alphanumeric() || c == '_')
}

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

    #[test]
    fn table_name_validation_accepts_valid_names() {
        let dir = tempdir().unwrap();
        let db_path = dir.path().join("t.duckdb");
        for name in ["foo", "foo_bar", "_x"] {
            let res = DuckDbSink::new(&db_path, name, DuckDbWriteMode::Replace);
            assert!(res.is_ok(), "expected {name:?} to be accepted");
        }
    }

    #[test]
    fn table_name_validation_rejects_invalid_names() {
        let dir = tempdir().unwrap();
        let db_path = dir.path().join("t.duckdb");
        for name in ["", "foo bar", "foo;DROP", "123foo", "foo'bar"] {
            let res = DuckDbSink::new(&db_path, name, DuckDbWriteMode::Replace);
            assert!(res.is_err(), "expected {name:?} to be rejected");
            match res.err().unwrap() {
                RpoError::Sink(msg) => {
                    assert!(
                        msg.contains("invalid table name"),
                        "unexpected error message: {msg}"
                    );
                }
                other => panic!("expected RpoError::Sink, got {other:?}"),
            }
        }
    }
}