hamelin_lib 0.19.0

use std::collections::HashMap;
use std::fmt::{Display, Formatter};
use std::hash::Hash;
use std::sync::Arc;

use anyhow::bail;
use ordermap::OrderMap;
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use tsify::Tsify;
use vecmap::VecMap;

use crate::provider::EnvironmentProvider;
use crate::tree::ast::dataset_identifier::QualifiedDatasetIdentifier;
use crate::tree::ast::identifier::SimpleIdentifier;
use crate::types::array::Array;
use crate::types::decimal_type::Decimal;
use crate::types::map::Map;
use crate::types::range::Range;
use crate::types::struct_type::Struct;
use crate::types::tuple::Tuple;
use crate::types::Type;

#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize, Tsify)]
#[tsify(into_wasm_abi, from_wasm_abi, hashmap_as_object)]
pub struct Catalog(pub HashMap<QualifiedDatasetIdentifier, Vec<Column>>);

#[derive(Default, Clone, Debug)]
pub struct CatalogBuilder {
    data_sets: HashMap<QualifiedDatasetIdentifier, Vec<Column>>,
}

impl CatalogBuilder {
    pub fn with(mut self, name: QualifiedDatasetIdentifier, columns: Vec<Column>) -> Self {
        self.data_sets.insert(name, columns);
        self
    }

    pub fn with_data_set(mut self, data_set: DataSetBuilder) -> Self {
        let (name, columns) = data_set.build();
        self.data_sets.insert(name, columns);
        self
    }

    pub fn build(self) -> Catalog {
        Catalog(self.data_sets)
    }
}

#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize, Tsify)]
#[tsify(into_wasm_abi, from_wasm_abi)]
pub struct Column {
    pub name: SimpleIdentifier,
    #[serde(rename = "type")]
    pub typ: HamelinType,
}

impl Column {
    pub fn new(name: &str, typ: HamelinType) -> Self {
        Column {
            name: name.into(),
            typ,
        }
    }
}

impl From<(SimpleIdentifier, Type)> for Column {
    fn from((name, typ): (SimpleIdentifier, Type)) -> Self {
        Column {
            name,
            typ: typ.into(),
        }
    }
}

impl From<&Struct> for Vec<Column> {
    fn from(s: &Struct) -> Self {
        s.iter()
            .map(|(name, typ)| Column::from((name.clone(), typ.clone())))
            .collect()
    }
}

impl From<Struct> for Vec<Column> {
    fn from(s: Struct) -> Self {
        Vec::from(&s)
    }
}

impl Display for Column {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        let t = Type::try_from(self.typ.clone()).map_err(|_| std::fmt::Error)?;
        write!(f, "  {}: {}", self.name, t)
    }
}

impl Hash for Column {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.name.hash(state);
        Type::try_from(self.typ.clone())
            .map(|t| t.hash(state))
            .unwrap_or_default();
    }
}

#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize, Tsify)]
#[tsify(into_wasm_abi, from_wasm_abi)]
#[serde(rename_all = "snake_case")]
pub enum HamelinType {
    Binary,
    Boolean,
    Interval,
    CalendarInterval,
    Int,
    Double,
    Rows,
    String,
    Timestamp,
    Unknown,
    Decimal {
        precision: i32,
        scale: i32,
    },
    Array {
        element_type: Box<HamelinType>,
    },
    Map {
        key_type: Box<HamelinType>,
        value_type: Box<HamelinType>,
    },
    Tuple {
        elements: Vec<HamelinType>,
    },
    Variant,
    Range {
        of: Box<HamelinType>,
    },
    RangeInclusive {
        of: Box<HamelinType>,
    },
    Struct(Vec<Column>),
}

impl HamelinType {
    pub fn to_sql(&self) -> Result<String, anyhow::Error> {
        Type::try_from(self.clone())
            .map(|t| t.to_sql().map(|s| s.to_string()))
            .unwrap_or(Err(anyhow::anyhow!("Failed to convert HamelinType to SQL")))
    }
}

impl From<Type> for HamelinType {
    fn from(value: Type) -> Self {
        match value {
            Type::Binary => HamelinType::Binary,
            Type::Boolean => HamelinType::Boolean,
            Type::Interval => HamelinType::Interval,
            Type::CalendarInterval => HamelinType::CalendarInterval,
            Type::Int => HamelinType::Int,
            Type::Double => HamelinType::Double,
            Type::Rows => HamelinType::Rows,
            Type::String => HamelinType::String,
            Type::Timestamp => HamelinType::Timestamp,
            Type::Unknown => HamelinType::Unknown,
            Type::Decimal(d) => HamelinType::Decimal {
                precision: d.precision,
                scale: d.scale,
            },
            Type::Array(a) => HamelinType::Array {
                element_type: Box::new(Self::from(a.element_type.as_ref().clone())),
            },
            Type::Map(m) => HamelinType::Map {
                key_type: Box::new(Self::from(m.key_type.as_ref().clone())),
                value_type: Box::new(Self::from(m.value_type.as_ref().clone())),
            },
            Type::Tuple(t) => HamelinType::Tuple {
                elements: t
                    .elements
                    .iter()
                    .map(|e| HamelinType::from(e.as_ref().clone()))
                    .collect(),
            },
            Type::Variant => HamelinType::Variant,
            Type::Range(r) => HamelinType::Range {
                of: Box::new(Self::from(r.of.as_ref().clone())),
            },
            Type::RangeInclusive(r) => HamelinType::RangeInclusive {
                of: Box::new(Self::from(r.of.as_ref().clone())),
            },
            Type::Struct(s) => HamelinType::Struct(Vec::from(&s)),
            // Function types are internal to expression type-checking and should never
            // appear in catalog types. Map to Unknown as a fallback.
            Type::Function(_) => HamelinType::Unknown,
        }
    }
}

impl From<Decimal> for HamelinType {
    fn from(value: Decimal) -> Self {
        HamelinType::Decimal {
            precision: value.precision,
            scale: value.scale,
        }
    }
}

impl From<Array> for HamelinType {
    fn from(value: Array) -> Self {
        HamelinType::Array {
            element_type: Box::new(value.element_type.as_ref().clone().into()),
        }
    }
}

impl From<Map> for HamelinType {
    fn from(value: Map) -> Self {
        HamelinType::Map {
            key_type: Box::new(value.key_type.as_ref().clone().into()),
            value_type: Box::new(value.value_type.as_ref().clone().into()),
        }
    }
}

impl From<Tuple> for HamelinType {
    fn from(value: Tuple) -> Self {
        HamelinType::Tuple {
            elements: value
                .elements
                .iter()
                .map(|e| HamelinType::from(e.as_ref().clone()))
                .collect(),
        }
    }
}

impl From<Range> for HamelinType {
    fn from(value: Range) -> Self {
        HamelinType::Range {
            of: Box::new(value.of.as_ref().clone().into()),
        }
    }
}

impl From<Struct> for HamelinType {
    fn from(value: Struct) -> Self {
        HamelinType::Struct(Vec::from(&value))
    }
}

#[derive(Debug, Clone)]
pub struct DataSetBuilder {
    pub name: QualifiedDatasetIdentifier,
    pub columns: Vec<Column>,
}

impl DataSetBuilder {
    pub fn new(name: QualifiedDatasetIdentifier) -> Self {
        DataSetBuilder {
            name,
            columns: vec![],
        }
    }

    pub fn with(mut self, name: SimpleIdentifier, typ: Type) -> Self {
        self.columns.push(Column {
            name,
            typ: typ.into(),
        });
        self
    }

    pub fn build(self) -> (QualifiedDatasetIdentifier, Vec<Column>) {
        (self.name, self.columns)
    }

    pub fn parse(
        self,
    ) -> anyhow::Result<(QualifiedDatasetIdentifier, OrderMap<SimpleIdentifier, Type>)> {
        Ok((
            self.name,
            self.columns
                .into_iter()
                .map(|c| Ok((c.name, c.typ.try_into()?)))
                .collect::<anyhow::Result<OrderMap<SimpleIdentifier, Type>>>()?,
        ))
    }
}

#[derive(Debug, Clone, Default)]
pub struct CatalogProvider {
    data_sets: Arc<RwLock<HashMap<QualifiedDatasetIdentifier, OrderMap<SimpleIdentifier, Type>>>>,
}

impl CatalogProvider {
    pub fn new(
        data_sets: HashMap<QualifiedDatasetIdentifier, OrderMap<SimpleIdentifier, Type>>,
    ) -> Self {
        CatalogProvider {
            data_sets: Arc::new(RwLock::new(data_sets)),
        }
    }

    pub fn set(&self, name: QualifiedDatasetIdentifier, columns: OrderMap<SimpleIdentifier, Type>) {
        let mut data_sets = self.data_sets.write();
        data_sets.insert(name, columns);
    }

    pub fn replace(&self, other: CatalogProvider) {
        let mut data_sets = self.data_sets.write();
        let other_data_sets = other.data_sets.read();
        *data_sets = other_data_sets.clone();
    }

    pub fn lookup(
        &self,
        name: &QualifiedDatasetIdentifier,
    ) -> Option<OrderMap<SimpleIdentifier, Type>> {
        let data_sets = self.data_sets.read();
        data_sets.get(name).cloned()
    }

    pub fn dataset_names(&self) -> Vec<QualifiedDatasetIdentifier> {
        let data_sets = self.data_sets.read();
        data_sets.keys().cloned().collect()
    }

    pub fn clear(&self) {
        let mut data_sets = self.data_sets.write();
        data_sets.clear();
    }

    pub fn unset(&self, name: &QualifiedDatasetIdentifier) {
        let mut data_sets = self.data_sets.write();
        data_sets.remove(name);
    }

    pub fn get_catalog(&self) -> Catalog {
        Catalog(
            self.data_sets
                .read()
                .iter()
                .map(|(k, v)| {
                    (
                        k.clone(),
                        v.iter()
                            .map(|(k, v)| Column::from((k.clone(), v.clone())))
                            .collect(),
                    )
                })
                .collect(),
        )
    }
}

impl TryFrom<Catalog> for CatalogProvider {
    type Error = anyhow::Error;

    fn try_from(value: Catalog) -> Result<Self, Self::Error> {
        let mut data_sets = HashMap::new();
        for (name, columns) in value.0.into_iter() {
            data_sets.insert(
                name,
                columns
                    .into_iter()
                    .map(|c| Ok((c.name, c.typ.try_into()?)))
                    .collect::<Result<OrderMap<_, _>, anyhow::Error>>()?,
            );
        }

        Ok(CatalogProvider::new(data_sets))
    }
}

impl PartialEq for CatalogProvider {
    fn eq(&self, other: &Self) -> bool {
        self.data_sets.read().eq(&*other.data_sets.read())
    }
}

impl Eq for CatalogProvider {}

impl EnvironmentProvider for CatalogProvider {
    fn reflect_columns(&self, name: &QualifiedDatasetIdentifier) -> anyhow::Result<Struct> {
        match self.lookup(name) {
            None => {
                bail!("Dataset not found: {}", name);
            }
            Some(data_set) => {
                let fields: VecMap<_, _> = data_set.into_iter().collect();
                Ok(Struct::new(fields))
            }
        }
    }

    fn reflect_datasets(&self) -> anyhow::Result<Vec<QualifiedDatasetIdentifier>> {
        Ok(self.dataset_names())
    }
}

impl Display for CatalogProvider {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        let data_sets = self.data_sets.read();
        write!(f, "Catalog:\n")?;
        for name in data_sets.keys() {
            write!(f, "{}:\n", name)?;
        }

        Ok(())
    }
}

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

    use crate::types::map::Map;
    use crate::types::{BOOLEAN, DOUBLE, INT, STRING, VARIANT};

    fn test_dataset(name: &str) -> QualifiedDatasetIdentifier {
        QualifiedDatasetIdentifier::from_canonical_str(&format!("test:{name}"))
            .expect("test dataset path")
    }

    /// Bounce through serialization.
    /// Useful to do this test as we do other tests.
    fn bounce(provider: CatalogProvider) {
        let catalog = provider.get_catalog();
        let serialized_cat = serde_json::to_string(&catalog).unwrap();
        let deserizlied_cat: Catalog = serde_json::from_str(&serialized_cat).unwrap();
        let deserialized_provider = CatalogProvider::try_from(deserizlied_cat).unwrap();
        assert_eq!(provider, deserialized_provider);
    }

    #[test]
    fn test_empty_catalog() {
        let catalog = CatalogProvider::default();
        assert_eq!(catalog.lookup(&test_dataset("missing")), None);
        bounce(catalog);
    }

    #[test]
    fn test_catalog_new() {
        let catalog = CatalogBuilder::default()
            .with_data_set(
                DataSetBuilder::new(test_dataset("test"))
                    .with(SimpleIdentifier::new("a"), INT)
                    .with(SimpleIdentifier::new("b"), STRING),
            )
            .build();

        let provider = CatalogProvider::try_from(catalog).unwrap();

        assert_eq!(
            provider.lookup(&test_dataset("test")),
            Some(
                vec![
                    (SimpleIdentifier::new("a"), INT),
                    (SimpleIdentifier::new("b"), STRING),
                ]
                .into_iter()
                .collect::<OrderMap<_, _>>()
            )
        );

        bounce(provider);
    }

    #[test]
    fn test_catalog_set_lookup() {
        let (data_set_name, data_set_cols) = DataSetBuilder::new(test_dataset("test"))
            .with(SimpleIdentifier::new("a"), INT)
            .with(SimpleIdentifier::new("b"), STRING)
            .parse()
            .unwrap();

        let provider = CatalogProvider::default();
        provider.set(data_set_name, data_set_cols);

        assert_eq!(
            provider.lookup(&test_dataset("test")),
            Some(
                vec![
                    (SimpleIdentifier::new("a"), INT),
                    (SimpleIdentifier::new("b"), STRING),
                ]
                .into_iter()
                .collect::<OrderMap<_, _>>()
            )
        );
    }

    #[test]
    fn test_catalog_complex() {
        let catalog = CatalogBuilder::default()
            .with_data_set(
                DataSetBuilder::new(test_dataset("test"))
                    .with(SimpleIdentifier::new("a"), INT)
                    .with(SimpleIdentifier::new("b"), STRING)
                    .with(SimpleIdentifier::new("c"), BOOLEAN)
                    .with(SimpleIdentifier::new("d"), DOUBLE)
                    .with(SimpleIdentifier::new("e"), VARIANT)
                    .with(
                        SimpleIdentifier::new("f"),
                        Struct::default()
                            .with_str("a", INT)
                            .with_str("b", STRING)
                            .with_str("c", BOOLEAN)
                            .with_str("d", DOUBLE)
                            .into(),
                    )
                    .with(SimpleIdentifier::new("f"), Map::new(STRING, INT).into()),
            )
            .build();

        let provider = CatalogProvider::try_from(catalog).unwrap();

        assert_eq!(
            provider.lookup(&test_dataset("test")),
            Some(
                vec![
                    (SimpleIdentifier::new("a"), INT,),
                    (SimpleIdentifier::new("b"), STRING,),
                    (SimpleIdentifier::new("c"), BOOLEAN,),
                    (SimpleIdentifier::new("d"), DOUBLE,),
                    (SimpleIdentifier::new("e"), VARIANT,),
                    (
                        SimpleIdentifier::new("f"),
                        Struct::default()
                            .with_str("a", INT)
                            .with_str("b", STRING)
                            .with_str("c", BOOLEAN)
                            .with_str("d", DOUBLE)
                            .into(),
                    ),
                    (SimpleIdentifier::new("f"), Map::new(STRING, INT).into(),),
                ]
                .into_iter()
                .collect::<OrderMap<_, _>>()
            )
        );

        bounce(provider);
    }

    #[test]
    fn test_keyword_column_names() {
        let dataset = test_dataset("keywords");
        let catalog = CatalogBuilder::default()
            .with_data_set(
                DataSetBuilder::new(dataset.clone())
                    .with(SimpleIdentifier::new("select"), INT)
                    .with(SimpleIdentifier::new("from"), STRING),
            )
            .build();

        let provider = CatalogProvider::try_from(catalog).unwrap();

        assert_eq!(
            provider.lookup(&dataset),
            Some(
                vec![
                    (SimpleIdentifier::new("select"), INT,),
                    (SimpleIdentifier::new("from"), STRING,),
                ]
                .into_iter()
                .collect::<OrderMap<_, _>>()
            )
        );

        bounce(provider);
    }

    #[test]
    pub fn test_struct_keyword_column_names() {
        let catalog = CatalogBuilder::default()
            .with_data_set(
                DataSetBuilder::new(test_dataset("let")).with(
                    SimpleIdentifier::new("select"),
                    Struct::default()
                        .with_str("as", INT)
                        .with_str("from", STRING)
                        .into(),
                ),
            )
            .build();

        let provider = CatalogProvider::try_from(catalog).unwrap();

        assert_eq!(
            provider.lookup(&test_dataset("let")),
            Some(
                vec![(
                    SimpleIdentifier::new("select"),
                    Struct::default()
                        .with_str("as", INT)
                        .with_str("from", STRING)
                        .into(),
                )]
                .into_iter()
                .collect::<OrderMap<_, _>>()
            )
        );

        bounce(provider);
    }

    #[test]
    pub fn test_catalog_replace() {
        let catalog1 = CatalogBuilder::default()
            .with_data_set(
                DataSetBuilder::new(test_dataset("first"))
                    .with(SimpleIdentifier::new("a"), INT)
                    .with(SimpleIdentifier::new("b"), STRING),
            )
            .build();

        let catalog2 = CatalogBuilder::default()
            .with_data_set(
                DataSetBuilder::new(test_dataset("second"))
                    .with(SimpleIdentifier::new("c"), DOUBLE)
                    .with(SimpleIdentifier::new("d"), BOOLEAN),
            )
            .build();

        let provider = CatalogProvider::try_from(catalog1.clone()).unwrap();
        assert_eq!(provider.get_catalog(), catalog1);

        provider.replace(CatalogProvider::try_from(catalog2).unwrap());

        assert_eq!(
            provider.get_catalog(),
            CatalogBuilder::default()
                .with_data_set(
                    DataSetBuilder::new(test_dataset("second"))
                        .with(SimpleIdentifier::new("c"), DOUBLE)
                        .with(SimpleIdentifier::new("d"), BOOLEAN),
                )
                .build()
        );
    }

    #[test]
    pub fn test_hamelin_to_sql() {
        assert_eq!(HamelinType::Boolean.to_sql().unwrap(), "BOOLEAN");
        assert_eq!(HamelinType::String.to_sql().unwrap(), "VARCHAR");
        assert_eq!(HamelinType::Double.to_sql().unwrap(), "DOUBLE");
        assert_eq!(
            HamelinType::Timestamp.to_sql().unwrap(),
            "TIMESTAMP(6) WITH TIME ZONE"
        );
        assert_eq!(
            HamelinType::Interval.to_sql().unwrap(),
            "INTERVAL DAY TO SECOND"
        );
        assert_eq!(HamelinType::Binary.to_sql().unwrap(), "VARBINARY");
        assert_eq!(
            HamelinType::Map {
                key_type: Box::new(HamelinType::String),
                value_type: Box::new(HamelinType::String)
            }
            .to_sql()
            .unwrap(),
            "MAP(VARCHAR, VARCHAR)"
        );
        assert_eq!(
            HamelinType::Array {
                element_type: Box::new(HamelinType::String)
            }
            .to_sql()
            .unwrap(),
            "ARRAY(VARCHAR)"
        );
        assert_eq!(
            HamelinType::Struct(vec![Column::new("a", HamelinType::String)])
                .to_sql()
                .unwrap(),
            "ROW(\"a\" VARCHAR)"
        );
        assert_eq!(
            HamelinType::Decimal {
                precision: 10,
                scale: 2
            }
            .to_sql()
            .unwrap(),
            "DECIMAL(10, 2)"
        );
        assert_eq!(HamelinType::Variant.to_sql().unwrap(), "JSON");
        assert_eq!(HamelinType::Rows.to_sql().unwrap(), "BIGINT");
    }

    #[test]
    fn test_calendar_interval_serialization() {
        let column = Column::new("interval_col", HamelinType::CalendarInterval);
        let catalog = Catalog(
            [(test_dataset("test"), vec![column.clone()])]
                .iter()
                .cloned()
                .collect(),
        );
        let serialized = serde_json::to_string(&catalog).unwrap();
        // Should contain "calendar_interval" as the type
        assert!(
            serialized.contains("\"calendar_interval\""),
            "Serialized catalog did not contain \"calendar_interval\": {}",
            serialized
        );
        // Round-trip
        let deserialized: Catalog = serde_json::from_str(&serialized).unwrap();
        let key = test_dataset("test");
        assert_eq!(deserialized.0[&key][0].typ, HamelinType::CalendarInterval);
    }
}