pub mod dtype;
pub(crate) mod preprocessing;
// pub mod metadata;
// pub mod schema;

pub use preprocessing::NullStrategy;

use std::collections::HashMap;
use std::sync::Arc;
use std::time::Instant;

use crate::Result;
use crate::dataframe::preprocessing::apply_null_strategy;
use arrow::{
    array::RecordBatch,
    datatypes::{Field, Schema},
    util::pretty::pretty_format_batches,
};

#[derive(Debug, Clone)]
pub struct DataFrame {
    // arrow_table: RecordBatch
    // metadata: DataFrameMetaData,
    record_batch: RecordBatch,
}

impl DataFrame {
    pub(crate) fn from_internal(record_batch: RecordBatch) -> Result<DataFrame> {
        autark_telemetry::init_from_env();
        let t0 = Instant::now();
        let record_batch = apply_null_strategy(&record_batch, NullStrategy::Drop)?;

        // shall be negligble.
        autark_telemetry::log::info!(
            target: "autark_dataframe::from_internal",
            "from_internal took {:?} (rows={}, cols={})",
            t0.elapsed(),
            record_batch.num_rows(),
            record_batch.num_columns()
        );

        Ok(DataFrame {
            record_batch,
            // metadata,
            // data_aux,
        })
    }

    pub fn rename_columns(&self, rename_map: &HashMap<String, String>) -> Result<DataFrame> {
        if rename_map.is_empty() {
            return Ok(self.clone());
        }

        let schema = self.record_batch.schema();
        let fields = schema
            .fields()
            .iter()
            .map(|field| {
                rename_map
                    .get(field.name())
                    .map(|name| Field::new(name, field.data_type().clone(), field.is_nullable()))
                    .unwrap_or_else(|| field.as_ref().clone())
            })
            .collect::<Vec<Field>>();

        let renamed_schema = Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone()));
        let renamed_batch =
            RecordBatch::try_new(renamed_schema, self.record_batch.columns().to_vec())?;
        Ok(DataFrame {
            record_batch: renamed_batch,
        })
    }

    pub fn record_batch(&self) -> RecordBatch {
        self.record_batch.clone()
    }
}

impl std::fmt::Display for DataFrame {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "{}",
            pretty_format_batches(&[self.record_batch.clone()])
                .expect("Error pretty formatting dataframe.")
        )?;
        Ok(())
    }
}
