#[derive(Clone, Debug)]
pub struct Schema {
    schema: arrow::datatypes::Schema,
}

macro_rules! impl_single_col {
    ($fn_name:ident, $dt:expr) => {
        pub fn $fn_name(self, name: impl Into<String>) -> Self {
            self.column(name, $dt)
        }
    };
}

impl Schema {
    pub fn new(schema: arrow::datatypes::Schema) -> Schema {
        Schema { schema }
    }

    pub fn concat(self, other: Self) -> Self {
        let mut fields = self.schema.fields().iter().cloned().collect::<Vec<_>>();
        fields.extend(other.schema.fields().iter().cloned());

        let mut metadata = self.schema.metadata().clone();
        metadata.extend(
            other
                .schema
                .metadata()
                .iter()
                .map(|(k, v)| (k.clone(), v.clone())),
        );

        Schema::new(arrow::datatypes::Schema::new_with_metadata(
            fields, metadata,
        ))
    }

    pub fn column(self, name: impl Into<String>, data_type: arrow::datatypes::DataType) -> Self {
        self.concat(Schema::new(arrow::datatypes::Schema::new(vec![
            arrow::datatypes::Field::new(name.into(), data_type, true),
        ])))
    }

    impl_single_col!(column_f64, arrow::datatypes::DataType::Float64);
    impl_single_col!(column_i64, arrow::datatypes::DataType::Int64);
    impl_single_col!(column_u64, arrow::datatypes::DataType::UInt64);
    impl_single_col!(column_bool, arrow::datatypes::DataType::Boolean);
    impl_single_col!(column_utf8, arrow::datatypes::DataType::Utf8);

    pub fn into_arrow(self) -> arrow::datatypes::Schema {
        self.schema
    }
}

impl From<arrow::datatypes::Schema> for Schema {
    fn from(schema: arrow::datatypes::Schema) -> Self {
        Schema::new(schema)
    }
}
