graphar-flight 0.1.2

//! FalkorDB GRAPH.QUERY response (Redis RESP) → Arrow RecordBatch.
//!
//! Response shape (non-compact mode):
//!   Array[
//!     Array[col_name, ...]          ← headers  (index 0)
//!     Array[Array[val, ...], ...]   ← data rows (index 1)
//!     Array[stat_string, ...]       ← stats     (index 2, ignored)
//!   ]

use std::sync::Arc;

use arrow_array::{
    ArrayRef, RecordBatch,
    builder::{BooleanBuilder, Float64Builder, Int64Builder, StringBuilder},
};
use arrow_schema::{DataType, Field, SchemaRef};
use redis::Value;

use crate::error::{FlightSqlError, Result};

/// Convert a raw `GRAPH.QUERY` Redis response to a `RecordBatch` using a
/// **pre-registered** schema to guide type coercion.
///
/// Schema fields must be listed in the same order as the `RETURN` clause.
pub fn response_to_batch(response: Value, schema: &SchemaRef) -> Result<RecordBatch> {
    let mut outer = match response {
        Value::Array(v) => v,
        other => {
            return Err(FlightSqlError::MalformedResponse(format!(
                "expected Array at top level, got {other:?}"
            )));
        }
    };

    if outer.len() < 2 {
        return Ok(RecordBatch::new_empty(schema.clone()));
    }

    let rows_val = outer.remove(1);
    let rows = match rows_val {
        Value::Array(r) => r,
        Value::Nil => vec![],
        other => {
            return Err(FlightSqlError::MalformedResponse(format!(
                "expected Array for data rows, got {other:?}"
            )));
        }
    };

    let ncols = schema.fields().len();
    let nrows = rows.len();

    let mut builders: Vec<ColumnBuilder> = schema
        .fields()
        .iter()
        .map(|f| ColumnBuilder::new(f.data_type(), nrows))
        .collect();

    for (row_idx, row_val) in rows.into_iter().enumerate() {
        let cells = match row_val {
            Value::Array(c) => c,
            other => {
                return Err(FlightSqlError::MalformedResponse(format!(
                    "row {row_idx}: expected Array, got {other:?}"
                )));
            }
        };

        if cells.len() != ncols {
            return Err(FlightSqlError::ColumnCountMismatch {
                schema: ncols,
                row: cells.len(),
            });
        }

        for (col_idx, cell) in cells.iter().enumerate() {
            builders[col_idx].append(cell, row_idx)?;
        }
    }

    let columns: Vec<ArrayRef> = builders.into_iter().map(|b| b.finish()).collect();
    Ok(RecordBatch::try_new(schema.clone(), columns)?)
}

/// Convert a raw `GRAPH.QUERY` response to a `RecordBatch` **without a
/// pre-registered schema**.
///
/// Column names are taken from the FalkorDB response header row. The Arrow
/// type for each column is inferred by inspecting the first non-null Redis
/// value in that column:
///
/// | Redis value | Arrow type |
/// |---|---|
/// | `Int` | `Int64` |
/// | `Double` | `Float64` |
/// | `Boolean` | `Boolean` |
/// | `BulkString` / `SimpleString` / other | `Utf8` |
/// | all-null column | `Utf8` (fallback) |
///
/// This means integer IDs like `_gar_id`, `src`, and `dst` are returned as
/// `Int64` by default — no registration needed.
pub fn response_to_batch_auto(response: Value) -> Result<RecordBatch> {
    let outer = match response {
        Value::Array(v) => v,
        other => {
            return Err(FlightSqlError::MalformedResponse(format!(
                "expected Array at top level, got {other:?}"
            )));
        }
    };

    if outer.len() < 2 {
        return Ok(RecordBatch::new_empty(Arc::new(
            arrow_schema::Schema::empty(),
        )));
    }

    let col_names = extract_string_vec(&outer[0])?;
    let rows: Vec<&Value> = match &outer[1] {
        Value::Array(r) => r.iter().collect(),
        Value::Nil => vec![],
        other => {
            return Err(FlightSqlError::MalformedResponse(format!(
                "expected Array for data rows, got {other:?}"
            )));
        }
    };

    let ncols = col_names.len();
    let nrows = rows.len();

    if ncols == 0 {
        return Ok(RecordBatch::new_empty(Arc::new(
            arrow_schema::Schema::empty(),
        )));
    }

    // Pass 1 — sniff the Arrow type for each column in a single row-major scan
    // that short-circuits once every column's type is known (M4).
    let col_types = sniff_column_types(&rows, ncols);

    let schema = Arc::new(arrow_schema::Schema::new(
        col_names
            .iter()
            .zip(&col_types)
            .map(|(n, dt)| Field::new(n.as_str(), dt.clone(), true))
            .collect::<Vec<_>>(),
    ));

    // Pass 2 — fill typed builders.
    let mut builders: Vec<ColumnBuilder> = col_types
        .iter()
        .map(|dt| ColumnBuilder::new(dt, nrows))
        .collect();

    for (row_idx, row_val) in rows.iter().enumerate() {
        let cells = match row_val {
            Value::Array(c) => c,
            other => {
                return Err(FlightSqlError::MalformedResponse(format!(
                    "row {row_idx}: expected Array, got {other:?}"
                )));
            }
        };
        // Validate the row width before filling builders. Without this a ragged
        // row (fewer cells than columns) silently under-fills the trailing
        // builders, and the mismatch only surfaces later as an opaque
        // `RecordBatch::try_new` length error. Mirror the schema-driven
        // `response_to_batch`, which already returns `ColumnCountMismatch`.
        if cells.len() != ncols {
            return Err(FlightSqlError::ColumnCountMismatch {
                schema: ncols,
                row: cells.len(),
            });
        }

        for (col_idx, cell) in cells.iter().enumerate() {
            builders[col_idx].append(cell, row_idx)?;
        }
    }

    let columns: Vec<ArrayRef> = builders.into_iter().map(|b| b.finish()).collect();
    Ok(RecordBatch::try_new(schema, columns)?)
}

/// Map a single Redis cell to the Arrow type it fixes for its column, or `None`
/// for a `Nil` cell (which leaves the column undetermined).
///
/// This is the *one* place the per-cell type rule lives; both the original
/// per-column sniff and the single-pass sniff agree because they call it.
#[inline]
fn cell_type(cell: &Value) -> Option<DataType> {
    match cell {
        Value::Nil => None, // null: keep the column undetermined
        Value::Int(_) => Some(DataType::Int64),
        Value::Double(_) => Some(DataType::Float64),
        Value::Boolean(_) => Some(DataType::Boolean),
        _ => Some(DataType::Utf8),
    }
}

/// Infer the Arrow type of every column in a **single row-major pass** that
/// stops as soon as all columns are determined (M4 — single-pass type sniff).
///
/// ## Equivalence to the previous per-column scan
///
/// The previous code called `sniff_column_type(rows, ci)` once per column,
/// each call walking `rows` top-to-bottom and returning the type of the
/// **first non-null cell** in that column (falling back to `Utf8` for an
/// all-null/empty column). Crucially the rule is *first-non-null wins* — there
/// is **no type promotion**: a column's type never changes once a non-null cell
/// is seen, regardless of what later rows contain (e.g. an `Int` first row
/// followed by a `Double` row still yields `Int64`; the builder then coerces
/// the later cells, see `redis_to_i64`). Because the type is a pure function of
/// the *first non-null cell per column* and is independent across columns, the
/// order in which we visit cells does not affect the result, so we can fuse the
/// `ncols` independent passes into one row-major pass:
///
/// * `types[ci]` is written exactly once — on the first non-null cell of column
///   `ci` — and never overwritten, matching "first non-null wins".
/// * columns whose every cell is `Nil` (or that have no cell in any row) are
///   never written and keep the `Utf8` default, matching the all-null fallback.
/// * once `remaining == 0` every column is fixed and no later cell could change
///   any answer, so we stop early — saving the wasted tail scan M4 targets.
///
/// Rows shorter than `ncols` (a malformed/ragged response) are handled exactly
/// as before: `cells.get(ci)` simply yields `None` for the missing trailing
/// columns, leaving them undetermined just like the per-column scan's
/// `cells.get(col_idx)` miss did.
fn sniff_column_types(rows: &[&Value], ncols: usize) -> Vec<DataType> {
    // Default for all-null / empty columns is Utf8 (the original fallback).
    let mut types = vec![DataType::Utf8; ncols];
    // `determined[ci]` flips true the moment column `ci`'s type is fixed.
    let mut determined = vec![false; ncols];
    let mut remaining = ncols;

    for row in rows {
        if remaining == 0 {
            break; // every column resolved — the rest of the scan is wasted work
        }
        let Value::Array(cells) = row else { continue };
        for ci in 0..ncols {
            if determined[ci] {
                continue; // already fixed by an earlier row — skip
            }
            if let Some(cell) = cells.get(ci)
                && let Some(dt) = cell_type(cell)
            {
                types[ci] = dt;
                determined[ci] = true;
                remaining -= 1;
            }
        }
    }
    types
}

// ── Per-column builder ────────────────────────────────────────────────────────

enum ColumnBuilder {
    Int64(Int64Builder),
    Float64(Float64Builder),
    Bool(BooleanBuilder),
    Utf8(StringBuilder),
}

impl ColumnBuilder {
    fn new(dt: &DataType, capacity: usize) -> Self {
        match dt {
            DataType::Int8
            | DataType::Int16
            | DataType::Int32
            | DataType::Int64
            | DataType::UInt8
            | DataType::UInt16
            | DataType::UInt32
            | DataType::UInt64 => Self::Int64(Int64Builder::with_capacity(capacity)),
            DataType::Float16 | DataType::Float32 | DataType::Float64 => {
                Self::Float64(Float64Builder::with_capacity(capacity))
            }
            DataType::Boolean => Self::Bool(BooleanBuilder::with_capacity(capacity)),
            _ => Self::Utf8(StringBuilder::with_capacity(capacity, capacity * 16)),
        }
    }

    fn append(&mut self, val: &Value, row: usize) -> Result<()> {
        match self {
            Self::Int64(b) => b.append_option(redis_to_i64(val, row)?),
            Self::Float64(b) => b.append_option(redis_to_f64(val, row)?),
            Self::Bool(b) => b.append_option(redis_to_bool(val, row)?),
            // Append directly into the builder; for the common valid-UTF-8
            // BulkString this borrows (Cow::Borrowed) with no owned-String copy.
            Self::Utf8(b) => append_redis_str(b, val),
        }
        Ok(())
    }

    fn finish(self) -> ArrayRef {
        match self {
            Self::Int64(mut b) => Arc::new(b.finish()),
            Self::Float64(mut b) => Arc::new(b.finish()),
            Self::Bool(mut b) => Arc::new(b.finish()),
            Self::Utf8(mut b) => Arc::new(b.finish()),
        }
    }
}

// ── Redis Value → scalar coercions ───────────────────────────────────────────

fn redis_to_i64(v: &Value, row: usize) -> Result<Option<i64>> {
    match v {
        Value::Nil => Ok(None),
        Value::Int(i) => Ok(Some(*i)),
        Value::Double(d) => Ok(Some(*d as i64)),
        Value::BulkString(b) => {
            let s = std::str::from_utf8(b).unwrap_or("");
            s.parse::<i64>()
                .map(Some)
                .map_err(|_| malformed(row, "expected i64", s))
        }
        Value::SimpleString(s) => s
            .parse::<i64>()
            .map(Some)
            .map_err(|_| malformed(row, "expected i64", s)),
        other => Err(malformed(
            row,
            "expected Int for Int64",
            &format!("{other:?}"),
        )),
    }
}

fn redis_to_f64(v: &Value, row: usize) -> Result<Option<f64>> {
    match v {
        Value::Nil => Ok(None),
        Value::Double(d) => Ok(Some(*d)),
        Value::Int(i) => Ok(Some(*i as f64)),
        Value::BulkString(b) => {
            let s = std::str::from_utf8(b).unwrap_or("");
            s.parse::<f64>()
                .map(Some)
                .map_err(|_| malformed(row, "expected f64", s))
        }
        Value::SimpleString(s) => s
            .parse::<f64>()
            .map(Some)
            .map_err(|_| malformed(row, "expected f64", s)),
        other => Err(malformed(
            row,
            "expected Double for Float64",
            &format!("{other:?}"),
        )),
    }
}

fn redis_to_bool(v: &Value, row: usize) -> Result<Option<bool>> {
    match v {
        Value::Nil => Ok(None),
        Value::Boolean(b) => Ok(Some(*b)),
        Value::Int(i) => Ok(Some(*i != 0)),
        Value::BulkString(b) => match b.as_slice() {
            b"true" | b"1" => Ok(Some(true)),
            b"false" | b"0" => Ok(Some(false)),
            other => Err(malformed(
                row,
                "expected bool",
                std::str::from_utf8(other).unwrap_or("?"),
            )),
        },
        other => Err(malformed(row, "expected Boolean", &format!("{other:?}"))),
    }
}

/// Append a redis value to a `StringBuilder` without an intermediate owned
/// `String` for the common cases (`from_utf8_lossy` is `Cow::Borrowed` when the
/// bytes are valid UTF-8, and the builder copies into its own buffer anyway).
fn append_redis_str(b: &mut StringBuilder, v: &Value) {
    match v {
        Value::Nil => b.append_null(),
        Value::BulkString(bytes) => b.append_value(String::from_utf8_lossy(bytes)),
        Value::SimpleString(s) => b.append_value(s),
        Value::Int(i) => b.append_value(i.to_string()),
        Value::Double(d) => b.append_value(d.to_string()),
        Value::Boolean(x) => b.append_value(if *x { "true" } else { "false" }),
        other => b.append_value(format!("{other:?}")),
    }
}

fn extract_string_vec(v: &Value) -> Result<Vec<String>> {
    match v {
        Value::Array(items) => items
            .iter()
            .map(|item| match item {
                Value::BulkString(b) => Ok(String::from_utf8_lossy(b).into_owned()),
                Value::SimpleString(s) => Ok(s.clone()),
                other => Ok(format!("{other:?}")),
            })
            .collect(),
        other => Err(FlightSqlError::MalformedResponse(format!(
            "expected Array for column names, got {other:?}"
        ))),
    }
}

fn malformed(row: usize, expected: &str, got: &str) -> FlightSqlError {
    FlightSqlError::MalformedResponse(format!("row {row}: {expected}, got '{got}'"))
}

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

    fn header(names: &[&str]) -> Value {
        Value::Array(names.iter().map(|n| Value::BulkString(n.as_bytes().to_vec())).collect())
    }

    fn row(cells: Vec<Value>) -> Value {
        Value::Array(cells)
    }

    /// A well-formed 2-column response round-trips through the schema-less path.
    #[test]
    fn auto_well_formed_rows_build_batch() {
        let resp = Value::Array(vec![
            header(&["id", "name"]),
            Value::Array(vec![
                row(vec![Value::Int(1), Value::BulkString(b"a".to_vec())]),
                row(vec![Value::Int(2), Value::BulkString(b"b".to_vec())]),
            ]),
        ]);
        let batch = response_to_batch_auto(resp).unwrap();
        assert_eq!(batch.num_rows(), 2);
        assert_eq!(batch.num_columns(), 2);
    }

    /// A ragged row (fewer cells than columns) must be rejected with a precise
    /// `ColumnCountMismatch`, not silently under-fill the trailing builder and
    /// surface later as an opaque `RecordBatch::try_new` length error.
    #[test]
    fn auto_short_row_is_column_count_mismatch() {
        let resp = Value::Array(vec![
            header(&["id", "name"]),
            Value::Array(vec![
                row(vec![Value::Int(1), Value::BulkString(b"a".to_vec())]),
                row(vec![Value::Int(2)]), // ragged: only 1 cell for 2 columns
            ]),
        ]);
        let err = response_to_batch_auto(resp).unwrap_err();
        match err {
            FlightSqlError::ColumnCountMismatch { schema, row } => {
                assert_eq!(schema, 2);
                assert_eq!(row, 1);
            }
            other => panic!("expected ColumnCountMismatch, got {other:?}"),
        }
    }

    /// An over-wide row (more cells than columns) is likewise a mismatch, where
    /// previously the extra cell was silently dropped.
    #[test]
    fn auto_long_row_is_column_count_mismatch() {
        let resp = Value::Array(vec![
            header(&["id"]),
            Value::Array(vec![row(vec![Value::Int(1), Value::Int(99)])]),
        ]);
        let err = response_to_batch_auto(resp).unwrap_err();
        assert!(matches!(
            err,
            FlightSqlError::ColumnCountMismatch { schema: 1, row: 2 }
        ));
    }

    /// A nested Cypher value (a `map`/`node`/`list` — anything Redis returns as a
    /// non-scalar `Value`) has no scalar Arrow type, so [`cell_type`] falls back
    /// to `Utf8` and the cell is rendered via its `Debug` form. This pins that
    /// documented fallback: `RETURN {name:'alice', age:30}` (a map) collapses to
    /// a `Utf8` column whose value is the map's debug string, rather than
    /// erroring or silently dropping the column.
    #[test]
    fn auto_nested_cypher_value_falls_back_to_utf8() {
        // A Cypher map value, as FalkorDB returns it over RESP.
        let cypher_map = Value::Map(vec![
            (
                Value::BulkString(b"name".to_vec()),
                Value::BulkString(b"alice".to_vec()),
            ),
            (Value::BulkString(b"age".to_vec()), Value::Int(30)),
        ]);
        // A Cypher list value (nested array).
        let cypher_list = Value::Array(vec![Value::Int(1), Value::Int(2), Value::Int(3)]);

        let resp = Value::Array(vec![
            header(&["id", "props", "tags"]),
            Value::Array(vec![row(vec![
                Value::Int(7),
                cypher_map.clone(),
                cypher_list.clone(),
            ])]),
        ]);
        let batch = response_to_batch_auto(resp).unwrap();
        assert_eq!(batch.num_rows(), 1);
        assert_eq!(batch.num_columns(), 3);

        // The nested columns sniff to Utf8 (the fallback), the scalar id stays Int64.
        assert_eq!(batch.schema().field(0).data_type(), &DataType::Int64);
        assert_eq!(batch.schema().field(1).data_type(), &DataType::Utf8);
        assert_eq!(batch.schema().field(2).data_type(), &DataType::Utf8);

        // The cell is the value's Debug rendering (what `append_redis_str` emits
        // for a non-scalar), so no data is lost — it is stringified, not dropped.
        let props = batch
            .column(1)
            .as_any()
            .downcast_ref::<arrow_array::StringArray>()
            .unwrap();
        assert_eq!(props.value(0), format!("{cypher_map:?}"));
        let tags = batch
            .column(2)
            .as_any()
            .downcast_ref::<arrow_array::StringArray>()
            .unwrap();
        assert_eq!(tags.value(0), format!("{cypher_list:?}"));
    }

    /// The same Utf8 fallback drives cell-type sniffing directly: a non-scalar
    /// `Value` fixes a column to `Utf8`, while `Nil` leaves it undetermined.
    #[test]
    fn cell_type_maps_nested_values_to_utf8() {
        assert_eq!(cell_type(&Value::Int(1)), Some(DataType::Int64));
        assert_eq!(cell_type(&Value::Double(1.0)), Some(DataType::Float64));
        assert_eq!(cell_type(&Value::Boolean(true)), Some(DataType::Boolean));
        assert_eq!(cell_type(&Value::Nil), None);
        // Nested Cypher shapes → Utf8.
        assert_eq!(cell_type(&Value::Map(vec![])), Some(DataType::Utf8));
        assert_eq!(cell_type(&Value::Array(vec![])), Some(DataType::Utf8));
        assert_eq!(
            cell_type(&Value::BulkString(b"x".to_vec())),
            Some(DataType::Utf8)
        );
    }
}