tellaro-query-language 1.3.8

//! Stats evaluator for TQL aggregation queries.
//!
//! Provides statistical aggregation functions with grouping support.

use crate::error::{Result, TqlError};
use crate::field_accessor;
use serde_json::{json, Value as JsonValue};
use std::collections::{BTreeMap, HashMap, HashSet};

/// Stats evaluator for aggregation queries
pub struct StatsEvaluator {
    _max_depth: usize,
}

/// Aggregation specification
#[derive(Debug, Clone)]
pub struct AggregationSpec {
    pub function: String,
    pub field: String,
    pub alias: Option<String>,
    pub params: HashMap<String, JsonValue>,
}

/// Stats query specification
#[derive(Debug, Clone)]
pub struct StatsQuery {
    pub aggregations: Vec<AggregationSpec>,
    pub group_by: Vec<String>,
}

impl Default for StatsEvaluator {
    fn default() -> Self {
        Self::new()
    }
}

impl StatsEvaluator {
    /// Create a new stats evaluator
    pub fn new() -> Self {
        Self { _max_depth: 100 }
    }

    /// Evaluate stats query against records
    ///
    /// Returns aggregated results in a structured format
    pub fn evaluate_stats(&self, records: &[JsonValue], query: &StatsQuery) -> Result<JsonValue> {
        if query.group_by.is_empty() {
            self.simple_aggregation(records, &query.aggregations)
        } else {
            self.grouped_aggregation(records, &query.aggregations, &query.group_by)
        }
    }

    /// Perform aggregation without grouping
    fn simple_aggregation(
        &self,
        records: &[JsonValue],
        aggregations: &[AggregationSpec],
    ) -> Result<JsonValue> {
        if aggregations.len() == 1 {
            // Single aggregation
            let agg = &aggregations[0];
            let value = self.calculate_aggregation(records, agg)?;

            Ok(json!({
                "type": "simple_aggregation",
                "function": agg.function,
                "field": agg.field,
                "alias": agg.alias,
                "value": value
            }))
        } else {
            // Multiple aggregations
            let mut results = HashMap::new();
            for agg in aggregations {
                let value = self.calculate_aggregation(records, agg)?;
                let key = agg
                    .alias
                    .clone()
                    .unwrap_or_else(|| format!("{}_{}", agg.function, agg.field));
                results.insert(key, value);
            }

            Ok(json!({
                "type": "multiple_aggregations",
                "results": results
            }))
        }
    }

    /// Perform aggregation with grouping
    fn grouped_aggregation(
        &self,
        records: &[JsonValue],
        aggregations: &[AggregationSpec],
        group_by_fields: &[String],
    ) -> Result<JsonValue> {
        // Group records by field values. BTreeMap ensures deterministic ordering of
        // results regardless of HashMap iteration order.
        let mut groups: BTreeMap<Vec<String>, Vec<&JsonValue>> = BTreeMap::new();

        for record in records {
            let mut key_parts = Vec::new();
            for field in group_by_fields {
                let value = match field_accessor::get_field(record, field) {
                    Ok(Some(v)) => v.clone(),
                    Ok(None) | Err(_) => JsonValue::Null,
                };
                let key_str = match value {
                    JsonValue::String(s) => s,
                    JsonValue::Number(n) => n.to_string(),
                    JsonValue::Bool(b) => b.to_string(),
                    JsonValue::Null => "null".to_string(),
                    _ => serde_json::to_string(&value).unwrap_or_default(),
                };
                key_parts.push(key_str);
            }

            groups.entry(key_parts).or_default().push(record);
        }

        // Calculate aggregations for each group
        let mut results = Vec::new();
        for (key_parts, group_records) in groups {
            let mut group_result: HashMap<String, JsonValue> = HashMap::new();

            // Add group key
            let mut key_map = HashMap::new();
            for (i, field) in group_by_fields.iter().enumerate() {
                key_map.insert(field.clone(), json!(key_parts[i]));
            }
            group_result.insert("key".to_string(), json!(key_map));
            group_result.insert("doc_count".to_string(), json!(group_records.len()));

            if aggregations.len() == 1 {
                // Single aggregation
                let agg = &aggregations[0];
                let owned_records: Vec<JsonValue> =
                    group_records.iter().map(|&r| r.clone()).collect();
                let value = self.calculate_aggregation(&owned_records, agg)?;
                let agg_key = agg.alias.clone().unwrap_or_else(|| agg.function.clone());
                group_result.insert(agg_key, value);
            } else {
                // Multiple aggregations
                let mut agg_results = HashMap::new();
                for agg in aggregations {
                    let owned_records: Vec<JsonValue> =
                        group_records.iter().map(|&r| r.clone()).collect();
                    let value = self.calculate_aggregation(&owned_records, agg)?;
                    let agg_key = agg
                        .alias
                        .clone()
                        .unwrap_or_else(|| format!("{}_{}", agg.function, agg.field));
                    agg_results.insert(agg_key, value);
                }
                group_result.insert("aggregations".to_string(), json!(agg_results));
            }

            results.push(json!(group_result));
        }

        Ok(json!({
            "type": "grouped_aggregation",
            "group_by": group_by_fields,
            "results": results
        }))
    }

    /// Calculate a single aggregation value
    fn calculate_aggregation(
        &self,
        records: &[JsonValue],
        agg_spec: &AggregationSpec,
    ) -> Result<JsonValue> {
        let func = &agg_spec.function;
        let field = &agg_spec.field;

        // Handle count(*)
        if func == "count" && field == "*" {
            return Ok(json!(records.len()));
        }

        // Extract field values
        let values: Vec<JsonValue> = records
            .iter()
            .filter_map(|record| match field_accessor::get_field(record, field) {
                Ok(Some(value)) => {
                    if value != &JsonValue::Null {
                        Some(value.clone())
                    } else {
                        None
                    }
                }
                _ => None,
            })
            .collect();

        // Calculate aggregation based on function
        match func.to_lowercase().as_str() {
            "count" => Ok(json!(values.len())),

            "unique_count" | "cardinality" => {
                let unique: HashSet<String> = values
                    .iter()
                    .map(|v| serde_json::to_string(v).unwrap_or_default())
                    .collect();
                Ok(json!(unique.len()))
            }

            "sum" => {
                let sum: f64 = values.iter().filter_map(|v| self.to_numeric(v)).sum();
                Ok(json!(sum))
            }

            "min" => {
                let min = values
                    .iter()
                    .filter_map(|v| self.to_numeric(v))
                    .min_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
                Ok(json!(min))
            }

            "max" => {
                let max = values
                    .iter()
                    .filter_map(|v| self.to_numeric(v))
                    .max_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
                Ok(json!(max))
            }

            "average" | "avg" | "mean" => {
                let numeric_values: Vec<f64> =
                    values.iter().filter_map(|v| self.to_numeric(v)).collect();

                if numeric_values.is_empty() {
                    Ok(JsonValue::Null)
                } else {
                    let avg = numeric_values.iter().sum::<f64>() / numeric_values.len() as f64;
                    Ok(json!(avg))
                }
            }

            "median" | "med" => {
                let mut numeric_values: Vec<f64> =
                    values.iter().filter_map(|v| self.to_numeric(v)).collect();

                if numeric_values.is_empty() {
                    return Ok(JsonValue::Null);
                }

                numeric_values
                    .sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
                let len = numeric_values.len();

                let median = if len.is_multiple_of(2) {
                    (numeric_values[len / 2 - 1] + numeric_values[len / 2]) / 2.0
                } else {
                    numeric_values[len / 2]
                };

                Ok(json!(median))
            }

            "std" | "stdev" | "standard_deviation" => {
                let numeric_values: Vec<f64> =
                    values.iter().filter_map(|v| self.to_numeric(v)).collect();

                if numeric_values.len() < 2 {
                    return Ok(JsonValue::Null);
                }

                let mean = numeric_values.iter().sum::<f64>() / numeric_values.len() as f64;
                let variance = numeric_values
                    .iter()
                    .map(|v| (v - mean).powi(2))
                    .sum::<f64>()
                    / (numeric_values.len() - 1) as f64;
                let std_dev = variance.sqrt();

                Ok(json!(std_dev))
            }

            "percentile" | "percentiles" | "p" | "pct" => {
                let mut numeric_values: Vec<f64> =
                    values.iter().filter_map(|v| self.to_numeric(v)).collect();

                if numeric_values.is_empty() {
                    return Ok(JsonValue::Null);
                }

                numeric_values
                    .sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));

                // Get percentile values from params (default to 50 for median)
                let percentile_values = agg_spec
                    .params
                    .get("percentile_values")
                    .and_then(|v| v.as_array())
                    .map(|arr| arr.iter().filter_map(|v| v.as_f64()).collect::<Vec<f64>>())
                    .unwrap_or_else(|| vec![50.0]);

                if percentile_values.len() == 1 {
                    // Single percentile
                    let p = self.calculate_percentile(&numeric_values, percentile_values[0]);
                    Ok(json!(p))
                } else {
                    // Multiple percentiles
                    let mut result = HashMap::new();
                    for p_val in percentile_values {
                        let p = self.calculate_percentile(&numeric_values, p_val);
                        result.insert(format!("p{}", p_val as i32), json!(p));
                    }
                    Ok(json!(result))
                }
            }

            "values" | "unique" => {
                // Return unique values
                let mut unique: Vec<String> = values
                    .iter()
                    .map(|v| match v {
                        JsonValue::String(s) => s.clone(),
                        JsonValue::Number(n) => n.to_string(),
                        JsonValue::Bool(b) => b.to_string(),
                        _ => serde_json::to_string(v).unwrap_or_default(),
                    })
                    .collect::<HashSet<_>>()
                    .into_iter()
                    .collect();

                unique.sort();
                Ok(json!(unique))
            }

            _ => Err(TqlError::ExecutionError(format!(
                "Unsupported aggregation function: {}",
                func
            ))),
        }
    }

    /// Convert value to numeric
    fn to_numeric(&self, value: &JsonValue) -> Option<f64> {
        match value {
            JsonValue::Number(n) => n.as_f64(),
            JsonValue::String(s) => s.parse::<f64>().ok(),
            JsonValue::Bool(b) => Some(if *b { 1.0 } else { 0.0 }),
            _ => None,
        }
    }

    /// Calculate percentile using nearest-rank method
    fn calculate_percentile(&self, sorted_values: &[f64], percentile: f64) -> Option<f64> {
        if sorted_values.is_empty() {
            return None;
        }

        let index = ((percentile / 100.0) * (sorted_values.len() as f64 - 1.0)).round() as usize;
        Some(sorted_values[index.min(sorted_values.len() - 1)])
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;

    #[test]
    fn test_simple_count() {
        let evaluator = StatsEvaluator::new();
        let records = vec![
            json!({"name": "Alice", "age": 30}),
            json!({"name": "Bob", "age": 25}),
            json!({"name": "Charlie", "age": 35}),
        ];

        let query = StatsQuery {
            aggregations: vec![AggregationSpec {
                function: "count".to_string(),
                field: "*".to_string(),
                alias: None,
                params: HashMap::new(),
            }],
            group_by: vec![],
        };

        let result = evaluator.evaluate_stats(&records, &query).unwrap();
        assert_eq!(result["value"], json!(3));
    }

    #[test]
    fn test_sum_aggregation() {
        let evaluator = StatsEvaluator::new();
        let records = vec![
            json!({"name": "Alice", "score": 90}),
            json!({"name": "Bob", "score": 85}),
            json!({"name": "Charlie", "score": 95}),
        ];

        let query = StatsQuery {
            aggregations: vec![AggregationSpec {
                function: "sum".to_string(),
                field: "score".to_string(),
                alias: None,
                params: HashMap::new(),
            }],
            group_by: vec![],
        };

        let result = evaluator.evaluate_stats(&records, &query).unwrap();
        assert_eq!(result["value"], json!(270.0));
    }

    #[test]
    fn test_average_aggregation() {
        let evaluator = StatsEvaluator::new();
        let records = vec![
            json!({"name": "Alice", "age": 30}),
            json!({"name": "Bob", "age": 20}),
            json!({"name": "Charlie", "age": 40}),
        ];

        let query = StatsQuery {
            aggregations: vec![AggregationSpec {
                function: "avg".to_string(),
                field: "age".to_string(),
                alias: None,
                params: HashMap::new(),
            }],
            group_by: vec![],
        };

        let result = evaluator.evaluate_stats(&records, &query).unwrap();
        assert_eq!(result["value"], json!(30.0));
    }

    #[test]
    fn test_min_max_aggregation() {
        let evaluator = StatsEvaluator::new();
        let records = vec![
            json!({"value": 10}),
            json!({"value": 50}),
            json!({"value": 30}),
        ];

        let query = StatsQuery {
            aggregations: vec![AggregationSpec {
                function: "min".to_string(),
                field: "value".to_string(),
                alias: Some("min_value".to_string()),
                params: HashMap::new(),
            }],
            group_by: vec![],
        };

        let result = evaluator.evaluate_stats(&records, &query).unwrap();
        assert_eq!(result["value"], json!(10.0));
    }

    #[test]
    fn test_grouped_aggregation() {
        let evaluator = StatsEvaluator::new();
        let records = vec![
            json!({"city": "NYC", "sales": 100}),
            json!({"city": "LA", "sales": 150}),
            json!({"city": "NYC", "sales": 200}),
            json!({"city": "LA", "sales": 250}),
        ];

        let query = StatsQuery {
            aggregations: vec![AggregationSpec {
                function: "sum".to_string(),
                field: "sales".to_string(),
                alias: None,
                params: HashMap::new(),
            }],
            group_by: vec!["city".to_string()],
        };

        let result = evaluator.evaluate_stats(&records, &query).unwrap();
        assert_eq!(result["type"], "grouped_aggregation");

        let results = result["results"].as_array().unwrap();
        assert_eq!(results.len(), 2);
    }

    #[test]
    fn test_multiple_aggregations() {
        let evaluator = StatsEvaluator::new();
        let records = vec![
            json!({"score": 90}),
            json!({"score": 85}),
            json!({"score": 95}),
        ];

        let query = StatsQuery {
            aggregations: vec![
                AggregationSpec {
                    function: "sum".to_string(),
                    field: "score".to_string(),
                    alias: Some("total".to_string()),
                    params: HashMap::new(),
                },
                AggregationSpec {
                    function: "avg".to_string(),
                    field: "score".to_string(),
                    alias: Some("average".to_string()),
                    params: HashMap::new(),
                },
            ],
            group_by: vec![],
        };

        let result = evaluator.evaluate_stats(&records, &query).unwrap();
        assert_eq!(result["type"], "multiple_aggregations");
        assert_eq!(result["results"]["total"], json!(270.0));
        assert_eq!(result["results"]["average"], json!(90.0));
    }

    #[test]
    fn test_median() {
        let evaluator = StatsEvaluator::new();
        let records = vec![
            json!({"value": 10}),
            json!({"value": 20}),
            json!({"value": 30}),
            json!({"value": 40}),
            json!({"value": 50}),
        ];

        let query = StatsQuery {
            aggregations: vec![AggregationSpec {
                function: "median".to_string(),
                field: "value".to_string(),
                alias: None,
                params: HashMap::new(),
            }],
            group_by: vec![],
        };

        let result = evaluator.evaluate_stats(&records, &query).unwrap();
        assert_eq!(result["value"], json!(30.0));
    }

    #[test]
    fn test_unique_count() {
        let evaluator = StatsEvaluator::new();
        let records = vec![
            json!({"city": "NYC"}),
            json!({"city": "LA"}),
            json!({"city": "NYC"}),
            json!({"city": "SF"}),
        ];

        let query = StatsQuery {
            aggregations: vec![AggregationSpec {
                function: "unique_count".to_string(),
                field: "city".to_string(),
                alias: None,
                params: HashMap::new(),
            }],
            group_by: vec![],
        };

        let result = evaluator.evaluate_stats(&records, &query).unwrap();
        assert_eq!(result["value"], json!(3));
    }
}