query-forge 0.9.0

use anyhow::{Context, Result, anyhow, bail};
use calamine::{Data, Reader, open_workbook_auto};

use crate::{QueryValue, TypeInferenceOptions};

use super::{SheetData, apply_inference_overrides};

pub(super) fn load_xlsx_sheet(
    workbook_path: &std::path::Path,
    requested_sheet: Option<&str>,
    inference_options: &TypeInferenceOptions,
    has_headers: bool,
) -> Result<SheetData> {
    let mut workbook = open_workbook_auto(workbook_path)
        .with_context(|| format!("failed to open {}", workbook_path.display()))?;

    let sheet_name = match requested_sheet {
        Some(name) => name.to_owned(),
        None => workbook
            .sheet_names()
            .first()
            .cloned()
            .ok_or_else(|| anyhow!("workbook does not contain any sheets"))?,
    };

    let range = workbook
        .worksheet_range(&sheet_name)
        .with_context(|| format!("failed to read sheet {sheet_name}"))?;

    if range.width() == 0 {
        bail!("sheet {sheet_name} is empty");
    }

    let mut rows = range.rows();
    let columns = if has_headers {
        let header_row = rows
            .next()
            .ok_or_else(|| anyhow!("sheet {sheet_name} is empty"))?;
        normalize_headers(header_row, range.width())
    } else {
        (1..=range.width())
            .map(|index| format!("column{index}"))
            .collect()
    };

    let data_rows = rows
        .map(|row| {
            (0..columns.len())
                .map(|index| convert_cell(row.get(index).unwrap_or(&Data::Empty)))
                .map(|value| apply_inference_overrides(value, inference_options))
                .collect::<Vec<_>>()
        })
        .filter(|row| row.iter().any(|value| !matches!(value, QueryValue::Null)))
        .collect();

    Ok(SheetData {
        original_name: sheet_name,
        columns,
        rows: data_rows,
    })
}

pub(super) fn normalize_headers(header_row: &[Data], width: usize) -> Vec<String> {
    let mut seen = std::collections::HashMap::new();

    (0..width)
        .map(|index| {
            let base = header_row
                .get(index)
                .map(cell_to_string)
                .filter(|value| !value.trim().is_empty())
                .unwrap_or_else(|| format!("column{}", index + 1));

            let count = seen
                .entry(base.clone())
                .and_modify(|count| *count += 1)
                .or_insert(1usize);

            if *count == 1 {
                base
            } else {
                format!("{base}_{count}")
            }
        })
        .collect()
}

fn convert_cell(cell: &Data) -> QueryValue {
    match cell {
        Data::Empty => QueryValue::Null,
        Data::Int(value) => QueryValue::Integer(*value),
        Data::Float(value) => QueryValue::Real(*value),
        Data::String(value) => QueryValue::Text(value.clone()),
        Data::Bool(value) => QueryValue::Integer(i64::from(*value)),
        Data::DateTime(value) => QueryValue::Text(value.to_string()),
        Data::DateTimeIso(value) => QueryValue::Text(value.clone()),
        Data::DurationIso(value) => QueryValue::Text(value.clone()),
        Data::Error(value) => QueryValue::Text(value.to_string()),
    }
}

fn cell_to_string(cell: &Data) -> String {
    match convert_cell(cell) {
        QueryValue::Null => String::new(),
        QueryValue::Integer(value) => value.to_string(),
        QueryValue::Real(value) => value.to_string(),
        QueryValue::Text(value) => value,
    }
}

#[cfg(test)]
mod tests {
    use calamine::Data;

    use super::normalize_headers;

    #[test]
    fn normalizes_duplicate_and_blank_headers() {
        let headers = vec![
            Data::String("name".into()),
            Data::String(String::new()),
            Data::String("name".into()),
        ];

        assert_eq!(
            normalize_headers(&headers, 3),
            vec!["name", "column2", "name_2"]
        );
    }
}