algeon-query-runtime 0.1.2

use std::collections::BTreeMap;

use query_planning::correlation::{PlanNodeId, QueryId, TraceId};
use serde_json::{Map, Value, json};

use super::{
    AggregateStrategyDecisionMetrics, BudgetedExecutionEvent, CompiledTaskId, DynamicFilterRuntimeMetrics,
    GraphAlgorithmRuntimeMetrics, JoinOperatorMetrics, JoinStrategyDecisionMetrics, OperatorMetrics, PlanTaskMapping,
    QueryMetrics, SourceReadParquetEvidenceMetrics, SourceReadRuntimeMetrics, TaskOperatorMetrics, TaskRuntimeMetrics,
};
use crate::observability::graph_observation::GraphMeasuredPeakRow;

/// Identifies the layout of the JSON document [`QueryMetrics::to_json`] emits.
pub const QUERY_METRICS_JSON_SCHEMA: &str = "query_engine.query_metrics";

impl QueryMetrics {
    /// Serializes query metrics as a nested, self-describing JSON document.
    ///
    /// The `metrics` object is generated by walking [`Self::descriptors`], so
    /// scalar JSON field names and descriptor names share one catalog.
    pub fn to_json(&self) -> Value {
        json!({
            "schema": QUERY_METRICS_JSON_SCHEMA,
            "query_id": self.query_id.as_str(),
            "trace_id": self.trace_id.as_ref().map(TraceId::as_str),
            "triage": self.triage_summary().to_json(),
            "completeness": self.completeness_json(),
            "descriptors": Self::descriptors().iter().map(|descriptor| descriptor.to_json()).collect::<Vec<_>>(),
            "metrics": self.described_metrics_json(),
            "rows": {
                "budgeted_executions": self.budgeted_execution_events.iter().map(BudgetedExecutionEvent::to_json).collect::<Vec<_>>(),
                "operators": self.operators.iter().map(OperatorMetrics::to_json).collect::<Vec<_>>(),
                "graph_claim_calibration": self
                    .graph_claim_calibration_rows
                    .iter()
                    .map(GraphMeasuredPeakRow::to_json)
                    .collect::<Vec<_>>(),
                "graph_runtime": self.graph_algorithm_runtime_rows.iter().map(GraphAlgorithmRuntimeMetrics::to_json).collect::<Vec<_>>(),
                "task_operators": self.task_operator_metrics.values().map(TaskOperatorMetrics::to_json).collect::<Vec<_>>(),
                "task_runtime": self.task_runtime_metrics.values().map(TaskRuntimeMetrics::to_json).collect::<Vec<_>>(),
                "plan_task_mappings": self.plan_task_mappings().iter().map(PlanTaskMapping::to_json).collect::<Vec<_>>(),
                "dynamic_filter_runtime": self.dynamic_filter_runtime_rows.iter().map(DynamicFilterRuntimeMetrics::to_json).collect::<Vec<_>>(),
                "source_read_runtime": self.source_read_runtime_rows.iter().map(SourceReadRuntimeMetrics::to_json).collect::<Vec<_>>(),
                "join_strategy_decisions": self.join_strategy_decisions.iter().map(JoinStrategyDecisionMetrics::to_json).collect::<Vec<_>>(),
                "aggregate_strategy_decisions": self.aggregate_strategy_decisions.iter().map(AggregateStrategyDecisionMetrics::to_json).collect::<Vec<_>>(),
            },
        })
    }

    /// Where these rows came from and whether they are all the detail there
    /// was. A value not projected from an observation owner cannot say, so its
    /// completeness is reported as unknown.
    fn completeness_json(&self) -> Value {
        let Some(projection) = &self.projection else {
            return json!({ "status": "unknown" });
        };
        let completeness = &projection.completeness;
        json!({
            "status": completeness.detail_status(),
            "revision": projection.revision,
            "sealed": projection.sealed,
            "capture_mode": completeness.capture_mode.as_str(),
            "truncated": completeness.truncated,
            "detail_rows_dropped": completeness.detail_rows_dropped,
            "totals_saturated": completeness.totals_saturated,
        })
    }

    fn described_metrics_json(&self) -> Value {
        let mut metrics = Map::new();
        for descriptor in Self::descriptors() {
            let value = self
                .metric_value_json(descriptor.name)
                .unwrap_or_else(|| panic!("QueryMetrics descriptor lacks JSON value: {}", descriptor.name));
            metrics.insert(
                descriptor.name.to_string(),
                json!({
                    "value": value,
                    "unit": descriptor.unit.as_str(),
                    "aggregation": descriptor.aggregation.as_str(),
                    "direction": descriptor.direction.map(|direction| direction.as_str()),
                    "healthy": descriptor.healthy,
                }),
            );
        }
        Value::Object(metrics)
    }

    fn metric_value_json(&self, name: &str) -> Option<Value> {
        Some(match name {
            "input_batches" => json!(self.input_batches),
            "input_rows" => json!(self.input_rows),
            "gpu_resident_input_handoffs" => json!(self.boundary.gpu_resident_input_handoffs),
            "host_to_device_uploads" => json!(self.boundary.host_to_device_uploads),
            "host_to_device_uploaded_rows" => json!(self.boundary.host_to_device_uploaded_rows),
            "host_to_device_uploaded_bytes" => json!(self.boundary.host_to_device_uploaded_bytes),
            "device_to_host_materializations" => json!(self.boundary.device_to_host_materializations),
            "device_to_host_materialized_rows" => json!(self.boundary.device_to_host_materialized_rows),
            "device_to_host_materialized_bytes" => json!(self.boundary.device_to_host_materialized_bytes),
            "device_to_device_copies" => json!(self.boundary.device_to_device_copies),
            "device_to_device_copied_bytes" => json!(self.boundary.device_to_device_copied_bytes),
            "max_task_observed_peak_bytes" => optional_u64_json(self.max_task_observed_peak_bytes),
            "max_task_observed_peak_bytes_observed" => json!(self.max_task_observed_peak_bytes.is_some()),
            "max_task_retained_bytes" => optional_u64_json(self.max_task_retained_bytes),
            "max_task_retained_bytes_observed" => json!(self.max_task_retained_bytes.is_some()),
            "native_graph_claim_calibration_rows" => json!(self.graph_claim_calibration_rows.len()),
            "native_graph_runtime_rows" => json!(self.graph_algorithm_runtime_rows.len()),
            "native_task_operator_rows" => json!(self.task_operator_metrics.len()),
            "native_task_runtime_rows" => json!(self.task_runtime_metrics.len()),
            "native_plan_task_mapping_rows" => json!(self.plan_task_mappings().len()),
            "native_dynamic_filter_runtime_rows" => json!(self.dynamic_filter_runtime_rows.len()),
            "native_source_read_runtime_rows" => json!(self.source_read_runtime_rows.len()),
            "peak_device_memory_bytes" => optional_u64_json(self.peak_device_memory_bytes),
            "resident_dataframe_peak_bytes" => optional_u64_json(self.resident_dataframe_peak_bytes),
            "resident_dataframe_peak_bytes_observed" => json!(self.resident_dataframe_peak_bytes.is_some()),
            "allocator_high_water_bytes" => optional_u64_json(self.allocator_high_water_bytes),
            "allocator_high_water_bytes_observed" => json!(self.allocator_high_water_bytes.is_some()),
            "source_chunks_planned" => json!(self.source_chunks_planned),
            "source_chunks_attempted" => json!(self.source_chunks_attempted),
            "source_chunks_committed" => json!(self.source_chunks_committed),
            "source_chunks_retried" => json!(self.source_chunks_retried),
            "source_read_prefetches_started" => json!(self.source_read_prefetches_started),
            "source_read_prefetches_completed" => json!(self.source_read_prefetches_completed),
            "source_read_prefetches_skipped" => json!(self.source_read_prefetches_skipped),
            "source_read_prefetches_abandoned" => json!(self.source_read_prefetches_abandoned),
            "row_groups_planned" => json!(self.row_groups_planned),
            "row_groups_committed" => json!(self.row_groups_committed),
            "partial_aggregate_count" => json!(self.partial_aggregate_count),
            "partial_aggregate_input_rows" => json!(self.partial_aggregate_input_rows),
            "partial_aggregate_rows" => json!(self.partial_aggregate_rows),
            "partial_aggregate_retained_bytes" => json!(self.partial_aggregate_retained_bytes),
            "final_merge_attempts" => json!(self.final_merge_attempts),
            "final_merge_levels" => json!(self.final_merge_levels),
            "final_merge_max_fan_in" => json!(self.final_merge_max_fan_in),
            "compiled_pipeline_attempts" => json!(self.compiled_pipeline_attempts),
            "compiled_pipeline_tasks_executed" => json!(self.compiled_pipeline_tasks_executed),
            "compiled_pipeline_host_arrow_sinks" => json!(self.compiled_pipeline_host_arrow_sinks),
            "parquet_files" => json!(self.parquet_files),
            "parquet_bytes_known" => json!(self.parquet_bytes_known),
            "row_groups_selected" => json!(self.row_groups_selected),
            "read_columns" => json!(self.read_columns),
            "retry_count" => json!(self.retry_count),
            "retry_resource_exhausted_count" => json!(self.retry_resource_exhausted_count),
            "dynamic_filter_produced_count" => json!(self.dynamic_filter_produced_count),
            "dynamic_filter_skipped_count" => json!(self.dynamic_filter_skipped_count),
            "dynamic_filter_applied_count" => json!(self.dynamic_filter_applied_count),
            "dynamic_filter_row_groups_pruned" => json!(self.dynamic_filter_row_groups_pruned),
            "dynamic_filter_source_bytes_pruned" => json!(self.dynamic_filter_source_bytes_pruned),
            "dynamic_filter_reader_predicate_rows_before" => {
                json!(self.dynamic_filter_reader_predicate_rows_before)
            },
            "dynamic_filter_reader_predicate_rows_after" => {
                json!(self.dynamic_filter_reader_predicate_rows_after)
            },
            "dynamic_filter_reader_predicate_rows_filtered" => {
                json!(self.dynamic_filter_reader_predicate_rows_filtered)
            },
            "dynamic_filter_post_decode_rows_before" => json!(self.dynamic_filter_post_decode_rows_before),
            "dynamic_filter_post_decode_rows_after" => json!(self.dynamic_filter_post_decode_rows_after),
            "dynamic_filter_post_decode_rows_filtered" => json!(self.dynamic_filter_post_decode_rows_filtered),
            "dynamic_filter_source_post_decode_rows_filtered" => {
                json!(self.dynamic_filter_source_post_decode_rows_filtered)
            },
            "dynamic_filter_filter_post_decode_rows_filtered" => {
                json!(self.dynamic_filter_filter_post_decode_rows_filtered)
            },
            "dynamic_filter_join_probe_post_decode_rows_filtered" => {
                json!(self.dynamic_filter_join_probe_post_decode_rows_filtered)
            },
            "hash_join_build_count" => json!(self.hash_join_build_count),
            "filtered_join_build_count" => json!(self.filtered_join_build_count),
            "one_shot_join_count" => json!(self.one_shot_join_count),
            "hash_join_hash_table_bytes" => json!(self.hash_join_hash_table_bytes),
            "filtered_join_hash_table_bytes" => json!(self.filtered_join_hash_table_bytes),
            "filtered_join_probe_state_bytes" => json!(self.filtered_join_probe_state_bytes),
            "mixed_join_size_data_bytes" => json!(self.mixed_join_size_data_bytes),
            "one_shot_join_index_bytes" => json!(self.one_shot_join_index_bytes),
            "one_shot_join_materialization_bytes" => json!(self.one_shot_join_materialization_bytes),
            "one_shot_join_output_bytes" => json!(self.one_shot_join_output_bytes),
            "join_gather_map_bytes" => json!(self.join_gather_map_bytes),
            "native_aggregate_strategy_rows" => json!(self.aggregate_strategy_decisions.len()),
            "native_join_strategy_rows" => json!(self.join_strategy_decisions.len()),
            "source_metadata_load_time" => json!(duration_ns(self.source_metadata_load_time)),
            "source_read_queue_wait_time" => json!(duration_ns(self.source_read_queue_wait_time)),
            "parquet_read_time" => json!(duration_ns(self.parquet_read_time)),
            "source_read_time" => json!(duration_ns(self.source_read_time)),
            "operator_time" => json!(duration_ns(self.operator_time)),
            "chunk_operator_time" => json!(duration_ns(self.chunk_operator_time)),
            "final_merge_time" => json!(duration_ns(self.final_merge_time)),
            "final_sort_time" => json!(duration_ns(self.final_sort_time)),
            "materialization_time" => json!(duration_ns(self.materialization_time)),
            _ => return None,
        })
    }
}

impl BudgetedExecutionEvent {
    /// Serializes this lower-level budget report as named JSON.
    pub fn to_json(&self) -> Value {
        json!({
            "query_id": optional_query_id(self.query_id.as_ref()),
            "operation": self.operation,
            "boundary": self.boundary,
            "budget_bytes": self.budget_bytes,
            "enforcement": self.enforcement,
            "status": self.status,
            "allocator_high_water_bytes": self.allocator_high_water_bytes,
            "retained_bytes": self.retained_bytes,
            "output_bytes": self.output_bytes,
            "output_rows": self.output_rows,
            "context_facts": self.context_facts,
        })
    }
}

impl OperatorMetrics {
    /// Serializes this operator metric row as named JSON.
    pub fn to_json(&self) -> Value {
        json!({
            "query_id": optional_query_id(self.query_id.as_ref()),
            "node_id": self.node_id,
            "operator_kind": self.operator_kind,
            "input_rows": self.input_rows,
            "input_bytes": self.input_bytes,
            "output_rows": self.output_rows,
            "output_bytes": self.output_bytes,
            "join": self.join.as_ref().map(JoinOperatorMetrics::to_json),
        })
    }
}

impl JoinOperatorMetrics {
    /// Serializes this join-specific operator payload as named JSON.
    pub fn to_json(&self) -> Value {
        json!({
            "query_id": optional_query_id(self.query_id.as_ref()),
            "join_type": self.join_type,
            "execution_strategy": self.execution_strategy,
            "key_count": self.key_count,
            "residual_predicate": self.residual_predicate,
            "null_equality": self.null_equality,
            "build_side": self.build_side,
            "probe_side": self.probe_side,
            "build_rows": self.build_rows,
            "build_bytes": self.build_bytes,
            "probe_rows": self.probe_rows,
            "probe_bytes": self.probe_bytes,
            "hash_table_bytes": self.hash_table_bytes,
            "filtered_probe_state_bytes": self.filtered_probe_state_bytes,
            "mixed_size_data_bytes": self.mixed_size_data_bytes,
            "gather_map_bytes": self.gather_map_bytes,
            "output_reservation_bytes": self.output_reservation_bytes,
            "output_rows": self.output_rows,
            "output_bytes": self.output_bytes,
        })
    }
}

impl GraphAlgorithmRuntimeMetrics {
    /// Serializes this graph algorithm runtime row as named JSON.
    pub fn to_json(&self) -> Value {
        json!({
            "query_id": optional_query_id(self.query_id.as_ref()),
            "node_id": self.node_id,
            "algorithm": self.algorithm,
            "runtime_outcome": self.runtime_outcome,
            "input_boundary": self.input_boundary,
            "output_boundary": self.output_boundary,
            "construction_policy": self.construction_policy,
            "source_column": self.source_column,
            "destination_column": self.destination_column,
            "weight_column": self.weight_column,
            "directed": self.directed,
            "renumber": self.renumber,
            "algorithm_parameters": self.algorithm_parameters,
            "algorithm_metadata": self.algorithm_metadata,
            "input_rows": self.input_rows,
            "input_columns": self.input_columns,
            "output_rows": self.output_rows,
            "output_columns": self.output_columns,
            "output_device_bytes": self.output_device_bytes,
            "graph_peak_bytes": self.graph_peak_bytes,
            "elapsed_ns": self.elapsed_ns,
            "gpu_resident_input_handoffs": self.gpu_resident_input_handoffs,
            "host_to_device_uploads": self.host_to_device_uploads,
            "device_to_host_materializations": self.device_to_host_materializations,
        })
    }
}

impl TaskOperatorMetrics {
    /// Serializes this per-operator task aggregate row as named JSON.
    pub fn to_json(&self) -> Value {
        json!({
            "query_id": optional_query_id(self.query_id.as_ref()),
            "operator_kind": self.operator_kind,
            "task_count": self.task_count,
            "elapsed_ns": duration_ns(self.elapsed_time),
            "max_observed_peak_bytes": self.max_observed_peak_bytes,
            "max_retained_bytes": self.max_retained_bytes,
        })
    }
}

impl TaskRuntimeMetrics {
    /// Serializes this per-compiled-task runtime row as named JSON.
    pub fn to_json(&self) -> Value {
        json!({
            "query_id": optional_query_id(self.query_id.as_ref()),
            "compiled_task_id": self.compiled_task_id,
            "operator_kind": self.operator_kind,
            "elapsed_ns": duration_ns(self.elapsed_time),
            "max_observed_peak_bytes": self.max_observed_peak_bytes,
            "max_retained_bytes": self.max_retained_bytes,
        })
    }
}

impl PlanTaskMapping {
    /// Serializes this plan-to-task mapping row as named JSON.
    pub fn to_json(&self) -> Value {
        json!({
            "query_id": self.query_id.as_str(),
            "plan_node_id": optional_plan_node_id(self.plan_node_id.as_ref()),
            "compiled_task_id": optional_compiled_task_id(self.compiled_task_id.as_ref()),
            "operator_kind": self.operator_kind,
            "datafusion_shape": self.datafusion_shape,
        })
    }
}

impl DynamicFilterRuntimeMetrics {
    /// Serializes this dynamic-filter runtime row as named JSON.
    pub fn to_json(&self) -> Value {
        json!({
            "query_id": optional_query_id(self.query_id.as_ref()),
            "dynamic_filter_id": self.dynamic_filter_id,
            "event": self.event,
            "producer_kind": self.producer_kind,
            "producer_node_id": self.producer_node_id,
            "consumer_kind": self.consumer_kind,
            "consumer_node_id": self.consumer_node_id,
            "consumer_lowering": self.consumer_lowering,
            "consumer_reduction_stage": self.consumer_reduction_stage,
            "filter_kind": self.filter_kind,
            "payload_value_count": self.payload_value_count,
            "payload_min_i64": self.payload_min_i64,
            "payload_max_i64": self.payload_max_i64,
            "producer_rows": self.producer_rows,
            "skipped_reason": self.skipped_reason,
            "selected_row_groups_before": self.selected_row_groups_before,
            "selected_row_groups_after": self.selected_row_groups_after,
            "selected_source_bytes_before": self.selected_source_bytes_before,
            "selected_source_bytes_after": self.selected_source_bytes_after,
            "source_file_index": self.source_file_index,
            "source_chunk_index": self.source_chunk_index,
            "reader_predicate_rows_before": self.reader_predicate_rows_before,
            "reader_predicate_rows_after": self.reader_predicate_rows_after,
            "reader_predicate_rows_filtered": self.reader_predicate_rows_filtered,
            "post_decode_rows_before": self.post_decode_rows_before,
            "post_decode_rows_after": self.post_decode_rows_after,
            "post_decode_rows_filtered": self.post_decode_rows_filtered,
            "row_groups_pruned": self.row_groups_pruned,
            "source_bytes_pruned": self.source_bytes_pruned,
        })
    }
}

impl SourceReadRuntimeMetrics {
    /// Serializes this source-read runtime row as named JSON.
    pub fn to_json(&self) -> Value {
        json!({
            "query_id": optional_query_id(self.query_id.as_ref()),
            "source_node_id": self.source_node_id,
            "read_kind": self.read_kind,
            "read_strategy": self.read_strategy,
            "parquet_read_evidence": self.parquet_read_evidence.to_json(),
            "batch_index": self.batch_index,
            "chunk_count": self.chunk_count,
            "first_chunk_index": self.first_chunk_index,
            "last_chunk_index": self.last_chunk_index,
            "first_file_index": self.first_file_index,
            "last_file_index": self.last_file_index,
            "file_count": self.file_count,
            "row_group_count": self.row_group_count,
            "read_columns_per_file": self.read_columns_per_file,
            "read_columns_total": self.read_columns_total,
            "known_file_bytes": self.known_file_bytes,
            "estimated_source_read_bytes": self.estimated_source_read_bytes,
            "known_row_group_bytes": self.known_row_group_bytes,
            "known_projected_compressed_bytes": self.known_projected_compressed_bytes,
            "known_projected_uncompressed_bytes": self.known_projected_uncompressed_bytes,
            "known_pure_filter_compressed_bytes": self.known_pure_filter_compressed_bytes,
            "known_pure_filter_uncompressed_bytes": self.known_pure_filter_uncompressed_bytes,
            "output_rows": self.output_rows,
            "reader_predicate": self.reader_predicate,
            "reader_dynamic_filter_count": self.reader_dynamic_filter_count,
            "source_batching_reason": self.source_batching_reason,
            "source_chunk_planning_mode": self.source_chunk_planning_mode,
            "source_chunk_planning_reason": self.source_chunk_planning_reason,
            "elapsed_ns": duration_ns(self.elapsed),
        })
    }
}

impl SourceReadParquetEvidenceMetrics {
    /// Serializes lower-layer Parquet evidence as named JSON.
    pub fn to_json(&self) -> Value {
        json!({
            "input_kind": self.input_kind,
            "remote_read_mode": self.remote_read_mode,
            "backend": self.backend,
            "backend_build_status": self.backend_build_status,
            "remote_bytes": self.remote_bytes,
            "remote_range_count": self.remote_range_count,
            "device_read_support": self.device_read_support,
            "pruning_input_row_groups": self.pruning_input_row_groups,
            "row_groups_after_stats_filter": self.row_groups_after_stats_filter,
            "row_groups_after_bloom_filter": self.row_groups_after_bloom_filter,
            "bloom_pruning_status": self.bloom_pruning_status,
        })
    }
}

impl JoinStrategyDecisionMetrics {
    /// Serializes this join-strategy diagnostic row as named JSON.
    pub fn to_json(&self) -> Value {
        json!({
            "query_id": optional_query_id(self.query_id.as_ref()),
            "node_id": self.node_id,
            "facts": facts_json(&self.facts),
        })
    }
}

impl AggregateStrategyDecisionMetrics {
    /// Serializes this aggregate-strategy diagnostic row as named JSON.
    pub fn to_json(&self) -> Value {
        json!({
            "query_id": optional_query_id(self.query_id.as_ref()),
            "node_id": self.node_id,
            "facts": facts_json(&self.facts),
        })
    }
}

fn optional_query_id(value: Option<&QueryId>) -> Option<&str> {
    value.map(QueryId::as_str)
}

fn optional_plan_node_id(value: Option<&PlanNodeId>) -> Option<&str> {
    value.map(PlanNodeId::as_str)
}

fn optional_compiled_task_id(value: Option<&CompiledTaskId>) -> Option<&str> {
    value.map(CompiledTaskId::as_str)
}

fn optional_u64_json(value: Option<u64>) -> Value {
    value.map_or(Value::Null, |value| json!(value))
}

fn duration_ns(value: std::time::Duration) -> u64 {
    value.as_nanos().min(u128::from(u64::MAX)) as u64
}

fn facts_json(facts: &BTreeMap<String, String>) -> Value {
    json!(facts)
}