tellaro-query-language 0.2.23

//! TQL stats AST to OpenSearch aggregation DSL translator.
//!
//! Converts TQL stats expressions (e.g., `| stats count(*) by event.code`)
//! into OpenSearch aggregation DSL with proper nested terms aggregations
//! for groupby, metric aggregation mappings, and `.keyword` field resolution.

use super::field_mappings::{FieldMappings, FieldType};
use crate::parser::{Aggregation, GroupBy, StatsNode};
use serde_json::{json, Value as JsonValue};

/// Default bucket size for grouped aggregations (matches Kibana/Grafana/Splunk defaults).
const DEFAULT_BUCKET_SIZE: usize = 5;

/// Map TQL aggregation function names to OpenSearch aggregation types.
fn map_aggregation(func: &str) -> Option<&'static str> {
    match func {
        "count" => Some("value_count"),
        "unique_count" | "cardinality" => Some("cardinality"),
        "sum" => Some("sum"),
        "min" => Some("min"),
        "max" => Some("max"),
        "average" | "avg" | "mean" => Some("avg"),
        "median" | "med" => Some("percentiles"),
        "std" | "stdev" | "standard_deviation" => Some("extended_stats"),
        "percentile" | "percentiles" | "p" | "pct" => Some("percentiles"),
        "percentile_rank" | "percentile_ranks" | "pct_rank" | "pct_ranks" => {
            Some("percentile_ranks")
        }
        "values" | "unique" => Some("terms"),
        _ => None,
    }
}

/// Resolve a field name to its `.keyword` variant if it's a text field.
///
/// OpenSearch cannot aggregate on text fields — terms aggregations require
/// keyword fields.
fn resolve_aggregation_field(field: &str, field_mappings: Option<&FieldMappings>) -> String {
    if let Some(mappings) = field_mappings {
        // If the base field is text type, use .keyword for aggregation
        if let Some(ft) = mappings.get_field_type(field) {
            if *ft == FieldType::Text {
                return format!("{}.keyword", field);
            }
        }
        // Also check if a .keyword subfield exists
        let keyword_field = format!("{}.keyword", field);
        if mappings.get_field_type(&keyword_field).is_some() {
            return keyword_field;
        }
    }
    field.to_string()
}

/// Build OpenSearch aggregation DSL from a TQL stats AST.
///
/// Returns a JSON object with the `aggs` key containing the aggregation DSL.
pub fn translate_stats(
    stats: &StatsNode,
    field_mappings: Option<&FieldMappings>,
) -> Result<JsonValue, String> {
    if stats.aggregations.is_empty() {
        return Err("No aggregations specified in stats query".to_string());
    }

    let aggs_dsl = if stats.group_by.is_empty() {
        build_simple_aggregations(&stats.aggregations)?
    } else {
        build_grouped_aggregations(&stats.aggregations, &stats.group_by, field_mappings)?
    };

    Ok(json!({ "aggs": aggs_dsl }))
}

/// Build simple aggregations without grouping.
fn build_simple_aggregations(aggregations: &[Aggregation]) -> Result<JsonValue, String> {
    let mut aggs = serde_json::Map::new();

    for (i, agg) in aggregations.iter().enumerate() {
        let func = agg.function.to_lowercase();
        let field = agg.field.as_deref().unwrap_or("*");
        let alias = agg
            .alias
            .clone()
            .unwrap_or_else(|| format!("{}_{}", func, i));

        let agg_dsl = build_single_aggregation(&func, field, agg)?;
        aggs.insert(alias, agg_dsl);
    }

    Ok(JsonValue::Object(aggs))
}

/// Build a single aggregation DSL object.
fn build_single_aggregation(
    func: &str,
    field: &str,
    agg: &Aggregation,
) -> Result<JsonValue, String> {
    // Special case: count(*)
    if func == "count" && field == "*" {
        return Ok(json!({ "value_count": { "field": "_id" } }));
    }

    let os_type =
        map_aggregation(func).ok_or_else(|| format!("Unknown aggregation function: {}", func))?;

    match func {
        "median" | "med" => Ok(json!({ "percentiles": { "field": field, "percents": [50.0] } })),
        "std" | "stdev" | "standard_deviation" => {
            Ok(json!({ "extended_stats": { "field": field } }))
        }
        "percentile" | "percentiles" | "p" | "pct" => {
            let percents = agg
                .percentile_values
                .as_ref()
                .cloned()
                .unwrap_or_else(|| vec![50.0]);
            Ok(json!({ "percentiles": { "field": field, "percents": percents } }))
        }
        "percentile_rank" | "percentile_ranks" | "pct_rank" | "pct_ranks" => {
            let values = agg.rank_values.as_ref().cloned().unwrap_or_default();
            if values.is_empty() {
                return Err("percentile_rank requires at least one value".to_string());
            }
            Ok(json!({ "percentile_ranks": { "field": field, "values": values } }))
        }
        "values" | "unique" => Ok(json!({ "terms": { "field": field, "size": 10000 } })),
        _ => {
            // Direct mapping (count, sum, min, max, avg, cardinality, etc.)
            Ok(json!({ os_type: { "field": field } }))
        }
    }
}

/// Build aggregations with grouping (nested terms aggregations).
fn build_grouped_aggregations(
    aggregations: &[Aggregation],
    group_by: &[GroupBy],
    field_mappings: Option<&FieldMappings>,
) -> Result<JsonValue, String> {
    // Start with innermost aggregations (the metric aggs)
    let inner_aggs = build_simple_aggregations(aggregations)?;

    // Check for top/bottom ordering modifiers
    let mut order_field: Option<String> = None;
    let mut order_direction = "desc";
    for (i, agg) in aggregations.iter().enumerate() {
        if let Some(ref modifier) = agg.modifier {
            let alias = agg
                .alias
                .clone()
                .unwrap_or_else(|| format!("{}_{}", agg.function, i));
            order_field = Some(alias);
            order_direction = if modifier == "top" { "desc" } else { "asc" };
            break;
        }
    }

    // Build nested terms aggregations in reverse order
    let mut current_aggs = inner_aggs;

    for (i, gb) in group_by.iter().rev().enumerate() {
        let bucket_size = gb.bucket_size.unwrap_or(DEFAULT_BUCKET_SIZE);
        let resolved_field = resolve_aggregation_field(&gb.field, field_mappings);

        let mut terms_agg = json!({
            "terms": {
                "field": resolved_field,
                "size": bucket_size
            }
        });

        // Add ordering on the outermost aggregation if modifier present
        if i == group_by.len() - 1 {
            if let Some(ref of_) = order_field {
                terms_agg["terms"]["order"] = json!({ of_: order_direction });
            }
        }

        // Nest sub-aggregations
        if current_aggs.as_object().is_some_and(|m| !m.is_empty()) {
            terms_agg["aggs"] = current_aggs;
        }

        let key = format!("group_by_{}", gb.field);
        current_aggs = json!({ key: terms_agg });
    }

    Ok(current_aggs)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::parser::TqlParser;

    fn parse_stats(query: &str) -> StatsNode {
        let parser = TqlParser::new();
        let ast = parser.parse(query).expect("parse failed");
        match ast {
            crate::parser::AstNode::StatsExpr(s) => s,
            crate::parser::AstNode::QueryWithStats(qws) => qws.stats,
            _ => panic!("expected stats AST, got {:?}", ast),
        }
    }

    #[test]
    fn test_simple_count_star() {
        let stats = parse_stats("| stats count(*)");
        let dsl = translate_stats(&stats, None).unwrap();
        let aggs = &dsl["aggs"];
        // Should have a value_count on _id
        assert!(aggs
            .as_object()
            .unwrap()
            .values()
            .any(|v| { v.get("value_count").is_some() }));
    }

    #[test]
    fn test_count_by_field() {
        let stats = parse_stats("| stats count(*) by event.code");
        let dsl = translate_stats(&stats, None).unwrap();
        let aggs = &dsl["aggs"];
        // Should have group_by_event.code with nested aggs
        let group = aggs.get("group_by_event.code").expect("missing group_by");
        assert!(group.get("terms").is_some());
        assert!(group.get("aggs").is_some());
    }

    #[test]
    fn test_multiple_group_by() {
        let stats = parse_stats("| stats count(*) by host.name, event.code");
        let dsl = translate_stats(&stats, None).unwrap();
        let aggs = &dsl["aggs"];
        // Outermost should be group_by_host.name
        let outer = aggs.get("group_by_host.name").expect("missing outer group");
        assert!(outer.get("terms").is_some());
        // Should have nested group_by_event.code
        let inner_aggs = outer.get("aggs").expect("missing inner aggs");
        assert!(inner_aggs.get("group_by_event.code").is_some());
    }

    #[test]
    fn test_avg_aggregation() {
        let stats = parse_stats("| stats avg(response_time)");
        let dsl = translate_stats(&stats, None).unwrap();
        let aggs = &dsl["aggs"];
        assert!(aggs
            .as_object()
            .unwrap()
            .values()
            .any(|v| { v.get("avg").is_some() }));
    }
}