ottyl/src/engine/snapshot/query.rs

//! Answering `EngineQuery`s: log/timeline listing (including the sort-helper-index path),
//! stats, declaration listing, bucket-value listing, and PromQL evaluation.

use crate::api::dto::{
    BucketValueDto, LogEntryDto, MetricDeclarationDto, ParamDeclarationDto,
    StatsDto, TimelineBucketDto,
};
use crate::consts::{DEFAULT_TIMELINE_BUCKET_MS, MERGE_WINDOW_MS};
use crate::engine::data_value::DataValue;
use crate::engine::index_key::IndexKey;
use crate::engine::message::{EngineQuery, EngineResponse};
use crate::engine::metric_primitive::{MetricPrimitive, Number};
use crate::engine::promql::{LabelMatcher, PromqlQuery};
use crate::engine::rate;
use crate::engine::snapshot::filter::CompiledFilter;
#[cfg(test)]
use crate::engine::snapshot::sort::SortHelperIndexEntry;
use crate::engine::snapshot::sort::{
    log_entry_params, CompiledParamDeclaration, SortHelperAccess, SortHelperKey,
};
use crate::engine::snapshot::EngineSnapshot;
use crate::engine::timestamp::Timestamp;
#[cfg(test)]
use std::collections::HashMap;
use std::collections::BTreeMap;

#[cfg(test)]
pub(crate) fn handle_query(
    snapshot: &EngineSnapshot,
    query: EngineQuery,
    sort_helper_indexes: &mut HashMap<String, SortHelperIndexEntry>,
) -> EngineResponse {
    handle_query_inner(snapshot, query, SortHelperAccess::Mutable(sort_helper_indexes))
}

/// Compiles `filter` against `snapshot`, or returns `Ok(None)` when there's no filter string
/// to compile. Shared by the `Logs` and `Timeline` query arms, which both bail out early with
/// `EngineResponse::Error` on a compile failure.
fn compile_optional_filter(
    snapshot: &EngineSnapshot,
    filter: Option<&str>,
    now: Timestamp,
) -> Result<Option<CompiledFilter>, String> {
    filter.map(|filter| snapshot.compile_filter_query(filter, now)).transpose()
}

/// Walks the sort-helper index for `sort_by` — rather than `snapshot.logs_range()`, which is
/// ordered by sequence, not by the param's value — collecting up to `limit` matching entries
/// in sort-by-param order.
///
/// Not range-bounded by since_sequence/start_ms the way the unsorted path is: `SortHelperKey::
/// WithParam` sorts by the param's value first (that's the point — it's what gives sort-by-param
/// order), so sequence/time bounds can't narrow a BTreeMap range seek here.
#[allow(clippy::too_many_arguments)]
fn logs_sorted_by_param(
    snapshot: &EngineSnapshot,
    sort_by: &str,
    sort_helper_access: SortHelperAccess<'_>,
    compiled_filter: Option<&CompiledFilter>,
    compiled_params: &[CompiledParamDeclaration],
    since_sequence: Option<u64>,
    start_ms: Option<i64>,
    end_ms: Option<i64>,
    sort_missing_fallback: bool,
    limit: usize,
) -> Vec<LogEntryDto> {
    let index = &snapshot.index;
    static EMPTY_HELPER: std::sync::LazyLock<BTreeMap<SortHelperKey, Timestamp>> =
        std::sync::LazyLock::new(BTreeMap::new);
    let helper_index = sort_helper_access.resolve(index, sort_by).unwrap_or(&EMPTY_HELPER);

    let mut entries = Vec::new();
    for (helper_key, timestamp) in helper_index.iter() {
        let sequence = match helper_key {
            SortHelperKey::WithParam(_, sequence) => *sequence,
            SortHelperKey::WithoutParam(sequence) => {
                if !sort_missing_fallback {
                    continue;
                }
                *sequence
            }
        };
        if !(since_sequence.is_none_or(|since| sequence > since)
            && start_ms.is_none_or(|start| timestamp.as_millis() >= start)
            && end_ms.is_none_or(|end| timestamp.as_millis() <= end))
        {
            continue;
        }
        let Some(DataValue::LogLiteral { text }) =
            index.get(&IndexKey::IngestedLog { sequence, timestamp: *timestamp })
        else {
            continue;
        };
        if !compiled_filter.is_none_or(|filter| filter.eval(sequence, *timestamp, text)) {
            continue;
        }
        entries.push(LogEntryDto {
            sequence,
            timestamp_ms: timestamp.as_millis(),
            text: text.clone(),
            params: log_entry_params(text, compiled_params),
        });
        if entries.len() >= limit {
            break;
        }
    }
    entries
}

pub(super) fn handle_query_inner(
    snapshot: &EngineSnapshot,
    query: EngineQuery,
    sort_helper_access: SortHelperAccess<'_>,
) -> EngineResponse {
    let approx_bytes_total = snapshot.approx_bytes_total;
    let ingest_lag_high_watermark_ms = snapshot.ingest_lag_high_watermark_ms;
    match query {
        EngineQuery::Logs(params) => {
            let compiled_filter =
                match compile_optional_filter(snapshot, params.filter.as_deref(), Timestamp::now()) {
                    Ok(compiled) => compiled,
                    Err(err) => return EngineResponse::Error(err),
                };
            if let Some(sort_by) = params.sort_by.as_deref() {
                if !snapshot.param_declarations_range().any(|(name, ..)| name == sort_by) {
                    return EngineResponse::Error(format!("unknown param '{sort_by}' for sort_by"));
                }
            }
            let limit = params.limit.unwrap_or(500).min(10_000);
            let compiled_params = snapshot.compiled_param_declarations();
            let since_bound = match params.since_sequence {
                Some(since) => std::ops::Bound::Excluded(since),
                None => std::ops::Bound::Unbounded,
            };
            let filtered = snapshot
                .logs_in_sequence_range((since_bound, std::ops::Bound::Unbounded))
                .filter(|(_, timestamp, _)| {
                    params.start_ms.is_none_or(|start| timestamp.as_millis() >= start)
                        && params.end_ms.is_none_or(|end| timestamp.as_millis() <= end)
                })
                .filter(|(sequence, timestamp, text)| {
                    compiled_filter
                        .as_ref()
                        .is_none_or(|filter| filter.eval(*sequence, *timestamp, text))
                });

            let entries: Vec<LogEntryDto> = match params.sort_by.as_deref() {
                None => filtered
                    .take(limit)
                    .map(|(sequence, timestamp, text)| LogEntryDto {
                        sequence,
                        timestamp_ms: timestamp.as_millis(),
                        text: text.to_string(),
                        params: log_entry_params(text, &compiled_params),
                    })
                    .collect(),
                Some(sort_by) => logs_sorted_by_param(
                    snapshot,
                    sort_by,
                    sort_helper_access,
                    compiled_filter.as_ref(),
                    &compiled_params,
                    params.since_sequence,
                    params.start_ms,
                    params.end_ms,
                    params.sort_missing_fallback.unwrap_or(true),
                    limit,
                ),
            };
            let entries: Vec<LogEntryDto> = entries.into_iter().take(limit).collect();
            EngineResponse::Logs(entries)
        }
        EngineQuery::Timeline(params) => {
            let compiled_filter =
                match compile_optional_filter(snapshot, params.filter.as_deref(), Timestamp::now()) {
                    Ok(compiled) => compiled,
                    Err(err) => return EngineResponse::Error(err),
                };
            let bucket_ms = params.bucket_ms.unwrap_or(DEFAULT_TIMELINE_BUCKET_MS).max(1);
            let mut buckets: BTreeMap<i64, u64> = BTreeMap::new();
            for (sequence, timestamp, text) in snapshot.logs_range().filter(|(_, timestamp, _)| {
                params.start_ms.is_none_or(|start| timestamp.as_millis() >= start)
                    && params.end_ms.is_none_or(|end| timestamp.as_millis() <= end)
            }) {
                if compiled_filter
                    .as_ref()
                    .is_some_and(|filter| !filter.eval(sequence, timestamp, text))
                {
                    continue;
                }
                let bucket_start = (timestamp.as_millis() / bucket_ms) * bucket_ms;
                *buckets.entry(bucket_start).or_insert(0) += 1;
            }
            let out = buckets
                .into_iter()
                .map(|(start_ms, count)| TimelineBucketDto { start_ms, count })
                .collect();
            EngineResponse::Timeline(out)
        }
        EngineQuery::Stats => {
            let mut log_count = 0u64;
            let mut first_sequence = None;
            let mut last_sequence = None;
            let mut first_timestamp_ms = None;
            let mut last_timestamp_ms = None;
            for (sequence, timestamp, _) in snapshot.logs_range() {
                log_count += 1;
                first_sequence.get_or_insert(sequence);
                last_sequence = Some(sequence);
                first_timestamp_ms.get_or_insert(timestamp.as_millis());
                last_timestamp_ms = Some(timestamp.as_millis());
            }
            EngineResponse::Stats(StatsDto {
                log_count,
                first_sequence,
                last_sequence,
                first_timestamp_ms,
                last_timestamp_ms,
                approx_bytes_total,
                ingest_lag_high_watermark_ms,
            })
        }
        EngineQuery::ParamDeclarations => {
            let declarations = snapshot.param_declarations_range()
                .map(|(name, message_match_regex, capture_group, json_field_path, kind, elapsed_since_last, auto_derived, _)| {
                    ParamDeclarationDto {
                        name: name.to_string(),
                        message_match_regex: message_match_regex.to_string(),
                        capture_group: capture_group.to_string(),
                        json_field_path: json_field_path.map(|path| path.to_string()),
                        kind,
                        elapsed_since_last,
                        auto_derived,
                        reserved: false,
                    }
                })
                .collect();
            EngineResponse::ParamDeclarations(declarations)
        }
        EngineQuery::MetricDeclarations => {
            let declarations = snapshot.metric_declarations_range()
                .map(|(name, message_match_regex, label_params, value_param, matched_labels, unmatched_labels)| {
                    MetricDeclarationDto {
                        name: name.to_string(),
                        message_match_regex: message_match_regex.to_string(),
                        label_params: label_params.to_vec(),
                        value_param: value_param.map(|s| s.to_string()),
                        matched_labels: matched_labels.cloned(),
                        unmatched_labels: unmatched_labels.cloned(),
                    }
                })
                .collect();
            EngineResponse::MetricDeclarations(declarations)
        }
        EngineQuery::BucketValues {
            metric_name,
            labels_filter,
            start_ms,
            end_ms,
            since_sequence,
            limit,
        } => {
            let limit = limit.unwrap_or(500).min(10_000);
            let values: Vec<BucketValueDto> = snapshot.bucket_values(&metric_name, None)
                .filter(|(labels, _)| {
                    labels_filter.as_ref().is_none_or(|filter| {
                        filter.iter().all(|(name, expected)| {
                            labels
                                .get(name)
                                .is_some_and(|value| &value.to_string() == expected)
                        })
                    })
                })
                .filter(|(_, (sequence, timestamp, _))| {
                    since_sequence.is_none_or(|since| *sequence > since)
                        && start_ms.is_none_or(|start| timestamp.as_millis() >= start)
                        && end_ms.is_none_or(|end| timestamp.as_millis() <= end)
                })
                .take(limit)
                .map(|(labels, (sequence, timestamp, value))| BucketValueDto {
                    labels: labels.clone(),
                    sequence: *sequence,
                    timestamp_ms: timestamp.as_millis(),
                    value: value.clone(),
                })
                .collect();
            EngineResponse::BucketValues(values)
        }
        EngineQuery::Promql { query, start_ms, end_ms } => {
            let response = PromqlQuery::parse(&query)
                .and_then(|parsed| snapshot.evaluate_promql(&parsed, start_ms, end_ms));
            EngineResponse::Promql(response)
        }
    }
}

pub(crate) fn coerce_to_f64(value: &MetricPrimitive) -> f64 {
    match value {
        MetricPrimitive::Number(n) => match n {
            Number::Int(v) => *v as f64,
            Number::Float(v) => *v,
        },
        MetricPrimitive::Timestamp(ts) => ts.as_millis() as f64,
        MetricPrimitive::Duration(ms) => *ms as f64,
        MetricPrimitive::String(_) => 1.0,
        MetricPrimitive::Boolean(b) => {
            if *b {
                1.0
            } else {
                0.0
            }
        }
    }
}

pub(crate) fn matches_label_matchers(
    labels: &BTreeMap<String, MetricPrimitive>,
    matchers: &[LabelMatcher],
) -> bool {
    matchers.iter().all(|matcher| {
        labels.get(&matcher.name).is_some_and(|value| value.to_string() == matcher.value)
    })
}

pub(crate) fn project_labels(
    labels: &BTreeMap<String, String>,
    by_labels: &[String],
) -> BTreeMap<String, String> {
    by_labels
        .iter()
        .filter_map(|name| labels.get(name).map(|value| (name.clone(), value.clone())))
        .collect()
}

pub(crate) fn windowed_rate_series(
    points: &[(i64, f64)],
    range_ms: i64,
    start_ms: Option<i64>,
    end_ms: Option<i64>,
) -> Vec<(i64, f64)> {
    let ts_points: Vec<(Timestamp, f64)> =
        points.iter().map(|(ms, value)| (Timestamp::from_millis(*ms), *value)).collect();
    let mut series: Vec<(i64, f64)> = rate::windowed_rate_series(&ts_points, range_ms)
        .into_iter()
        .map(|(ts, value)| (ts.as_millis(), value))
        .collect();

    if let Some(end) = end_ms {
        if series.last().is_none_or(|(ts, _)| *ts < end) {
            let value = rate::rate_at(&ts_points, Timestamp::from_millis(end), range_ms);
            series.push((end, value));
        }
    }
    if let Some(start) = start_ms {
        if series.first().is_none_or(|(ts, _)| *ts > start) {
            let value = rate::rate_at(&ts_points, Timestamp::from_millis(start), range_ms);
            series.insert(0, (start, value));
        }
    }
    series
}

pub(crate) fn merge_series_windowed(sub_series: &[Vec<(i64, f64)>], is_avg: bool) -> Vec<(i64, f64)> {
    if sub_series.len() <= 1 {
        return sub_series.first().cloned().unwrap_or_default();
    }
    let mut grouped: BTreeMap<i64, Vec<f64>> = BTreeMap::new();
    for points in sub_series {
        for (timestamp_ms, value) in points {
            let window_start = timestamp_ms.div_euclid(MERGE_WINDOW_MS) * MERGE_WINDOW_MS;
            grouped.entry(window_start).or_default().push(*value);
        }
    }
    grouped
        .into_iter()
        .map(|(timestamp_ms, values)| {
            let sum: f64 = values.iter().sum();
            let value = if is_avg { sum / values.len() as f64 } else { sum };
            (timestamp_ms, value)
        })
        .collect()
}

pub(crate) fn merge_series_locf(sub_series: &[Vec<(i64, f64)>], is_avg: bool) -> Vec<(i64, f64)> {
    let mut all_timestamps: std::collections::BTreeSet<i64> = std::collections::BTreeSet::new();
    for points in sub_series {
        all_timestamps.extend(points.iter().map(|(ts, _)| *ts));
    }
    let mut cursors = vec![0usize; sub_series.len()];
    let mut last_values = vec![0.0f64; sub_series.len()];
    all_timestamps
        .into_iter()
        .map(|ts| {
            for (i, points) in sub_series.iter().enumerate() {
                while cursors[i] < points.len() && points[cursors[i]].0 <= ts {
                    last_values[i] = points[cursors[i]].1;
                    cursors[i] += 1;
                }
            }
            let sum: f64 = last_values.iter().sum();
            let value = if is_avg { sum / last_values.len() as f64 } else { sum };
            (ts, value)
        })
        .collect()
}

pub(crate) fn ottyl_reserved_numeric_value_fn(metric_name: &str) -> Option<fn(u64, i64) -> f64> {
    match metric_name {
        "ottyl_seq" => Some(|sequence, _timestamp_ms| sequence as f64),
        "ottyl_ts" => Some(|_sequence, timestamp_ms| timestamp_ms as f64),
        _ => None,
    }
}