tellaro-query-language 1.3.8

//! Post-processing for operations not supported by OpenSearch.
//!
//! This module applies mutators and filters to query results that couldn't
//! be pushed down to OpenSearch. It is essential for TQL queries with `scan_all=True`
//! where mutators like `is_private`, `is_global`, `lowercase`, etc. need to be
//! applied to filter or transform results after they are fetched from OpenSearch.

use super::error::Result;
use crate::evaluator::TqlEvaluator;
use crate::field_accessor;
use crate::mutators::{self, Mutator};
use crate::parser::{Mutator as MutatorSpec, TqlParser};
use serde_json::Value as JsonValue;

/// Post-processor for applying operations on query results.
///
/// The PostProcessor handles operations that cannot be pushed to OpenSearch,
/// such as field mutators (is_private, is_global, lowercase, etc.) and
/// complex filtering conditions.
pub struct PostProcessor {
    parser: TqlParser,
    evaluator: TqlEvaluator,
}

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

impl PostProcessor {
    /// Create a new PostProcessor
    pub fn new() -> Self {
        Self {
            parser: TqlParser::new(),
            evaluator: TqlEvaluator::new(),
        }
    }

    /// Apply post-processing to OpenSearch results using a TQL query.
    ///
    /// This method parses the TQL query, extracts mutators, applies them to
    /// transform field values, and then filters based on the query conditions.
    ///
    /// # Arguments
    ///
    /// * `results` - Query results from OpenSearch (typically from `_source` field)
    /// * `tql_query` - The original TQL query string
    ///
    /// # Returns
    ///
    /// Filtered and transformed results
    ///
    /// # Example
    ///
    /// ```ignore
    /// use tql::opensearch::PostProcessor;
    ///
    /// let results = vec![
    ///     json!({"source": {"ip": "192.168.1.1"}}),
    ///     json!({"source": {"ip": "8.8.8.8"}}),
    /// ];
    ///
    /// let processor = PostProcessor::new();
    /// let filtered = processor.process_results(results, "source.ip | is_private eq true")?;
    /// assert_eq!(filtered.len(), 1); // Only the 192.168.x.x IP
    /// ```
    pub fn process_results(
        &self,
        results: Vec<JsonValue>,
        tql_query: &str,
    ) -> Result<Vec<JsonValue>> {
        // Parse the TQL query
        let ast = self.parser.parse(tql_query).map_err(|e| {
            super::error::OpenSearchError::TranslationError(format!("Failed to parse TQL: {}", e))
        })?;

        // Filter results using the evaluator
        // The evaluator handles mutators during evaluation
        let filtered: Vec<JsonValue> = results
            .into_iter()
            .filter(|record| self.evaluator.evaluate(&ast, record).unwrap_or(false))
            .collect();

        Ok(filtered)
    }

    /// Apply post-processing with enrichment (mutators modify the output records).
    ///
    /// This is similar to `process_results` but the returned records will have
    /// mutator transformations applied to them (e.g., fields converted to lowercase).
    ///
    /// # Arguments
    ///
    /// * `results` - Query results from OpenSearch
    /// * `tql_query` - The original TQL query string
    ///
    /// # Returns
    ///
    /// Filtered and enriched results
    pub fn process_results_with_enrichment(
        &self,
        results: Vec<JsonValue>,
        tql_query: &str,
    ) -> Result<Vec<JsonValue>> {
        // Parse the TQL query
        let ast = self.parser.parse(tql_query).map_err(|e| {
            super::error::OpenSearchError::TranslationError(format!("Failed to parse TQL: {}", e))
        })?;

        // Filter and enrich results
        let enriched = self
            .evaluator
            .filter_and_enrich(&ast, &results)
            .map_err(|e| {
                super::error::OpenSearchError::TranslationError(format!(
                    "Failed to apply post-processing: {}",
                    e
                ))
            })?;

        Ok(enriched)
    }

    /// Apply a list of mutators to a specific field in a list of records.
    ///
    /// This is a lower-level method for applying mutators directly.
    ///
    /// # Arguments
    ///
    /// * `results` - Query results from OpenSearch
    /// * `mutators` - List of mutator instances to apply
    /// * `field` - Field name to apply mutators to (supports dot notation)
    ///
    /// # Returns
    ///
    /// Records with mutators applied to the specified field
    pub fn apply_mutators(
        results: Vec<JsonValue>,
        mutators: &[Box<dyn Mutator>],
        field: &str,
    ) -> Result<Vec<JsonValue>> {
        if mutators.is_empty() {
            return Ok(results);
        }

        let mut processed = Vec::with_capacity(results.len());

        for mut record in results {
            // Get the field value
            let field_value = match field_accessor::get_field(&record, field) {
                Ok(Some(value)) => value.clone(),
                _ => {
                    // Field doesn't exist, keep the record as-is
                    processed.push(record);
                    continue;
                }
            };

            // Apply mutators in sequence
            let mut current_value = field_value;
            let mut is_enrichment_mutator = false;
            for mutator in mutators {
                // Check if this is an enrichment mutator (like nslookup, geoip_lookup)
                if mutator.is_enrichment() {
                    is_enrichment_mutator = true;
                }
                current_value = mutator.apply(field, &record, &current_value).map_err(|e| {
                    super::error::OpenSearchError::TranslationError(format!(
                        "Mutator {} failed: {}",
                        mutator.name(),
                        e
                    ))
                })?;
            }

            // Handle enrichment mutators specially - they return enrichment structures
            // that should be applied to separate fields, not overwrite the original field
            if is_enrichment_mutator {
                // Check if this is a TQL enrichment structure
                if let Some(enrichment) = current_value.get("_tql_enrichment") {
                    // Apply enrichment data to the appropriate fields
                    if let Some(enrichment_type) = enrichment.get("type").and_then(|v| v.as_str()) {
                        match enrichment_type {
                            "dns" => {
                                // Apply DNS enrichment
                                if let Some(domain_field) =
                                    enrichment.get("domain_field").and_then(|v| v.as_str())
                                {
                                    if let Some(domain) = enrichment.get("domain") {
                                        let _ = field_accessor::set_field(
                                            &mut record,
                                            domain_field,
                                            domain.clone(),
                                        );
                                    }
                                }
                                if let Some(dns_field) =
                                    enrichment.get("dns_field").and_then(|v| v.as_str())
                                {
                                    if let Some(dns) = enrichment.get("dns") {
                                        let _ = field_accessor::set_field(
                                            &mut record,
                                            dns_field,
                                            dns.clone(),
                                        );
                                    }
                                }
                            }
                            "geo" => {
                                // Apply geo enrichment (future implementation)
                            }
                            _ => {}
                        }
                    }
                }
                // Don't overwrite the original field for enrichment mutators
            } else {
                // Update the field in the record for non-enrichment mutators
                if let Err(_e) = field_accessor::set_field(&mut record, field, current_value) {
                    // If we can't set the field, continue with the original record
                    // This can happen for deeply nested fields that don't exist
                }
            }

            processed.push(record);
        }

        Ok(processed)
    }

    /// Apply mutators and filter based on a comparison value.
    ///
    /// This is used when the TQL query has a mutator that transforms a value
    /// and then compares it (e.g., `source.ip | is_private eq true`).
    ///
    /// # Arguments
    ///
    /// * `results` - Query results from OpenSearch
    /// * `field` - Field name to apply mutators to
    /// * `mutator_specs` - List of mutator specifications
    /// * `operator` - Comparison operator
    /// * `compare_value` - Value to compare against
    ///
    /// # Returns
    ///
    /// Filtered results
    pub fn apply_mutators_and_filter(
        results: Vec<JsonValue>,
        field: &str,
        mutator_specs: &[MutatorSpec],
        operator: &str,
        compare_value: &JsonValue,
    ) -> Result<Vec<JsonValue>> {
        if mutator_specs.is_empty() {
            // No mutators, shouldn't really happen but handle it
            return Ok(results);
        }

        // Create mutator instances
        let mutators: Vec<Box<dyn Mutator>> = mutator_specs
            .iter()
            .map(|spec| mutators::create_mutator(&spec.name, None))
            .collect::<std::result::Result<Vec<_>, _>>()
            .map_err(|e| {
                super::error::OpenSearchError::TranslationError(format!(
                    "Failed to create mutator: {}",
                    e
                ))
            })?;

        let mut filtered = Vec::new();

        for record in results {
            // Get the field value
            let field_value = match field_accessor::get_field(&record, field) {
                Ok(Some(value)) => value.clone(),
                _ => continue, // Field doesn't exist, skip
            };

            // Apply mutators in sequence
            let mut current_value = field_value;
            let mut mutator_failed = false;

            for mutator in &mutators {
                match mutator.apply(field, &record, &current_value) {
                    Ok(new_value) => current_value = new_value,
                    Err(_) => {
                        mutator_failed = true;
                        break;
                    }
                }
            }

            if mutator_failed {
                continue;
            }

            // Extract the comparison value from enrichment mutators
            // Enrichment mutators return a structure with _tql_return_value
            let compare_source = if let Some(return_value) = current_value.get("_tql_return_value")
            {
                return_value.clone()
            } else {
                current_value
            };

            // Compare the mutated value
            if Self::compare_values(&compare_source, operator, compare_value) {
                filtered.push(record);
            }
        }

        Ok(filtered)
    }

    /// Compare two JSON values using the specified operator.
    fn compare_values(left: &JsonValue, operator: &str, right: &JsonValue) -> bool {
        match operator.to_lowercase().as_str() {
            "eq" | "=" | "==" => left == right,
            "ne" | "!=" | "<>" => left != right,
            "gt" | ">" => Self::compare_numeric(left, right, |l, r| l > r),
            "gte" | ">=" => Self::compare_numeric(left, right, |l, r| l >= r),
            "lt" | "<" => Self::compare_numeric(left, right, |l, r| l < r),
            "lte" | "<=" => Self::compare_numeric(left, right, |l, r| l <= r),
            "contains" => Self::string_contains(left, right),
            "startswith" => Self::string_starts_with(left, right),
            "endswith" => Self::string_ends_with(left, right),
            _ => false,
        }
    }

    /// Compare numeric values
    fn compare_numeric<F>(left: &JsonValue, right: &JsonValue, compare: F) -> bool
    where
        F: Fn(f64, f64) -> bool,
    {
        let left_num = match left {
            JsonValue::Number(n) => n.as_f64(),
            _ => None,
        };
        let right_num = match right {
            JsonValue::Number(n) => n.as_f64(),
            _ => None,
        };

        match (left_num, right_num) {
            (Some(l), Some(r)) => compare(l, r),
            _ => false,
        }
    }

    /// Check if left string contains right string
    fn string_contains(left: &JsonValue, right: &JsonValue) -> bool {
        match (left.as_str(), right.as_str()) {
            (Some(l), Some(r)) => l.contains(r),
            _ => false,
        }
    }

    /// Check if left string starts with right string
    fn string_starts_with(left: &JsonValue, right: &JsonValue) -> bool {
        match (left.as_str(), right.as_str()) {
            (Some(l), Some(r)) => l.starts_with(r),
            _ => false,
        }
    }

    /// Check if left string ends with right string
    fn string_ends_with(left: &JsonValue, right: &JsonValue) -> bool {
        match (left.as_str(), right.as_str()) {
            (Some(l), Some(r)) => l.ends_with(r),
            _ => false,
        }
    }

    /// Filter results based on conditions not supported by OpenSearch.
    ///
    /// # Arguments
    ///
    /// * `results` - Query results from OpenSearch
    /// * `filter_fn` - Filter function to apply
    ///
    /// # Returns
    ///
    /// Filtered results
    pub fn filter_results<F>(results: Vec<JsonValue>, filter_fn: F) -> Result<Vec<JsonValue>>
    where
        F: Fn(&JsonValue) -> bool,
    {
        Ok(results.into_iter().filter(filter_fn).collect())
    }
}

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

    #[test]
    fn test_filter_results() {
        let results = vec![json!({"age": 25}), json!({"age": 30}), json!({"age": 35})];

        let filtered = PostProcessor::filter_results(results, |r| {
            r.get("age").and_then(|v| v.as_i64()).unwrap_or(0) > 25
        })
        .unwrap();

        assert_eq!(filtered.len(), 2);
    }

    #[test]
    fn test_process_results_simple() {
        let results = vec![
            json!({"name": "John", "age": 30}),
            json!({"name": "Jane", "age": 25}),
            json!({"name": "Bob", "age": 35}),
        ];

        let processor = PostProcessor::new();
        let filtered = processor.process_results(results, "age > 25").unwrap();

        assert_eq!(filtered.len(), 2);
    }

    #[test]
    fn test_process_results_with_mutators() {
        let results = vec![
            json!({"name": "JOHN"}),
            json!({"name": "jane"}),
            json!({"name": "BOB"}),
        ];

        let processor = PostProcessor::new();
        let filtered = processor
            .process_results(results, "name | lowercase eq 'john'")
            .unwrap();

        assert_eq!(filtered.len(), 1);
        assert_eq!(filtered[0]["name"], "JOHN");
    }

    #[test]
    fn test_apply_mutators_to_field() {
        let results = vec![json!({"message": "  hello  "}), json!({"message": "world"})];

        let mutators: Vec<Box<dyn Mutator>> = vec![mutators::create_mutator("trim", None).unwrap()];

        let processed = PostProcessor::apply_mutators(results, &mutators, "message").unwrap();

        assert_eq!(processed[0]["message"], "hello");
        assert_eq!(processed[1]["message"], "world");
    }

    #[test]
    fn test_apply_mutators_and_filter() {
        let results = vec![
            json!({"ip": "192.168.1.1"}),
            json!({"ip": "8.8.8.8"}),
            json!({"ip": "10.0.0.1"}),
        ];

        let mutator_specs = vec![MutatorSpec {
            name: "is_private".to_string(),
            args: vec![],
        }];

        let filtered = PostProcessor::apply_mutators_and_filter(
            results,
            "ip",
            &mutator_specs,
            "eq",
            &json!(true),
        )
        .unwrap();

        // Should only include private IPs (192.168.x.x and 10.x.x.x)
        assert_eq!(filtered.len(), 2);
    }

    #[test]
    fn test_compare_values_equality() {
        assert!(PostProcessor::compare_values(
            &json!("hello"),
            "eq",
            &json!("hello")
        ));
        assert!(!PostProcessor::compare_values(
            &json!("hello"),
            "eq",
            &json!("world")
        ));
        assert!(PostProcessor::compare_values(
            &json!(true),
            "eq",
            &json!(true)
        ));
        assert!(PostProcessor::compare_values(&json!(42), "eq", &json!(42)));
    }

    #[test]
    fn test_compare_values_numeric() {
        assert!(PostProcessor::compare_values(&json!(10), "gt", &json!(5)));
        assert!(!PostProcessor::compare_values(&json!(5), "gt", &json!(10)));
        assert!(PostProcessor::compare_values(&json!(10), "gte", &json!(10)));
        assert!(PostProcessor::compare_values(&json!(5), "lt", &json!(10)));
        assert!(PostProcessor::compare_values(&json!(10), "lte", &json!(10)));
    }

    #[test]
    fn test_compare_values_string_operations() {
        assert!(PostProcessor::compare_values(
            &json!("hello world"),
            "contains",
            &json!("world")
        ));
        assert!(PostProcessor::compare_values(
            &json!("hello world"),
            "startswith",
            &json!("hello")
        ));
        assert!(PostProcessor::compare_values(
            &json!("hello world"),
            "endswith",
            &json!("world")
        ));
    }

    #[test]
    fn test_process_results_with_is_private() {
        let results = vec![
            json!({"source": {"ip": "192.168.1.1"}}),
            json!({"source": {"ip": "8.8.8.8"}}),
            json!({"source": {"ip": "10.0.0.1"}}),
        ];

        let processor = PostProcessor::new();
        let filtered = processor
            .process_results(results, "source.ip | is_private eq true")
            .unwrap();

        // Should only include private IPs
        assert_eq!(filtered.len(), 2);
    }

    #[test]
    fn test_process_results_compound_query() {
        let results = vec![
            json!({"name": "ADMIN", "age": 35}),
            json!({"name": "USER", "age": 25}),
            json!({"name": "admin", "age": 40}),
        ];

        let processor = PostProcessor::new();
        let filtered = processor
            .process_results(results, "name | lowercase eq 'admin' AND age > 30")
            .unwrap();

        assert_eq!(filtered.len(), 2); // Both ADMIN (35) and admin (40) match
    }
}