ggsql 0.4.1

//! DataFrame to JSON conversion utilities for Vega-Lite writer
//!
//! This module handles converting Arrow DataFrames to Vega-Lite JSON data values,
//! including temporal type handling and binned data transformations.

use crate::array_util::*;
use crate::plot::scale::ScaleTypeKind;

/// Column name for row index (used to preserve data order in Vega-Lite)
pub(super) const ROW_INDEX_COLUMN: &str = "__ggsql_row_index__";
// ArrayElement is used for temporal parsing
#[allow(unused_imports)]
use crate::plot::ArrayElement;
use crate::plot::ParameterValue;
use crate::{naming, AestheticValue, DataFrame, GgsqlError, Plot, Result};
use arrow::array::{Array, ArrayRef};
use arrow::datatypes::{DataType, TimeUnit};
use serde_json::{json, Map, Value};
use std::collections::HashMap;

/// Temporal type for binned date/datetime/time columns
#[derive(Debug, Clone, Copy, PartialEq)]
pub(super) enum TemporalType {
    Date,
    DateTime,
    Time,
}

/// Convert DataFrame to Vega-Lite data values (array of objects)
pub(super) fn dataframe_to_values(df: &DataFrame) -> Result<Vec<Value>> {
    let mut values = Vec::new();
    let height = df.height();
    let column_names = df.get_column_names();
    let columns = df.get_columns();

    for row_idx in 0..height {
        let mut row_obj = Map::new();

        for (col_idx, col_name) in column_names.iter().enumerate() {
            let column = columns.get(col_idx).ok_or_else(|| {
                GgsqlError::WriterError(format!("Failed to get column {}", col_name))
            })?;

            // Get value from array and convert to JSON Value
            let value = series_value_at(column, row_idx)?;
            row_obj.insert(col_name.to_string(), value);
        }

        values.push(Value::Object(row_obj));
    }

    Ok(values)
}

/// Get a single value from an arrow array at a given index as JSON Value
pub(super) fn series_value_at(array: &ArrayRef, idx: usize) -> Result<Value> {
    if array.is_null(idx) {
        return Ok(Value::Null);
    }

    match array.data_type() {
        DataType::Int8 => Ok(json!(as_i8(array)?.value(idx))),
        DataType::Int16 => Ok(json!(as_i16(array)?.value(idx))),
        DataType::Int32 => Ok(json!(as_i32(array)?.value(idx))),
        DataType::Int64 => Ok(json!(as_i64(array)?.value(idx))),
        DataType::UInt8 => Ok(json!(as_u8(array)?.value(idx))),
        DataType::UInt16 => Ok(json!(as_u16(array)?.value(idx))),
        DataType::UInt32 => Ok(json!(as_u32(array)?.value(idx))),
        DataType::UInt64 => Ok(json!(as_u64(array)?.value(idx))),
        DataType::Float32 => Ok(json!(as_f32(array)?.value(idx))),
        DataType::Float64 => Ok(json!(as_f64(array)?.value(idx))),
        DataType::Boolean => Ok(json!(as_bool(array)?.value(idx))),
        DataType::Utf8 => {
            // Keep strings as strings (don't parse to numbers)
            Ok(json!(as_str(array)?.value(idx)))
        }
        DataType::Date32 => {
            // Convert days since epoch to ISO date string: "YYYY-MM-DD"
            let days = as_date32(array)?.value(idx);
            let unix_epoch = chrono::NaiveDate::from_ymd_opt(1970, 1, 1).unwrap();
            let date = unix_epoch + chrono::Duration::days(days as i64);
            Ok(json!(date.format("%Y-%m-%d").to_string()))
        }
        DataType::Timestamp(time_unit, _) => {
            // Convert timestamp to ISO datetime: "YYYY-MM-DDTHH:MM:SS.sssZ"
            let timestamp = as_timestamp_us(array).map(|a| a.value(idx)).or_else(|_| {
                // Try casting to microsecond timestamp first
                let cast = cast_array(array, &DataType::Timestamp(TimeUnit::Microsecond, None))?;
                Ok(as_timestamp_us(&cast)?.value(idx))
            })?;
            // timestamp is in microseconds for TimestampMicrosecondArray
            let micros = match time_unit {
                TimeUnit::Microsecond => timestamp,
                TimeUnit::Millisecond => timestamp * 1_000,
                TimeUnit::Nanosecond => timestamp / 1_000,
                TimeUnit::Second => timestamp * 1_000_000,
            };
            let secs = micros / 1_000_000;
            let nsecs = ((micros % 1_000_000) * 1000) as u32;
            let dt = chrono::DateTime::<chrono::Utc>::from_timestamp(secs, nsecs)
                .unwrap_or_else(|| chrono::DateTime::<chrono::Utc>::from_timestamp(0, 0).unwrap());
            Ok(json!(dt.format("%Y-%m-%dT%H:%M:%S%.3fZ").to_string()))
        }
        DataType::Time64(_) => {
            // Convert nanoseconds since midnight to ISO time: "HH:MM:SS.sss"
            let nanos = as_time64_ns(array)?.value(idx);
            let hours = nanos / 3_600_000_000_000;
            let minutes = (nanos % 3_600_000_000_000) / 60_000_000_000;
            let seconds = (nanos % 60_000_000_000) / 1_000_000_000;
            let millis = (nanos % 1_000_000_000) / 1_000_000;
            Ok(json!(format!(
                "{:02}:{:02}:{:02}.{:03}",
                hours, minutes, seconds, millis
            )))
        }
        _ => {
            // Fallback: convert to string
            Ok(json!(value_to_string(array, idx)))
        }
    }
}

/// Given a bin center value and breaks array, return (bin_start, bin_end).
/// Find the bin interval that contains the given value.
///
/// The breaks array contains bin edges [e0, e1, e2, ...].
/// Returns the (lower, upper) edges of the bin containing the value.
/// Uses half-open intervals [lower, upper) except for the last bin which is [lower, upper].
pub(super) fn find_bin_for_value(value: f64, breaks: &[f64]) -> Option<(f64, f64)> {
    let n = breaks.len();
    if n < 2 {
        return None;
    }

    for i in 0..n - 1 {
        let lower = breaks[i];
        let upper = breaks[i + 1];
        let is_last_bin = i == n - 2;

        // Use [lower, upper) for all bins except the last which uses [lower, upper]
        let in_bin = if is_last_bin {
            value >= lower && value <= upper
        } else {
            value >= lower && value < upper
        };

        if in_bin {
            return Some((lower, upper));
        }
    }
    None
}

/// Convert DataFrame to Vega-Lite data values with bin columns.
///
/// For columns with binned scales, this replaces the center value with bin_start
/// and adds a corresponding bin_end column.
pub(super) fn dataframe_to_values_with_bins(
    df: &DataFrame,
    binned_columns: &HashMap<String, Vec<f64>>,
) -> Result<Vec<Value>> {
    let mut values = Vec::new();
    let height = df.height();
    let column_names = df.get_column_names();
    let columns = df.get_columns();

    for row_idx in 0..height {
        let mut row_obj = Map::new();

        for (col_idx, col_name) in column_names.iter().enumerate() {
            let column = columns.get(col_idx).ok_or_else(|| {
                GgsqlError::WriterError(format!("Failed to get column {}", col_name))
            })?;

            // Get value from array and convert to JSON Value
            let value = series_value_at(column, row_idx)?;

            // Check if this column has binned data
            let col_name_str = col_name.to_string();
            if let Some(breaks) = binned_columns.get(&col_name_str) {
                // Check if this is a temporal string (date/datetime/time)
                let temporal_info = value.as_str().and_then(parse_temporal_string);

                // Get value as f64 - works for numeric columns or parsed temporal strings
                let numeric_value = value.as_f64().or_else(|| temporal_info.map(|(val, _)| val));

                if let Some(val) = numeric_value {
                    if let Some((start, end)) = find_bin_for_value(val, breaks) {
                        // Replace value with bin_start, preserving original value type
                        if let Some((_, temporal_type)) = temporal_info {
                            // Temporal column - format bin edges as ISO strings
                            let start_str = format_temporal(start, temporal_type);
                            let end_str = format_temporal(end, temporal_type);
                            row_obj.insert(col_name_str.clone(), json!(start_str));
                            row_obj.insert(naming::bin_end_column(&col_name_str), json!(end_str));
                        } else {
                            // Numeric column - use raw values
                            row_obj.insert(col_name_str.clone(), json!(start));
                            row_obj.insert(naming::bin_end_column(&col_name_str), json!(end));
                        }
                        continue;
                    }
                }
            }

            // Not binned or couldn't resolve edges - use original value
            row_obj.insert(col_name.to_string(), value);
        }

        values.push(Value::Object(row_obj));
    }

    Ok(values)
}

/// Detect the temporal type of a string value.
/// Returns the parsed numeric value and the type.
///
/// Uses ArrayElement's parsing methods which support comprehensive format variations.
pub(super) fn parse_temporal_string(s: &str) -> Option<(f64, TemporalType)> {
    // Try date first (YYYY-MM-DD) - must check before datetime since dates are shorter
    if let Some(ArrayElement::Date(days)) = ArrayElement::from_date_string(s) {
        return Some((days as f64, TemporalType::Date));
    }
    // Try datetime (various ISO formats with/without timezone)
    if let Some(ArrayElement::DateTime(micros)) = ArrayElement::from_datetime_string(s) {
        return Some((micros as f64, TemporalType::DateTime));
    }
    // Try time (HH:MM:SS[.sss])
    if let Some(ArrayElement::Time(nanos)) = ArrayElement::from_time_string(s) {
        return Some((nanos as f64, TemporalType::Time));
    }
    None
}

/// Format a numeric temporal value back to ISO string.
pub(super) fn format_temporal(value: f64, temporal_type: TemporalType) -> String {
    match temporal_type {
        TemporalType::Date => ArrayElement::date_to_iso(value as i32),
        TemporalType::DateTime => ArrayElement::datetime_to_iso(value as i64),
        TemporalType::Time => ArrayElement::time_to_iso(value as i64),
    }
}

/// Collect binned column information from spec.
///
/// Returns a map of column name -> breaks array for all columns with binned scales.
/// The column name uses the aesthetic-prefixed format (e.g., `__ggsql_aes_x__`) since
/// that's what appears in the DataFrame after query execution.
///
/// Only x and y aesthetics are collected since only those have x2/y2 counterparts
/// in Vega-Lite for representing bin ranges.
pub(super) fn collect_binned_columns(spec: &Plot) -> HashMap<String, Vec<f64>> {
    let mut binned_columns: HashMap<String, Vec<f64>> = HashMap::new();
    let aesthetic_ctx = spec.get_aesthetic_context();

    for scale in &spec.scales {
        // Only x and y aesthetics support bin ranges (x2/y2) in Vega-Lite
        if !aesthetic_ctx.is_primary_internal(&scale.aesthetic) {
            continue;
        }

        // Check if this is a binned scale
        let is_binned = scale
            .scale_type
            .as_ref()
            .map(|st| st.scale_type_kind() == ScaleTypeKind::Binned)
            .unwrap_or(false);

        if !is_binned {
            continue;
        }

        // Get breaks array from scale properties
        if let Some(ParameterValue::Array(breaks)) = scale.properties.get("breaks") {
            let break_values: Vec<f64> = breaks.iter().filter_map(|e| e.to_f64()).collect();

            if break_values.len() >= 2 {
                // Insert the aesthetic column name (what's in the DataFrame after execution)
                let aes_col_name = naming::aesthetic_column(&scale.aesthetic);
                binned_columns.insert(aes_col_name, break_values.clone());

                // Also insert mappings for original column names (for unit tests and
                // cases where the full pipeline isn't used)
                for layer in &spec.layers {
                    if let Some(AestheticValue::Column { name: col, .. }) =
                        layer.mappings.aesthetics.get(&scale.aesthetic)
                    {
                        binned_columns.insert(col.clone(), break_values.clone());
                    }
                }
            }
        }
    }

    binned_columns
}

/// Check if an aesthetic has a binned scale in the spec.
pub(super) fn is_binned_aesthetic(aesthetic: &str, spec: &Plot) -> bool {
    let aesthetic_ctx = spec.get_aesthetic_context();
    let primary = aesthetic_ctx
        .primary_internal_position(aesthetic)
        .unwrap_or(aesthetic);
    spec.find_scale(primary)
        .and_then(|s| s.scale_type.as_ref())
        .map(|st| st.scale_type_kind() == ScaleTypeKind::Binned)
        .unwrap_or(false)
}

/// Unify multiple datasets into a single dataset with source identification.
///
/// This concatenates all layer datasets into one unified dataset, adding a
/// `__ggsql_source__` field to each row that identifies which layer's data
/// the row belongs to. Each layer then uses a Vega-Lite transform filter
/// to select its data.
///
/// # Arguments
/// * `datasets` - Map of dataset key to Vega-Lite JSON values array
///
/// # Returns
/// Unified array of all rows with source identification
pub(super) fn unify_datasets(datasets: &Map<String, Value>) -> Result<Vec<Value>> {
    // 1. Collect all unique column names across all datasets
    let mut all_columns: std::collections::HashSet<String> = std::collections::HashSet::new();
    for (_key, values) in datasets {
        if let Some(arr) = values.as_array() {
            for row in arr {
                if let Some(obj) = row.as_object() {
                    for col_name in obj.keys() {
                        all_columns.insert(col_name.clone());
                    }
                }
            }
        }
    }

    // 2. For each dataset, for each row:
    //    - Include all columns (null for missing)
    //    - Add __ggsql_source__ field with dataset key
    //    - Add __ggsql_row_index__ field for order preservation
    let mut unified = Vec::new();
    let mut row_index: usize = 0;
    for (key, values) in datasets {
        if let Some(arr) = values.as_array() {
            for row in arr {
                if let Some(obj) = row.as_object() {
                    let mut new_row = Map::new();

                    // Include all columns from union schema (null for missing)
                    for col_name in &all_columns {
                        let value = obj.get(col_name).cloned().unwrap_or(Value::Null);
                        new_row.insert(col_name.clone(), value);
                    }

                    // Add source identifier
                    new_row.insert(naming::SOURCE_COLUMN.to_string(), json!(key));

                    // Add row index for order preservation (used by line/path/polygon geoms)
                    new_row.insert(ROW_INDEX_COLUMN.to_string(), json!(row_index));
                    row_index += 1;

                    unified.push(Value::Object(new_row));
                }
            }
        }
    }

    Ok(unified)
}