hamelin_executor 0.9.4

use std::fmt::Display;

use comfy_table::{presets::UTF8_FULL, ContentArrangement, Table};
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};

use hamelin_lib::catalog::{Column, HamelinType};

#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct ResultSet {
    pub columns: Vec<Column>,
    pub rows: Vec<Vec<Value>>,
}

impl PartialEq for ResultSet {
    fn eq(&self, other: &Self) -> bool {
        if self.columns != other.columns {
            return false;
        }
        if self.rows.len() != other.rows.len() {
            return false;
        }
        for (self_row, other_row) in self.rows.iter().zip(other.rows.iter()) {
            if self_row.len() != other_row.len() {
                return false;
            }
            for (col_idx, (self_val, other_val)) in
                self_row.iter().zip(other_row.iter()).enumerate()
            {
                let col_type = self.columns.get(col_idx).map(|c| &c.typ);
                if !values_equal_for_type(self_val, other_val, col_type) {
                    return false;
                }
            }
        }
        true
    }
}

impl Eq for ResultSet {}

/// Compare two values, using order-neutral JSON comparison for Variant columns.
/// Variant values are serialized as JSON strings, and key order may differ between backends.
/// This function parses the strings and compares the parsed values.
fn values_equal_for_type(a: &Value, b: &Value, col_type: Option<&HamelinType>) -> bool {
    if matches!(col_type, Some(HamelinType::Variant)) {
        if let (Value::String(a_str), Value::String(b_str)) = (a, b) {
            if let (Ok(a_json), Ok(b_json)) = (
                serde_json::from_str::<Value>(a_str),
                serde_json::from_str::<Value>(b_str),
            ) {
                return a_json == b_json;
            }
        }
    }
    a == b
}

impl ResultSet {
    pub fn new(columns: Vec<Column>, rows: Vec<Vec<Value>>) -> Self {
        Self { columns, rows }
    }

    pub fn builder() -> ResultSetBuilder {
        ResultSetBuilder::default()
    }

    pub fn zipped_rows(
        &self,
    ) -> impl Iterator<Item = impl Iterator<Item = (&Value, &Column)> + '_> + '_ {
        self.rows
            .iter()
            .map(|row| row.iter().zip(self.columns.iter()))
    }

    pub fn rows_with_keys(&self) -> impl Iterator<Item = Map<String, Value>> + '_ {
        self.zipped_rows().map(|row| {
            let mut map = Map::new();
            for (value, column) in row {
                map.insert(
                    column.name.to_string(),
                    transform_for_display(value, &column.typ),
                );
            }
            map
        })
    }

    pub fn rendered_rows(&self) -> impl Iterator<Item = impl Iterator<Item = String> + '_> + '_ {
        self.zipped_rows().map(|row| {
            row.map(|(value, col)| {
                let ht = transform_for_display(value, &col.typ);
                serde_json::to_string_pretty(&ht).unwrap_or("cannot deserialize".to_string())
            })
        })
    }
}

impl Display for ResultSet {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let mut table = Table::new();
        table
            .load_preset(UTF8_FULL)
            .set_content_arrangement(ContentArrangement::Dynamic)
            .set_header(self.columns.iter().map(|c| c.name.to_string()))
            .add_rows(self.rendered_rows());

        #[cfg(target_arch = "wasm32")]
        table.set_width(100);

        write!(f, "{table}")
    }
}

/// Transform a serde_json Value to a version that is appropriate for display.
///
/// The result set format elides unnecessary type information for structs and such, but
/// this makes the format not great for showing to a human. This function transforms
/// the internal result set format to a format more appropriate for display, but one
/// that that still conforms to serde json.
fn transform_for_display(value: &Value, field_type: &HamelinType) -> Value {
    // TODO: Just go straight to string here? Render hamelin literals.
    match (value, field_type) {
        (Value::Array(arr), HamelinType::Tuple { elements }) => Value::Array(
            arr.iter()
                .zip(elements.iter())
                .map(|(v, t)| transform_for_display(v, t))
                .collect(),
        ),

        (Value::Array(arr), HamelinType::Array { element_type }) => Value::Array(
            arr.iter()
                .map(|v| transform_for_display(v, element_type))
                .collect(),
        ),

        (Value::Object(o), HamelinType::Map { value_type, .. }) => {
            let map = o.iter().fold(Map::new(), |mut acc, (k, v)| {
                acc.insert(k.to_string(), transform_for_display(v, value_type));
                acc
            });
            Value::Object(map)
        }

        (Value::Array(arr), HamelinType::Struct(fields)) => {
            let map =
                fields
                    .iter()
                    .zip(arr.iter())
                    .fold(Map::new(), |mut acc, (col, field_val)| {
                        acc.insert(
                            col.name.to_string(),
                            transform_for_display(field_val, &col.typ),
                        );
                        acc
                    });
            Value::Object(map)
        }

        (Value::Array(arr), HamelinType::Range { of } | HamelinType::RangeInclusive { of })
            if arr.len() == 2 =>
        {
            let op = if matches!(field_type, HamelinType::RangeInclusive { .. }) {
                "..="
            } else {
                ".."
            };
            match (&arr[0], &arr[1], of.as_ref()) {
                (
                    Value::String(start),
                    Value::String(end),
                    HamelinType::Timestamp | HamelinType::Interval | HamelinType::Decimal { .. },
                ) => Value::String(format!("{}{}{}", start, op, end)),
                (start, end, _) => Value::String(format!("{}{}{}", start, op, end)),
            }
        }

        _ => value.clone(),
    }
}

#[derive(Debug, Default)]
pub struct ResultSetBuilder {
    columns: Vec<Column>,
    rows: Vec<Vec<Value>>,
}

impl ResultSetBuilder {
    pub fn with_column(mut self, name: &str, typ: HamelinType) -> Self {
        self.columns.push(Column::new(name, typ));
        self
    }

    pub fn with_row(mut self, row: Vec<Value>) -> Self {
        self.rows.push(row);
        self
    }

    pub fn build(self) -> ResultSet {
        ResultSet {
            columns: self.columns,
            rows: self.rows,
        }
    }
}

#[cfg(test)]
mod tests {
    use comfy_table::{presets::UTF8_FULL, ContentArrangement, Table};
    use pretty_assertions::assert_eq;
    use serde_json::{json, Value};

    use hamelin_lib::{
        catalog::{Column, HamelinType},
        types::{
            array::Array, decimal_type::Decimal, map::Map, range::Range, struct_type::Struct,
            tuple::Tuple, Type,
        },
    };

    use super::ResultSet;

    #[test]
    fn test_serialize() {
        assert_eq!(
            serde_json::to_string_pretty(&ResultSet {
                columns: vec![
                    Column {
                        name: "i".into(),
                        typ: HamelinType::Int,
                    },
                    Column {
                        name: "s".into(),
                        typ: HamelinType::String,
                    },
                    Column {
                        name: "ts".into(),
                        typ: HamelinType::Timestamp,
                    },
                    Column {
                        name: "b".into(),
                        typ: HamelinType::Boolean,
                    },
                    Column {
                        name: "dec".into(),
                        typ: HamelinType::Decimal {
                            precision: 3,
                            scale: 2
                        },
                    },
                    Column {
                        name: "d".into(),
                        typ: HamelinType::Double,
                    },
                    Column {
                        name: "arr".into(),
                        typ: HamelinType::Array {
                            element_type: Box::new(HamelinType::Int)
                        },
                    },
                    Column {
                        name: "tuple".into(),
                        typ: HamelinType::Tuple {
                            elements: vec![HamelinType::Int, HamelinType::String]
                        }
                    },
                    Column {
                        name: "struct".into(),
                        typ: HamelinType::from(Type::from(
                            Struct::default()
                                .with_str("x", Type::Int)
                                .with_str("y", Type::String)
                        )),
                    },
                    Column {
                        name: "map".into(),
                        typ: HamelinType::Map {
                            key_type: Box::new(HamelinType::Int),
                            value_type: Box::new(HamelinType::String)
                        },
                    },
                    Column {
                        name: "interval".into(),
                        typ: HamelinType::Interval,
                    },
                    Column {
                        name: "range".into(),
                        typ: HamelinType::Range {
                            of: Box::new(HamelinType::Timestamp),
                        },
                    },
                    Column {
                        name: "null".into(),
                        typ: HamelinType::Unknown,
                    }
                ],
                rows: vec![vec![
                    json!(123),
                    json!("hello"),
                    json!("2024-01-01T00:00:00+00:00"),
                    json!(true),
                    json!("3.14"),
                    json!(3.14),
                    json!([1, 2, 3]),
                    json!([1, "a"]),
                    json!([1, "a"]),
                    json!({"1": "a", "2": "b"}),
                    json!("1 00:00:00.000"),
                    json!(["2024-01-01T00:00:00+00:00", "2025-01-01T00:00:00+00:00"]),
                    json!("`null`")
                ]],
            })
            .unwrap(),
            r#"{
  "columns": [
    {
      "name": "i",
      "type": "int"
    },
    {
      "name": "s",
      "type": "string"
    },
    {
      "name": "ts",
      "type": "timestamp"
    },
    {
      "name": "b",
      "type": "boolean"
    },
    {
      "name": "dec",
      "type": {
        "decimal": {
          "precision": 3,
          "scale": 2
        }
      }
    },
    {
      "name": "d",
      "type": "double"
    },
    {
      "name": "arr",
      "type": {
        "array": {
          "element_type": "int"
        }
      }
    },
    {
      "name": "tuple",
      "type": {
        "tuple": {
          "elements": [
            "int",
            "string"
          ]
        }
      }
    },
    {
      "name": "struct",
      "type": {
        "struct": [
          {
            "name": "x",
            "type": "int"
          },
          {
            "name": "y",
            "type": "string"
          }
        ]
      }
    },
    {
      "name": "map",
      "type": {
        "map": {
          "key_type": "int",
          "value_type": "string"
        }
      }
    },
    {
      "name": "interval",
      "type": "interval"
    },
    {
      "name": "range",
      "type": {
        "range": {
          "of": "timestamp"
        }
      }
    },
    {
      "name": "`null`",
      "type": "unknown"
    }
  ],
  "rows": [
    [
      123,
      "hello",
      "2024-01-01T00:00:00+00:00",
      true,
      "3.14",
      3.14,
      [
        1,
        2,
        3
      ],
      [
        1,
        "a"
      ],
      [
        1,
        "a"
      ],
      {
        "1": "a",
        "2": "b"
      },
      "1 00:00:00.000",
      [
        "2024-01-01T00:00:00+00:00",
        "2025-01-01T00:00:00+00:00"
      ],
      "`null`"
    ]
  ]
}"#
        );
    }

    #[test]
    fn test_display() -> anyhow::Result<()> {
        let mut expected = Table::new();
        expected
            .load_preset(UTF8_FULL)
            .set_content_arrangement(ContentArrangement::Dynamic)
            .set_header(vec![
                "i", "s", "ts", "b", "dec", "d", "arr", "tuple", "struct", "map", "interval",
                "range", "`null`",
            ])
            .add_rows(vec![vec![
                "123".to_string(),
                "\"hello\"".to_string(),
                "\"2024-01-01T00:00:00+00:00\"".to_string(),
                "true".to_string(),
                "\"3.14\"".to_string(),
                "3.14".to_string(),
                "[\n  1,\n  2,\n  3\n]".to_string(),
                "[\n  1,\n  \"a\"\n]".to_string(),
                "{\n  \"x\": 1,\n  \"y\": \"a\"\n}".to_string(),
                "{\n  \"1\": \"a\",\n  \"2\": \"b\"\n}".to_string(),
                "\"1 00:00:00.000\"".to_string(),
                json!("2024-01-01T00:00:00+00:00..2025-01-01T00:00:00+00:00").to_string(),
                "null".to_string(),
            ]]);

        assert_eq!(
            ResultSet {
                columns: vec![
                    Column {
                        name: "i".into(),
                        typ: HamelinType::Int,
                    },
                    Column {
                        name: "s".into(),
                        typ: HamelinType::String,
                    },
                    Column {
                        name: "ts".into(),
                        typ: HamelinType::Timestamp,
                    },
                    Column {
                        name: "b".into(),
                        typ: HamelinType::Boolean,
                    },
                    Column {
                        name: "dec".into(),
                        typ: HamelinType::from(Type::Decimal(Decimal::new(3, 2)?)),
                    },
                    Column {
                        name: "d".into(),
                        typ: HamelinType::Double,
                    },
                    Column {
                        name: "arr".into(),
                        typ: HamelinType::from(Type::Array(Array::new(Type::Int))),
                    },
                    Column {
                        name: "tuple".into(),
                        typ: HamelinType::from(Type::Tuple(Tuple::new(vec![
                            Type::Int,
                            Type::String
                        ]))),
                    },
                    Column {
                        name: "struct".into(),
                        typ: HamelinType::from(Type::Struct(
                            Struct::default()
                                .with_str("x", Type::Int)
                                .with_str("y", Type::String)
                        )),
                    },
                    Column {
                        name: "map".into(),
                        typ: HamelinType::from(Type::Map(Map::new(Type::Int, Type::String)))
                    },
                    Column {
                        name: "interval".into(),
                        typ: HamelinType::Interval,
                    },
                    Column {
                        name: "range".into(),
                        typ: HamelinType::from(Type::Range(Range::new(Type::Timestamp))),
                    },
                    Column {
                        name: "null".into(),
                        typ: HamelinType::Unknown,
                    }
                ],
                rows: vec![vec![
                    json!(123),
                    json!("hello"),
                    json!("2024-01-01T00:00:00+00:00"),
                    json!(true),
                    json!("3.14"),
                    json!(3.14),
                    json!([1, 2, 3]),
                    json!([1, "a"]),
                    json!([1, "a"]),
                    json!({"1": "a", "2": "b"}),
                    json!("1 00:00:00.000"),
                    json!(["2024-01-01T00:00:00+00:00", "2025-01-01T00:00:00+00:00"]),
                    Value::Null,
                ]],
            }
            .to_string(),
            expected.to_string(),
        );
        Ok(())
    }
}