query-forge 0.9.0

use anyhow::{Context, Result, bail};
use serde_json::Value as JsonValue;
use std::{
    collections::{HashMap, hash_map::Entry},
    io::{BufRead, BufReader},
};

use crate::{QueryValue, TypeInferenceOptions};

use super::SheetData;

pub(super) fn load_jsonl_sheet(
    jsonl_path: &std::path::Path,
    requested_sheet: Option<&str>,
    inference_options: &TypeInferenceOptions,
) -> Result<SheetData> {
    if let Some(selector) = requested_sheet {
        bail!(
            "JSONL input {} does not support selector '{selector}'. Remove ':{selector}' from --input for JSONL files.",
            jsonl_path.display()
        );
    }

    let file = std::fs::File::open(jsonl_path)
        .with_context(|| format!("failed to read {}", jsonl_path.display()))?;
    let reader = BufReader::new(file);

    let mut columns = Vec::<String>::new();
    let mut column_indices = HashMap::<String, usize>::new();
    let mut saw_row = false;

    for (line_index, raw_line) in reader.lines().enumerate() {
        let raw_line = raw_line.with_context(|| {
            format!(
                "failed to read JSONL line {} from {}",
                line_index + 1,
                jsonl_path.display()
            )
        })?;
        let object = parse_jsonl_object(jsonl_path, line_index + 1, raw_line)?;
        let Some(object) = object else {
            continue;
        };
        saw_row = true;

        for key in object.keys().cloned() {
            if let Entry::Vacant(entry) = column_indices.entry(key) {
                let index = columns.len();
                columns.push(entry.key().clone());
                entry.insert(index);
            }
        }
    }

    if !saw_row {
        bail!("JSONL input {} is empty", jsonl_path.display());
    }

    let file = std::fs::File::open(jsonl_path)
        .with_context(|| format!("failed to read {}", jsonl_path.display()))?;
    let reader = BufReader::new(file);
    let mut rows = Vec::<Vec<QueryValue>>::new();

    // Discovering the full schema first lets the data pass allocate each row once.
    for (line_index, raw_line) in reader.lines().enumerate() {
        let raw_line = raw_line.with_context(|| {
            format!(
                "failed to read JSONL line {} from {}",
                line_index + 1,
                jsonl_path.display()
            )
        })?;
        let object = parse_jsonl_object(jsonl_path, line_index + 1, raw_line)?;
        let Some(object) = object else {
            continue;
        };

        let mut row = vec![QueryValue::Null; columns.len()];
        for (key, value) in object {
            if let Some(index) = column_indices.get(&key).copied() {
                row[index] = json_to_query_value(value, inference_options);
            }
        }
        rows.push(row);
    }

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

pub(super) fn load_jsonl_sheet_from_str(
    jsonl_path: &std::path::Path,
    content: &str,
    requested_sheet: Option<&str>,
    inference_options: &TypeInferenceOptions,
) -> Result<SheetData> {
    if let Some(selector) = requested_sheet {
        bail!(
            "JSONL input {} does not support selector '{selector}'. Remove ':{selector}' from --input for JSONL files.",
            jsonl_path.display()
        );
    }

    let mut objects = Vec::new();
    let mut columns = Vec::<String>::new();
    let mut column_indices = HashMap::<String, usize>::new();

    for (line_index, raw_line) in content.lines().enumerate() {
        let object = parse_jsonl_object(jsonl_path, line_index + 1, raw_line.to_owned())?;
        let Some(object) = object else {
            continue;
        };

        for key in object.keys().cloned() {
            if let Entry::Vacant(entry) = column_indices.entry(key) {
                let index = columns.len();
                columns.push(entry.key().clone());
                entry.insert(index);
            }
        }

        objects.push(object);
    }

    if objects.is_empty() {
        bail!("JSONL input {} is empty", jsonl_path.display());
    }

    let rows = objects
        .into_iter()
        .map(|object| {
            let mut row = vec![QueryValue::Null; columns.len()];
            for (key, value) in object {
                if let Some(index) = column_indices.get(&key).copied() {
                    row[index] = json_to_query_value(value, inference_options);
                }
            }
            row
        })
        .collect();

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

fn parse_jsonl_object(
    jsonl_path: &std::path::Path,
    line_number: usize,
    raw_line: String,
) -> Result<Option<serde_json::Map<String, JsonValue>>> {
    let line = raw_line.trim();
    if line.is_empty() {
        return Ok(None);
    }

    let json_value = serde_json::from_str::<JsonValue>(line).with_context(|| {
        format!(
            "failed to parse JSONL line {} in {}",
            line_number,
            jsonl_path.display()
        )
    })?;

    let JsonValue::Object(object) = json_value else {
        bail!(
            "JSONL line {} in {} is not an object",
            line_number,
            jsonl_path.display()
        );
    };

    Ok(Some(object))
}

fn json_to_query_value(value: JsonValue, inference_options: &TypeInferenceOptions) -> QueryValue {
    use super::apply_inference_overrides;
    use super::parse_scalar_value;

    match value {
        JsonValue::Null => QueryValue::Null,
        JsonValue::Bool(flag) => {
            apply_inference_overrides(QueryValue::Integer(i64::from(flag)), inference_options)
        }
        JsonValue::Number(number) => {
            if let Some(integer) = number.as_i64() {
                apply_inference_overrides(QueryValue::Integer(integer), inference_options)
            } else if let Some(real) = number.as_f64() {
                apply_inference_overrides(QueryValue::Real(real), inference_options)
            } else {
                QueryValue::Text(number.to_string())
            }
        }
        JsonValue::String(text) => parse_scalar_value(&text, inference_options),
        JsonValue::Array(_) | JsonValue::Object(_) => QueryValue::Text(value.to_string()),
    }
}