tellaro-query-language 2.0.0

//! OpenSearch query executor for TQL.
//!
//! This module provides the core `execute_opensearch()` functionality that executes
//! TQL queries against OpenSearch with automatic post-processing for mutators that
//! cannot be pushed down to the database (is_private, is_global, geo, etc.).
//!
//! # Features
//!
//! - TQL to OpenSearch DSL conversion
//! - Automatic scroll API for scan_all mode (unlimited results)
//! - Post-processing for mutators not supported by OpenSearch
//! - Health status tracking
//! - Query analysis and optimization
//!
//! # Example
//!
//! ```no_run
//! use tellaro_query_language::opensearch::{OpenSearchConfig, TqlExecutor, ExecuteOptions};
//!
//! #[tokio::main]
//! async fn main() -> Result<(), Box<dyn std::error::Error>> {
//!     let config = OpenSearchConfig::from_env()?;
//!     let executor = TqlExecutor::new(config)?;
//!
//!     let result = executor.execute_opensearch(
//!         "source.ip | is_private eq true",
//!         "endpoint-*",
//!         ExecuteOptions::default().with_scan_all(true),
//!     ).await?;
//!
//!     println!("Found {} results", result.total);
//!     Ok(())
//! }
//! ```

use super::error::{OpenSearchError, Result};
use super::field_mappings::FieldMappings;
use super::post_processor::PostProcessor;
use super::query_builder::QueryBuilder;
use super::OpenSearchConfig;
use crate::Tql;
use opensearch::OpenSearch;
use serde::{Deserialize, Serialize};
use serde_json::{json, Value as JsonValue};

/// Source of field mappings used for TQL query generation.
///
/// This enum tracks where field mappings came from, enabling clear debug logging
/// to understand whether intelligent field selection was used and from what source.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum MappingSource {
    /// Mappings provided directly via `with_field_mappings()` (e.g., from index templates)
    Provided {
        /// Number of field mappings provided
        field_count: usize,
    },
    /// Mappings fetched from OpenSearch index via `with_mappings_from_index()`
    FetchedFromIndex {
        /// Index pattern used to fetch mappings
        index: String,
        /// Number of field mappings fetched
        field_count: usize,
    },
    /// No mappings available (using raw field names without intelligent selection)
    None {
        /// Reason why no mappings are available
        reason: String,
    },
}

impl Default for MappingSource {
    fn default() -> Self {
        Self::None {
            reason: "Not configured".to_string(),
        }
    }
}

impl std::fmt::Display for MappingSource {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            MappingSource::Provided { field_count } => {
                write!(f, "Provided ({} fields)", field_count)
            }
            MappingSource::FetchedFromIndex { index, field_count } => {
                write!(f, "FetchedFromIndex '{}' ({} fields)", index, field_count)
            }
            MappingSource::None { reason } => {
                write!(f, "None ({})", reason)
            }
        }
    }
}

/// How hard OpenSearch should work to count matches.
///
/// The wire form is a sum type (`true` / `false` / an integer), which is why
/// this is an enum rather than a `bool`: a `bool` cannot express the bounded
/// form, and the bounded form is the one an operator actually wants on a
/// cluster where an exact count of a broad query is expensive.
///
/// `false` is deliberately not offered. It removes `hits.total` from the
/// response entirely, and the only observable effect through this API would be
/// `ExecuteResult::opensearch_total` reporting a confident `0` -- which is the
/// bug this type exists to close, not a mode worth supporting.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum TrackTotalHits {
    /// Count every match. `hits.total.relation` comes back `eq`.
    ///
    /// The default, for two reasons. The Python implementation has always sent
    /// `track_total_hits=True`
    /// (`src/tql/core_components/opensearch_operations.py`), so anything else
    /// leaves the two implementations answering differently for the same query.
    /// And the alternative default is "report the 10,000 window as though it
    /// were the truth", which is the defect (tql#170).
    ///
    /// It is not free: OpenSearch counts the full match set rather than
    /// stopping at 10,000. On a narrow query that costs nothing; on a
    /// match-everything query over a large index it is a full postings scan.
    /// Callers who care can opt into `UpTo(n)`.
    #[default]
    Exact,
    /// Count exactly up to `n`, then stop and report `{"value": n, "relation": "gte"}`.
    /// `UpTo(10_000)` reproduces OpenSearch's own default, and therefore this
    /// executor's behaviour before the option existed.
    UpTo(u64),
}

impl TrackTotalHits {
    fn to_json(self) -> JsonValue {
        match self {
            Self::Exact => json!(true),
            Self::UpTo(n) => json!(n),
        }
    }
}

/// Whether an `opensearch_total` is exact or a floor.
///
/// Mirrors OpenSearch's `hits.total.relation`. This is the field whose entire
/// purpose is to say "this number is a lower bound, not a count" -- without it
/// a 40,000-hit search and a 10,000-hit search are byte-identical on the wire.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
#[serde(rename_all = "lowercase")]
pub enum TotalRelation {
    /// Exact: there are precisely `opensearch_total` matches.
    Eq,
    /// A floor: there are AT LEAST `opensearch_total` matches.
    ///
    /// This is the serde default deliberately. A payload produced by a TQL
    /// older than this change carries no relation field, which means its
    /// producer never asked for exact counting and its total may well be the
    /// 10,000 window. "At least N" is true whether the count was exact or
    /// capped; "exactly N" would be a claim of precision we invented. A default
    /// of `Eq` here would reintroduce the exact bug across a version skew.
    #[default]
    Gte,
}

impl TotalRelation {
    /// The OpenSearch wire spelling.
    ///
    /// Lives here so consumers -- notably tellaro-agent's search proxy, which
    /// hardcoded `"eq"` -- never hand-roll the mapping.
    pub fn as_wire_str(self) -> &'static str {
        match self {
            Self::Eq => "eq",
            Self::Gte => "gte",
        }
    }

    fn from_wire(s: &str) -> Self {
        if s.eq_ignore_ascii_case("eq") {
            Self::Eq
        } else {
            Self::Gte
        }
    }
}

/// Read `hits.total` out of an OpenSearch search response.
///
/// One implementation for both send paths. It was duplicated verbatim, and
/// both copies projected `value` out of the object and dropped `relation` on
/// the floor -- the whole of tql#170.
///
/// A bare number is the ES2-era shape: self-consistently an exact count, so it
/// maps to `Eq`. An object with no `relation` key maps to `Eq` for the same
/// reason -- the only producer that omits it is a caller synthesizing a
/// response from documents it counted itself.
///
/// Note this is the OPPOSITE default from `TotalRelation`'s serde default, and
/// deliberately so: they answer different questions. Here the producer counted
/// the documents in front of us; there, a producer we never met declined to say.
/// Stamp the total-tracking option onto an outgoing search body.
///
/// Pure and shared by both send paths, so a unit test can assert what goes on
/// the wire without a cluster or an HTTP mock. The two send paths previously
/// built their bodies inline and neither ever set this key.
fn apply_total_tracking(body: &mut JsonValue, options: &ExecuteOptions) {
    body["track_total_hits"] = options.track_total_hits.to_json();
}

fn parse_hits_total(response_body: &JsonValue) -> (usize, TotalRelation) {
    let Some(total) = response_body.get("hits").and_then(|h| h.get("total")) else {
        return (0, TotalRelation::Eq);
    };
    match total {
        JsonValue::Object(obj) => (
            obj.get("value").and_then(|v| v.as_u64()).unwrap_or(0) as usize,
            obj.get("relation")
                .and_then(|r| r.as_str())
                .map_or(TotalRelation::Eq, TotalRelation::from_wire),
        ),
        other => (other.as_u64().unwrap_or(0) as usize, TotalRelation::Eq),
    }
}

/// Result of executing a TQL query against OpenSearch.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExecuteResult {
    /// Matching documents (post-processed if applicable)
    pub results: Vec<JsonValue>,

    /// Total number of results after post-processing
    pub total: usize,

    /// Total documents from OpenSearch before post-processing
    pub opensearch_total: usize,

    /// Whether `opensearch_total` is exact or a lower bound.
    ///
    /// Additive with a serde default so an older consumer keeps deserializing
    /// and an older payload keeps loading. Changing `opensearch_total`'s type
    /// instead would have been the tidier data model and the wrong trade: it
    /// breaks tellaro-agent's search proxy at compile time and forces a
    /// lockstep release, for a field nothing downstream reads yet.
    #[serde(default)]
    pub opensearch_total_relation: TotalRelation,

    /// Whether post-processing was applied
    pub post_processing_applied: bool,

    /// Query health status: "green", "yellow", or "red"
    pub health_status: String,

    /// List of health issues/warnings
    pub health_reasons: Vec<String>,

    /// Scroll/scan information
    pub scan_info: Option<ScanInfo>,

    /// Error message if execution failed (None on success)
    pub error: Option<String>,

    /// The generated OpenSearch DSL query (for debugging)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub dsl_query: Option<JsonValue>,

    /// Source of field mappings used for query generation (for debugging)
    #[serde(default)]
    pub mapping_source: MappingSource,

    /// Aggregation results for stats queries (raw OpenSearch aggregations)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub aggregations: Option<JsonValue>,
}

impl Default for ExecuteResult {
    fn default() -> Self {
        Self {
            results: Vec::new(),
            total: 0,
            opensearch_total: 0,
            opensearch_total_relation: TotalRelation::default(),
            post_processing_applied: false,
            health_status: "green".to_string(),
            health_reasons: Vec::new(),
            scan_info: None,
            error: None,
            dsl_query: None,
            mapping_source: MappingSource::default(),
            aggregations: None,
        }
    }
}

impl ExecuteResult {
    /// Create an error result
    pub fn error(message: impl Into<String>) -> Self {
        let msg = message.into();
        Self {
            health_status: "red".to_string(),
            health_reasons: vec![msg.clone()],
            error: Some(msg),
            // Zero results, exactly. The `Gte` default exists for
            // cross-version deserialization, not for values we construct
            // ourselves -- rendering a failed query as "at least 0 matches"
            // would be a strange thing to put on the wire.
            opensearch_total_relation: TotalRelation::Eq,
            ..Default::default()
        }
    }

    /// True when `opensearch_total` is an exact count rather than a floor.
    pub fn opensearch_total_is_exact(&self) -> bool {
        self.opensearch_total_relation == TotalRelation::Eq
    }
}

/// Information about scroll/scan operations.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ScanInfo {
    /// Number of scroll batches executed
    pub batches: usize,

    /// Total documents scrolled
    pub total_scrolled: usize,

    /// Scroll batch size
    pub scroll_size: usize,

    /// Whether scroll was used
    pub scroll_used: bool,
}

/// Options for executing a TQL query.
#[derive(Debug, Clone)]
pub struct ExecuteOptions {
    /// Whether to fetch all results using scroll API
    pub scan_all: bool,

    /// Scroll batch size (default: 10000)
    pub scroll_size: usize,

    /// Scroll timeout (default: "5m")
    pub scroll_timeout: String,

    /// Sort specification
    pub sort: Option<Vec<JsonValue>>,

    /// Time range filter
    pub time_range: Option<TimeRange>,

    /// Timestamp field name (default: "@timestamp")
    pub timestamp_field: String,

    /// Maximum results to return (ignored if scan_all is true)
    pub size: usize,

    /// How exactly to count total matches (default: `Exact`).
    ///
    /// Unset, OpenSearch applies its own default of 10,000 and answers
    /// `{"value": 10000, "relation": "gte"}` for anything larger (tql#170).
    pub track_total_hits: TrackTotalHits,
}

impl Default for ExecuteOptions {
    fn default() -> Self {
        Self {
            scan_all: false,
            scroll_size: 10000,
            scroll_timeout: "5m".to_string(),
            sort: None,
            time_range: None,
            timestamp_field: "@timestamp".to_string(),
            size: 10000,
            track_total_hits: TrackTotalHits::default(),
        }
    }
}

impl ExecuteOptions {
    /// Enable scan_all mode for unlimited results
    pub fn with_scan_all(mut self, scan_all: bool) -> Self {
        self.scan_all = scan_all;
        self
    }

    /// Set scroll batch size
    pub fn with_scroll_size(mut self, size: usize) -> Self {
        self.scroll_size = size;
        self
    }

    /// Set scroll timeout
    pub fn with_scroll_timeout(mut self, timeout: impl Into<String>) -> Self {
        self.scroll_timeout = timeout.into();
        self
    }

    /// Set time range filter
    pub fn with_time_range(mut self, gte: impl Into<String>, lt: impl Into<String>) -> Self {
        self.time_range = Some(TimeRange {
            gte: gte.into(),
            lt: lt.into(),
        });
        self
    }

    /// Set timestamp field
    pub fn with_timestamp_field(mut self, field: impl Into<String>) -> Self {
        self.timestamp_field = field.into();
        self
    }

    /// Set sort specification
    pub fn with_sort(mut self, sort: Vec<JsonValue>) -> Self {
        self.sort = Some(sort);
        self
    }

    /// Set max results (only used when scan_all is false)
    pub fn with_size(mut self, size: usize) -> Self {
        self.size = size;
        self
    }

    /// Set how exactly total matches are counted.
    ///
    /// Use `TrackTotalHits::UpTo(n)` to bound the counting cost on a very large
    /// index; the reported total then comes back with a `Gte` relation, which
    /// is the honest answer rather than a cheaper wrong one.
    pub fn with_track_total_hits(mut self, track: TrackTotalHits) -> Self {
        self.track_total_hits = track;
        self
    }
}

/// Time range for filtering.
#[derive(Debug, Clone)]
pub struct TimeRange {
    /// Greater than or equal to
    pub gte: String,
    /// Less than
    pub lt: String,
}

/// TQL query executor for OpenSearch.
///
/// Supports field mappings for intelligent query generation. When field mappings
/// are provided, the executor will use them to select appropriate field variants
/// (e.g., using `.keyword` subfield for exact matches on text fields).
pub struct TqlExecutor {
    client: OpenSearch,
    tql: Tql,
    query_builder: QueryBuilder,
    post_processor: PostProcessor,
    field_mappings: Option<FieldMappings>,
    /// Tracks where field mappings came from for debugging
    mapping_source: MappingSource,
}

impl TqlExecutor {
    /// Create a new TQL executor from configuration.
    pub fn new(config: OpenSearchConfig) -> Result<Self> {
        let client = config.create_client()?;
        Ok(Self {
            client,
            tql: Tql::new(),
            query_builder: QueryBuilder::new(None),
            post_processor: PostProcessor::new(),
            field_mappings: None,
            mapping_source: MappingSource::default(),
        })
    }

    /// Create a new TQL executor from an existing OpenSearch client.
    pub fn from_client(client: OpenSearch) -> Self {
        Self {
            client,
            tql: Tql::new(),
            query_builder: QueryBuilder::new(None),
            post_processor: PostProcessor::new(),
            field_mappings: None,
            mapping_source: MappingSource::default(),
        }
    }

    /// Set field mappings for intelligent query generation.
    ///
    /// Field mappings enable the executor to use appropriate field variants
    /// based on the operator being used. For example, for exact matches on
    /// text fields, it will use the `.keyword` subfield if available.
    ///
    /// # Arguments
    ///
    /// * `mappings` - Field mappings extracted from OpenSearch index mappings
    ///
    /// # Example
    ///
    /// ```ignore
    /// use tellaro_query_language::opensearch::{TqlExecutor, FieldMappings};
    ///
    /// let mappings_response = client.indices().get_mapping().send().await?;
    /// let field_mappings = FieldMappings::from_opensearch_response(mappings_response)?;
    ///
    /// let executor = TqlExecutor::from_client(client)
    ///     .with_field_mappings(field_mappings);
    /// ```
    pub fn with_field_mappings(mut self, mappings: FieldMappings) -> Self {
        let field_count = mappings.len();
        self.field_mappings = Some(mappings.clone());
        self.query_builder = QueryBuilder::new(Some(mappings));
        self.mapping_source = MappingSource::Provided { field_count };
        self
    }

    /// Fetch field mappings from an OpenSearch index and configure the executor.
    ///
    /// This is a convenience method that fetches mappings from the specified
    /// index pattern and configures the executor to use them.
    ///
    /// # Arguments
    ///
    /// * `index` - Index pattern to fetch mappings from (e.g., "endpoint-windows-*")
    ///
    /// # Returns
    ///
    /// Self with field mappings configured, or unchanged if fetching fails.
    pub async fn with_mappings_from_index(mut self, index: &str) -> Self {
        match self.fetch_field_mappings(index).await {
            Ok(mappings) => {
                let field_count = mappings.len();
                self.field_mappings = Some(mappings.clone());
                self.query_builder = QueryBuilder::new(Some(mappings));
                self.mapping_source = MappingSource::FetchedFromIndex {
                    index: index.to_string(),
                    field_count,
                };
            }
            Err(e) => {
                self.mapping_source = MappingSource::None {
                    reason: format!("Failed to fetch from '{}': {}", index, e),
                };
                tracing::warn!(
                    "Failed to fetch field mappings for index '{}': {}. Using default query generation.",
                    index, e
                );
            }
        }
        self
    }

    /// Get the source of field mappings being used.
    ///
    /// Returns information about where the field mappings came from,
    /// useful for debugging and logging.
    pub fn get_mapping_source(&self) -> &MappingSource {
        &self.mapping_source
    }

    /// Fetch field mappings from an OpenSearch index.
    ///
    /// # Arguments
    ///
    /// * `index` - Index pattern to fetch mappings from
    ///
    /// # Returns
    ///
    /// FieldMappings extracted from the index, or an error if fetching fails.
    pub async fn fetch_field_mappings(&self, index: &str) -> Result<FieldMappings> {
        let response = self
            .client
            .indices()
            .get_mapping(opensearch::indices::IndicesGetMappingParts::Index(&[index]))
            .send()
            .await
            .map_err(|e| OpenSearchError::MappingError(e.to_string()))?;

        let response_body = response.json::<JsonValue>().await.map_err(|e| {
            OpenSearchError::MappingError(format!("Failed to parse mapping response: {}", e))
        })?;

        FieldMappings::from_opensearch_response(response_body)
            .map_err(|e| OpenSearchError::MappingError(format!("Failed to parse mappings: {}", e)))
    }

    /// Execute a TQL query against OpenSearch.
    ///
    /// This is the main entry point that handles:
    /// 1. Parsing the TQL query
    /// 2. Checking for post-processing mutators
    /// 3. Converting to OpenSearch DSL
    /// 4. Executing with scroll API if scan_all is true
    /// 5. Applying post-processing for mutators not supported by OpenSearch
    ///
    /// # Arguments
    ///
    /// * `query` - TQL query string
    /// * `index` - Index pattern to search (e.g., "endpoint-*")
    /// * `options` - Execution options
    ///
    /// # Returns
    ///
    /// ExecuteResult containing matching documents and metadata
    pub async fn execute_opensearch(
        &self,
        query: &str,
        index: &str,
        options: ExecuteOptions,
    ) -> Result<ExecuteResult> {
        // Parse the TQL query
        let ast = match self.tql.parse(query) {
            Ok(ast) => ast,
            Err(e) => {
                return Ok(ExecuteResult::error(format!("Failed to parse TQL: {}", e)));
            }
        };

        // Check if query requires post-processing
        let needs_post_processing = self.tql.ast_has_post_processing_mutators(&ast);

        // Determine execution strategy
        let use_scroll = options.scan_all || needs_post_processing;

        // Build OpenSearch DSL query
        let dsl_query = match self.query_builder.build_query(&ast) {
            Ok(q) => q,
            Err(e) => {
                return Ok(ExecuteResult::error(format!(
                    "Failed to build DSL query: {}",
                    e
                )));
            }
        };

        // Add time range filter if specified
        let final_query = self.add_time_range_filter(dsl_query, &options);

        // Debug log the generated DSL query. Uses tracing::debug! (not an
        // unconditional eprintln) so it is filtered out at INFO and above, and
        // when emitted is captured at DEBUG level by the host's tracing
        // subscriber instead of being mis-surfaced as ERROR on stderr.
        tracing::debug!(
            "Generated OpenSearch DSL:\n{}",
            serde_json::to_string_pretty(&final_query)
                .unwrap_or_else(|_| "Failed to serialize".to_string())
        );

        // Execute the query
        if use_scroll {
            self.execute_with_scroll(query, index, final_query, options, needs_post_processing)
                .await
        } else {
            self.execute_simple(query, index, final_query, options, needs_post_processing)
                .await
        }
    }

    /// Add time range filter to the query.
    fn add_time_range_filter(&self, mut query: JsonValue, options: &ExecuteOptions) -> JsonValue {
        if let Some(ref time_range) = options.time_range {
            let range_filter = json!({
                "range": {
                    &options.timestamp_field: {
                        "gte": &time_range.gte,
                        "lt": &time_range.lt,
                        "format": "strict_date_optional_time"
                    }
                }
            });

            // Add to existing bool filter or create new bool query
            if let Some(query_obj) = query.get_mut("query") {
                if let Some(bool_query) = query_obj.get_mut("bool") {
                    // Add to existing filter array
                    if let Some(filter_arr) = bool_query.get_mut("filter") {
                        if let Some(arr) = filter_arr.as_array_mut() {
                            arr.push(range_filter);
                        }
                    } else {
                        bool_query["filter"] = json!([range_filter]);
                    }
                } else {
                    // Wrap existing query in bool with filter
                    let existing = query_obj.take();
                    *query_obj = json!({
                        "bool": {
                            "must": [existing],
                            "filter": [range_filter]
                        }
                    });
                }
            } else {
                // No query object, create one
                query["query"] = json!({
                    "bool": {
                        "filter": [range_filter]
                    }
                });
            }
        }

        query
    }

    /// Execute query with scroll API for unlimited results.
    async fn execute_with_scroll(
        &self,
        tql_query: &str,
        index: &str,
        query: JsonValue,
        options: ExecuteOptions,
        needs_post_processing: bool,
    ) -> Result<ExecuteResult> {
        let mut all_hits: Vec<JsonValue> = Vec::new();
        let mut batches = 0;

        // Build scroll query
        let mut scroll_query = query.clone();
        scroll_query["size"] = json!(options.scroll_size);
        // The scroll path always asks for an exact count, ignoring any bound the
        // caller set. Two independent reasons, and the first is not negotiable:
        //
        //  * OpenSearch REFUSES an integer track_total_hits in a scroll
        //    context -- "Validation Failed: 1: disabling [track_total_hits] is
        //    not allowed in a scroll context" -- and rejects the whole search.
        //    Forwarding ExecuteOptions verbatim here turned every scan_all with
        //    a bound into a hard error. Measured against a live 2.19.4 cluster;
        //    no amount of asserting on the emitted body would have found it.
        //  * A bound is meaningless here anyway. scan_all fetches the entire
        //    match set, so the caller is already paying for a full pass.
        //
        // It is still sent rather than skipped: the loop breaks without
        // erroring on several paths, so all_hits.len() is a lower bound after a
        // possibly-partial scroll, not a total. Reporting that as exact would
        // re-create tql#170 one level down.
        scroll_query["track_total_hits"] = TrackTotalHits::Exact.to_json();

        // Add sort if specified
        if let Some(ref sort) = options.sort {
            scroll_query["sort"] = json!(sort);
        } else {
            // Default sort by timestamp descending
            scroll_query["sort"] = json!([{&options.timestamp_field: {"order": "desc"}}]);
        }

        // Initial scroll search
        let response = self
            .client
            .search(opensearch::SearchParts::Index(&[index]))
            .scroll(&options.scroll_timeout)
            .body(scroll_query)
            .send()
            .await
            .map_err(|e| OpenSearchError::SearchError(e.to_string()))?;

        let response_body = response.json::<JsonValue>().await.map_err(|e| {
            OpenSearchError::SearchError(format!("Failed to parse response: {}", e))
        })?;

        // Check for errors
        if let Some(error) = response_body.get("error") {
            return Ok(ExecuteResult::error(format!("OpenSearch error: {}", error)));
        }

        // Extract scroll ID and hits
        let mut scroll_id = response_body
            .get("_scroll_id")
            .and_then(|s| s.as_str())
            .map(|s| s.to_string());

        let (opensearch_total, opensearch_total_relation) = parse_hits_total(&response_body);

        // Collect first batch
        if let Some(hits) = response_body
            .get("hits")
            .and_then(|h| h.get("hits"))
            .and_then(|h| h.as_array())
        {
            all_hits.extend(hits.clone());
            batches += 1;
        }

        // Continue scrolling until no more results
        while let Some(ref current_scroll_id) = scroll_id {
            // Don't continue if we got fewer results than scroll_size (last batch)
            let last_batch_size = response_body
                .get("hits")
                .and_then(|h| h.get("hits"))
                .and_then(|h| h.as_array())
                .map(|a| a.len())
                .unwrap_or(0);

            if last_batch_size == 0 {
                break;
            }

            // Scroll request
            let scroll_response = self
                .client
                .scroll(opensearch::ScrollParts::None)
                .scroll(&options.scroll_timeout)
                .body(json!({
                    "scroll_id": current_scroll_id
                }))
                .send()
                .await
                .map_err(|e| OpenSearchError::ScrollError(e.to_string()))?;

            let scroll_body = scroll_response.json::<JsonValue>().await.map_err(|e| {
                OpenSearchError::ScrollError(format!("Failed to parse scroll response: {}", e))
            })?;

            // Check for errors - don't fail, we have partial results
            if scroll_body.get("error").is_some() {
                break;
            }

            // Extract hits from scroll response
            let hits = scroll_body
                .get("hits")
                .and_then(|h| h.get("hits"))
                .and_then(|h| h.as_array());

            match hits {
                Some(batch) if !batch.is_empty() => {
                    all_hits.extend(batch.clone());
                    batches += 1;

                    // Update scroll_id for next iteration
                    scroll_id = scroll_body
                        .get("_scroll_id")
                        .and_then(|s| s.as_str())
                        .map(|s| s.to_string());
                }
                _ => {
                    // No more results
                    break;
                }
            }
        }

        // Clear scroll context
        if let Some(ref final_scroll_id) = scroll_id {
            let _ = self
                .client
                .clear_scroll(opensearch::ClearScrollParts::None)
                .body(json!({ "scroll_id": [final_scroll_id] }))
                .send()
                .await;
        }

        // Extract _source documents and merge _id/_score metadata (like Python TQL does)
        // The agent layer is responsible for wrapping into full OpenSearch hit format
        let source_docs: Vec<JsonValue> = all_hits
            .iter()
            .filter_map(|hit| {
                let mut doc = hit.get("_source").cloned()?;
                // Merge _id and _score into the document (Python TQL behavior)
                if let Some(obj) = doc.as_object_mut() {
                    if let Some(id) = hit.get("_id") {
                        obj.insert("_id".to_string(), id.clone());
                    }
                    if let Some(score) = hit.get("_score") {
                        obj.insert("_score".to_string(), score.clone());
                    }
                }
                Some(doc)
            })
            .collect();

        // Apply post-processing if needed
        let (final_results, post_processing_applied) = if needs_post_processing {
            // PROPAGATE. This arm used to answer a post-processing failure with
            // the UNFILTERED documents, which is the same swallow
            // `process_results` carried, one layer up and pointing the other
            // way: a query the engine refuses would have returned every
            // document in the result window instead of none.
            //
            // It was inert while `process_results` could only fail on an
            // unparseable query -- a query that reached here has already
            // translated, so the parse cannot fail twice. Making that function
            // report the branch's fail-closed refusals is exactly what would
            // have made this arm live, and a refusal rendered as "here is
            // everything" is worse than the zero hits it replaces.
            (
                self.post_processor
                    .process_results(source_docs, tql_query)?,
                true,
            )
        } else {
            (source_docs, false)
        };

        let total = final_results.len();

        // The scroll loop breaks without erroring on several paths (a failed
        // scroll request, an empty batch, the zero-batch guard). If we came out
        // holding fewer documents than the count promised, we did not scan the
        // whole match set -- so what we have is a floor, whatever OpenSearch
        // said. Reporting `eq` here would be the same confident-wrong-answer
        // shape tql#170 is about, just one level down.
        let opensearch_total_relation = if all_hits.len() < opensearch_total {
            TotalRelation::Gte
        } else {
            opensearch_total_relation
        };

        // Determine health status
        let (health_status, health_reasons) = if needs_post_processing {
            (
                "yellow".to_string(),
                vec!["Query requires post-processing mutators".to_string()],
            )
        } else {
            ("green".to_string(), Vec::new())
        };

        Ok(ExecuteResult {
            results: final_results,
            total,
            opensearch_total,
            opensearch_total_relation,
            post_processing_applied,
            health_status,
            health_reasons,
            scan_info: Some(ScanInfo {
                batches,
                total_scrolled: all_hits.len(),
                scroll_size: options.scroll_size,
                scroll_used: true,
            }),
            error: None,
            dsl_query: Some(query.clone()),
            mapping_source: self.mapping_source.clone(),
            aggregations: None, // Scroll doesn't return aggregations
        })
    }

    /// Execute simple query without scroll.
    async fn execute_simple(
        &self,
        tql_query: &str,
        index: &str,
        query: JsonValue,
        options: ExecuteOptions,
        needs_post_processing: bool,
    ) -> Result<ExecuteResult> {
        // Store original query for debugging
        let original_query = query.clone();

        // Build query with size (don't override size=0 for stats queries)
        let mut final_query = query;
        let is_stats_query = final_query.get("aggs").is_some();
        if !is_stats_query {
            final_query["size"] = json!(options.size);
        }
        // Outside the is_stats_query guard on purpose: a stats query runs with
        // size=0, but its hits.total is still what a consumer reports as the
        // number of matching documents, so it needs an honest one too.
        apply_total_tracking(&mut final_query, &options);

        // Add sort if specified
        if let Some(ref sort) = options.sort {
            final_query["sort"] = json!(sort);
        }

        // Execute search
        let response = self
            .client
            .search(opensearch::SearchParts::Index(&[index]))
            .body(final_query)
            .send()
            .await
            .map_err(|e| OpenSearchError::SearchError(e.to_string()))?;

        let response_body = response.json::<JsonValue>().await.map_err(|e| {
            OpenSearchError::SearchError(format!("Failed to parse response: {}", e))
        })?;

        // Check for errors
        if let Some(error) = response_body.get("error") {
            return Ok(ExecuteResult::error(format!("OpenSearch error: {}", error)));
        }

        let (opensearch_total, opensearch_total_relation) = parse_hits_total(&response_body);

        // Extract _source documents and merge _id/_score metadata (like Python TQL does)
        // The agent layer is responsible for wrapping into full OpenSearch hit format
        let source_docs: Vec<JsonValue> = response_body
            .get("hits")
            .and_then(|h| h.get("hits"))
            .and_then(|h| h.as_array())
            .map(|hits| {
                hits.iter()
                    .filter_map(|hit| {
                        let mut doc = hit.get("_source").cloned()?;
                        // Merge _id and _score into the document (Python TQL behavior)
                        if let Some(obj) = doc.as_object_mut() {
                            if let Some(id) = hit.get("_id") {
                                obj.insert("_id".to_string(), id.clone());
                            }
                            if let Some(score) = hit.get("_score") {
                                obj.insert("_score".to_string(), score.clone());
                            }
                        }
                        Some(doc)
                    })
                    .collect()
            })
            .unwrap_or_default();

        // Apply post-processing if needed
        let (final_results, post_processing_applied) = if needs_post_processing {
            // PROPAGATE. This arm used to answer a post-processing failure with
            // the UNFILTERED documents, which is the same swallow
            // `process_results` carried, one layer up and pointing the other
            // way: a query the engine refuses would have returned every
            // document in the result window instead of none.
            //
            // It was inert while `process_results` could only fail on an
            // unparseable query -- a query that reached here has already
            // translated, so the parse cannot fail twice. Making that function
            // report the branch's fail-closed refusals is exactly what would
            // have made this arm live, and a refusal rendered as "here is
            // everything" is worse than the zero hits it replaces.
            (
                self.post_processor
                    .process_results(source_docs, tql_query)?,
                true,
            )
        } else {
            (source_docs, false)
        };

        let total = final_results.len();

        // Extract aggregations for stats queries
        let aggregations = response_body.get("aggregations").cloned();

        Ok(ExecuteResult {
            results: final_results,
            total,
            opensearch_total,
            opensearch_total_relation,
            post_processing_applied,
            health_status: "green".to_string(),
            health_reasons: Vec::new(),
            scan_info: None,
            error: None,
            dsl_query: Some(original_query),
            mapping_source: self.mapping_source.clone(),
            aggregations,
        })
    }

    /// Analyze a TQL query for post-processing requirements.
    ///
    /// Useful for debugging and understanding query health before execution.
    pub fn analyze_query(&self, query: &str) -> QueryAnalysis {
        match self.tql.parse(query) {
            Ok(ast) => {
                let has_post_processing = self.tql.ast_has_post_processing_mutators(&ast);
                let mutators = Tql::extract_mutators_from_ast(&ast);

                QueryAnalysis {
                    has_post_processing,
                    health_status: if has_post_processing {
                        "yellow".to_string()
                    } else {
                        "green".to_string()
                    },
                    health_reasons: if has_post_processing {
                        vec!["Query contains post-processing mutators".to_string()]
                    } else {
                        Vec::new()
                    },
                    post_processing_mutators: mutators
                        .into_iter()
                        .flat_map(|(_, muts)| muts.into_iter().map(|m| m.name))
                        .collect(),
                    error: None,
                }
            }
            Err(e) => QueryAnalysis {
                has_post_processing: false,
                health_status: "red".to_string(),
                health_reasons: vec![format!("Parse error: {}", e)],
                post_processing_mutators: Vec::new(),
                error: Some(e.to_string()),
            },
        }
    }
}

/// Result of analyzing a TQL query.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QueryAnalysis {
    /// Whether the query requires post-processing
    pub has_post_processing: bool,

    /// Health status: "green", "yellow", or "red"
    pub health_status: String,

    /// Health issues/warnings
    pub health_reasons: Vec<String>,

    /// List of mutators requiring post-processing
    pub post_processing_mutators: Vec<String>,

    /// Error message if analysis failed
    pub error: Option<String>,
}

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

    #[test]
    fn test_execute_options_default() {
        let options = ExecuteOptions::default();
        assert!(!options.scan_all);
        assert_eq!(options.scroll_size, 10000);
        assert_eq!(options.scroll_timeout, "5m");
        assert_eq!(options.timestamp_field, "@timestamp");
    }

    #[test]
    fn test_execute_options_builder() {
        let options = ExecuteOptions::default()
            .with_scan_all(true)
            .with_scroll_size(5000)
            .with_scroll_timeout("10m")
            .with_timestamp_field("timestamp")
            .with_time_range("2024-01-01T00:00:00Z", "2024-01-02T00:00:00Z");

        assert!(options.scan_all);
        assert_eq!(options.scroll_size, 5000);
        assert_eq!(options.scroll_timeout, "10m");
        assert_eq!(options.timestamp_field, "timestamp");
        assert!(options.time_range.is_some());
    }

    // ---- tql#170: the total is counted honestly and reported honestly ----

    #[test]
    fn the_simple_body_asks_for_an_exact_total_by_default() {
        // Before this option existed the key was simply absent, so OpenSearch
        // applied its own default of 10,000 and every larger result set came
        // back as a flat 10000.
        let mut body = json!({"query": {"match_all": {}}});
        apply_total_tracking(&mut body, &ExecuteOptions::default());
        assert_eq!(body["track_total_hits"], json!(true));
    }

    #[test]
    fn a_bounded_count_is_sent_as_the_integer() {
        let mut body = json!({"query": {"match_all": {}}});
        let opts = ExecuteOptions::default().with_track_total_hits(TrackTotalHits::UpTo(50_000));
        apply_total_tracking(&mut body, &opts);
        assert_eq!(body["track_total_hits"], json!(50_000));
    }

    #[test]
    fn applying_total_tracking_leaves_the_rest_of_the_body_alone() {
        let mut body = json!({"query": {"match_all": {}}, "size": 10, "aggs": {"a": {}}});
        apply_total_tracking(&mut body, &ExecuteOptions::default());
        assert_eq!(body["size"], json!(10));
        assert!(body.get("aggs").is_some());
        assert_eq!(body["query"], json!({"match_all": {}}));
    }

    #[test]
    fn a_capped_total_is_read_as_a_floor_not_a_count() {
        // The exact response OpenSearch returns when it stops counting.
        let (total, relation) =
            parse_hits_total(&json!({"hits": {"total": {"value": 10000, "relation": "gte"}}}));
        assert_eq!(total, 10_000);
        assert_eq!(relation, TotalRelation::Gte);
    }

    #[test]
    fn an_exact_total_is_read_as_exact() {
        let (total, relation) =
            parse_hits_total(&json!({"hits": {"total": {"value": 11500, "relation": "eq"}}}));
        assert_eq!(total, 11_500);
        assert_eq!(relation, TotalRelation::Eq);
    }

    #[test]
    fn a_total_with_no_relation_is_exact() {
        // Both the pre-7.0 flat form and a self-counted synthetic response.
        // The producer counted the documents in front of it, so it is exact.
        assert_eq!(
            parse_hits_total(&json!({"hits": {"total": {"value": 7}}})),
            (7, TotalRelation::Eq)
        );
        assert_eq!(
            parse_hits_total(&json!({"hits": {"total": 7}})),
            (7, TotalRelation::Eq)
        );
    }

    #[test]
    fn a_missing_total_is_zero_and_exact() {
        assert_eq!(
            parse_hits_total(&json!({"hits": {}})),
            (0, TotalRelation::Eq)
        );
        assert_eq!(parse_hits_total(&json!({})), (0, TotalRelation::Eq));
    }

    #[test]
    fn an_unrecognised_relation_fails_closed_to_a_floor() {
        // Anything that is not exactly "eq" is treated as a lower bound, so a
        // relation OpenSearch adds later cannot be mistaken for a count.
        let (_, relation) = parse_hits_total(
            &json!({"hits": {"total": {"value": 5, "relation": "something_new"}}}),
        );
        assert_eq!(relation, TotalRelation::Gte);
    }

    #[test]
    fn the_wire_spelling_round_trips() {
        assert_eq!(TotalRelation::Eq.as_wire_str(), "eq");
        assert_eq!(TotalRelation::Gte.as_wire_str(), "gte");
        assert_eq!(
            serde_json::to_value(TotalRelation::Eq).unwrap(),
            json!("eq")
        );
        assert_eq!(
            serde_json::to_value(TotalRelation::Gte).unwrap(),
            json!("gte")
        );
    }

    #[test]
    fn a_payload_from_an_older_tql_is_not_assumed_exact() {
        // Build the "old" payload by serializing a current result and deleting
        // the field, rather than hand-writing JSON -- that way this keeps
        // testing version skew as the struct grows, instead of rotting into a
        // test of whether someone updated a string literal.
        let mut payload = serde_json::to_value(ExecuteResult {
            opensearch_total: 10_000,
            opensearch_total_relation: TotalRelation::Eq,
            ..Default::default()
        })
        .expect("serialize");
        payload
            .as_object_mut()
            .expect("object")
            .remove("opensearch_total_relation");

        let parsed: ExecuteResult =
            serde_json::from_value(payload).expect("an older payload must still deserialize");

        // No relation field means the producer never asked for exact counting,
        // so its total may well be the 10,000 window. "At least N" is true
        // either way; "exactly N" would be precision we invented.
        assert_eq!(parsed.opensearch_total_relation, TotalRelation::Gte);
        assert!(!parsed.opensearch_total_is_exact());
    }

    #[test]
    fn an_error_result_reports_exactly_zero() {
        let result = ExecuteResult::error("boom");
        assert_eq!(result.opensearch_total, 0);
        assert!(result.opensearch_total_is_exact());
    }

    #[test]
    fn test_execute_result_error() {
        let result = ExecuteResult::error("Test error");
        assert_eq!(result.health_status, "red");
        assert_eq!(result.error, Some("Test error".to_string()));
        assert!(result.health_reasons.contains(&"Test error".to_string()));
    }
}