query-forge 0.9.0

use std::io::Read;

use anyhow::{Context, Result, bail};

use crate::{QueryValue, TypeInferenceOptions};

use super::{SheetData, normalize_text_headers, parse_scalar_value};

pub(super) fn load_csv_sheet(
    csv_path: &std::path::Path,
    requested_sheet: Option<&str>,
    inference_options: &TypeInferenceOptions,
    has_headers: bool,
) -> Result<SheetData> {
    let reader = csv::ReaderBuilder::new()
        .has_headers(has_headers)
        .from_path(csv_path)
        .with_context(|| format!("failed to open {}", csv_path.display()))?;

    load_csv_sheet_from_reader(
        csv_path,
        reader,
        requested_sheet,
        inference_options,
        has_headers,
    )
}

pub(super) fn load_csv_sheet_from_str(
    csv_path: &std::path::Path,
    content: &str,
    requested_sheet: Option<&str>,
    inference_options: &TypeInferenceOptions,
    has_headers: bool,
) -> Result<SheetData> {
    let reader = csv::ReaderBuilder::new()
        .has_headers(has_headers)
        .from_reader(content.as_bytes());

    load_csv_sheet_from_reader(
        csv_path,
        reader,
        requested_sheet,
        inference_options,
        has_headers,
    )
}

fn load_csv_sheet_from_reader<R: Read>(
    csv_path: &std::path::Path,
    mut reader: csv::Reader<R>,
    requested_sheet: Option<&str>,
    inference_options: &TypeInferenceOptions,
    has_headers: bool,
) -> Result<SheetData> {
    if let Some(selector) = requested_sheet {
        bail!(
            "CSV input {} does not support selector '{selector}'. Remove ':{selector}' from --input for CSV files.",
            csv_path.display()
        );
    }

    let mut columns = if has_headers {
        let headers = reader
            .headers()
            .with_context(|| format!("failed to read headers from {}", csv_path.display()))?
            .iter()
            .map(str::to_owned)
            .collect::<Vec<_>>();
        normalize_text_headers(&headers)
    } else {
        Vec::new()
    };

    let mut inferred_width = columns.len();
    let mut rows = Vec::new();
    for record in reader.records() {
        let record = record
            .with_context(|| format!("failed to read CSV record from {}", csv_path.display()))?;
        let target_width = if has_headers {
            columns.len()
        } else {
            inferred_width = inferred_width.max(record.len());
            record.len()
        };
        let row = (0..target_width)
            .map(|index| {
                let value = record.get(index).unwrap_or("");
                parse_scalar_value(value, inference_options)
            })
            .collect::<Vec<_>>();

        if row.iter().any(|value| !matches!(value, QueryValue::Null)) {
            rows.push(row);
        }
    }

    if !has_headers {
        columns = (1..=inferred_width)
            .map(|index| format!("column{index}"))
            .collect();
        for row in &mut rows {
            row.resize(inferred_width, QueryValue::Null);
        }
    }

    if columns.is_empty() {
        bail!("CSV input {} is empty", csv_path.display());
    }

    Ok(SheetData {
        original_name: "csv".to_owned(),
        columns,
        rows,
    })
}