use crate::{error::Error, xlsx::XlsxReader};
use polars::prelude::*;

impl From<PolarsError> for Error {
    fn from(value: PolarsError) -> Self {
        Error::PolarsError(value.to_string())
    }
}

impl XlsxReader {
    pub fn to_lazyframe(mut self) -> Result<LazyFrame, Error> {
        self.finish()?;
        let (num_cols, num_rows) = self.get_worksheet_dimensions();

        let mut columns: Vec<Column> = (0..num_cols)
            .map(|i| Column::new(format!("column_{i}").into(), Vec::<String>::new()))
            .collect();

        const CHUNK_SIZE: usize = 1000;
        let mut chunk: Vec<Vec<String>> = Vec::with_capacity(CHUNK_SIZE);

        for current_row in 0..num_rows {
            let row = self.get_row(current_row);

            chunk.push(row);

            if chunk.len() >= CHUNK_SIZE {
                append_chunk(&mut columns, &chunk);
                chunk.clear();
            }
        }

        if !chunk.is_empty() {
            append_chunk(&mut columns, &chunk);
        }

        let df =
            DataFrame::new(num_rows as usize, columns).map_err(|e| Error::from(e.to_string()))?;

        Ok(df.lazy())
    }
}

fn append_chunk(columns: &mut [Column], chunk: &[Vec<String>]) {
    for (col_idx, column) in columns.iter_mut().enumerate() {
        let chunk_data: Vec<String> = chunk.iter().map(|row| row[col_idx].clone()).collect();

        let chunk_column = Column::new("temp".into(), chunk_data);
        column.append(&chunk_column).unwrap();
    }
}
