query-forge 0.9.0

use anyhow::{Context, Result, anyhow, bail};
use serde_json::Value as JsonValue;

use crate::{ExtractionOptions, JsonMode, QueryValue, TypeInferenceOptions};

use super::{SheetData, apply_inference_overrides, parse_scalar_value, rows_maps_to_sheet_data};

pub(super) fn load_json_sheet(
    json_path: &std::path::Path,
    requested_sheet: Option<&str>,
    inference_options: &TypeInferenceOptions,
    extraction_options: &ExtractionOptions,
) -> Result<SheetData> {
    let content = std::fs::read_to_string(json_path)
        .with_context(|| format!("failed to read {}", json_path.display()))?;
    load_json_sheet_from_str(
        json_path,
        &content,
        requested_sheet,
        inference_options,
        extraction_options,
    )
}

pub(super) fn load_json_sheet_from_str(
    json_path: &std::path::Path,
    content: &str,
    requested_sheet: Option<&str>,
    inference_options: &TypeInferenceOptions,
    extraction_options: &ExtractionOptions,
) -> Result<SheetData> {
    let root = serde_json::from_str::<JsonValue>(&content)
        .with_context(|| format!("failed to parse JSON {}", json_path.display()))?;

    let scope = if let Some(sheet_key) = requested_sheet {
        let JsonValue::Object(mut object) = root else {
            bail!(
                "JSON input {} uses selector '{sheet_key}', but key selection requires a top-level object. Remove the selector or provide a JSON object at the root.",
                json_path.display()
            );
        };

        let mut available_keys = object.keys().cloned().collect::<Vec<_>>();
        available_keys.sort();

        object.remove(sheet_key).ok_or_else(|| {
            let available_suffix = if available_keys.is_empty() {
                " The top-level object has no keys.".to_owned()
            } else {
                format!(" Available keys: {}.", available_keys.join(", "))
            };
            anyhow!(
                "JSON key '{sheet_key}' not found in {}.{available_suffix}",
                json_path.display()
            )
        })?
    } else {
        root
    };

    let rows_maps = match extraction_options.json_mode {
        JsonMode::Array => json_scope_to_rows(scope, inference_options),
        JsonMode::Object => json_scope_to_rows_object(scope, inference_options),
        JsonMode::Flatten => json_scope_to_rows_flatten(scope, inference_options),
    };
    if rows_maps.is_empty() {
        if let Some(sheet_key) = requested_sheet {
            bail!(
                "JSON selector '{sheet_key}' in {} resolved to an empty table (no rows)",
                json_path.display()
            );
        }
        bail!("JSON input {} is empty", json_path.display());
    }

    Ok(rows_maps_to_sheet_data(
        rows_maps,
        requested_sheet.unwrap_or("json").to_owned(),
    ))
}

fn json_scope_to_rows(
    scope: JsonValue,
    inference_options: &TypeInferenceOptions,
) -> Vec<Vec<(String, QueryValue)>> {
    match scope {
        JsonValue::Array(items) => items
            .into_iter()
            .map(|item| json_item_to_row(item, inference_options))
            .collect::<Vec<_>>(),
        JsonValue::Object(object) => vec![
            object
                .into_iter()
                .map(|(key, value)| (key, json_to_query_value(value, inference_options)))
                .collect(),
        ],
        scalar => vec![vec![(
            "value".to_owned(),
            json_to_query_value(scalar, inference_options),
        )]],
    }
}

fn json_item_to_row(
    item: JsonValue,
    inference_options: &TypeInferenceOptions,
) -> Vec<(String, QueryValue)> {
    match item {
        JsonValue::Object(object) => object
            .into_iter()
            .map(|(key, value)| (key, json_to_query_value(value, inference_options)))
            .collect(),
        scalar => vec![(
            "value".to_owned(),
            json_to_query_value(scalar, inference_options),
        )],
    }
}

fn json_to_query_value(value: JsonValue, inference_options: &TypeInferenceOptions) -> QueryValue {
    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()),
    }
}

fn json_scope_to_rows_object(
    scope: JsonValue,
    inference_options: &TypeInferenceOptions,
) -> Vec<Vec<(String, QueryValue)>> {
    match scope {
        JsonValue::Object(object) => object
            .into_iter()
            .map(|(key, value)| {
                vec![
                    ("key".to_owned(), QueryValue::Text(key)),
                    (
                        "value".to_owned(),
                        json_to_query_value(value, inference_options),
                    ),
                ]
            })
            .collect(),
        JsonValue::Array(items) => items
            .into_iter()
            .flat_map(|item| json_scope_to_rows_object(item, inference_options))
            .collect(),
        scalar => vec![vec![
            ("key".to_owned(), QueryValue::Text("value".to_owned())),
            (
                "value".to_owned(),
                json_to_query_value(scalar, inference_options),
            ),
        ]],
    }
}

fn json_scope_to_rows_flatten(
    scope: JsonValue,
    inference_options: &TypeInferenceOptions,
) -> Vec<Vec<(String, QueryValue)>> {
    match scope {
        JsonValue::Array(items) => items
            .into_iter()
            .map(|item| {
                let mut row = Vec::new();
                flatten_json_value("", item, &mut row, inference_options);
                row
            })
            .filter(|row| !row.is_empty())
            .collect(),
        _ => {
            let mut row = Vec::new();
            flatten_json_value("", scope, &mut row, inference_options);
            if row.is_empty() { vec![] } else { vec![row] }
        }
    }
}

fn flatten_json_value(
    prefix: &str,
    value: JsonValue,
    result: &mut Vec<(String, QueryValue)>,
    inference_options: &TypeInferenceOptions,
) {
    match value {
        JsonValue::Object(obj) => {
            for (k, v) in obj {
                let key = if prefix.is_empty() {
                    k
                } else {
                    format!("{prefix}.{k}")
                };
                flatten_json_value(&key, v, result, inference_options);
            }
        }
        JsonValue::Array(arr) => {
            for (i, v) in arr.into_iter().enumerate() {
                let key = if prefix.is_empty() {
                    i.to_string()
                } else {
                    format!("{prefix}.{i}")
                };
                flatten_json_value(&key, v, result, inference_options);
            }
        }
        _ => {
            let key = if prefix.is_empty() {
                "value".to_owned()
            } else {
                prefix.to_owned()
            };
            result.push((key, json_to_query_value(value, inference_options)));
        }
    }
}