velesdb-core 5.0.0

//! ORDER BY clause execution for VelesQL queries.
//!
//! Handles multi-column sorting with support for:
//! - Metadata field sorting (ASC/DESC)
//! - similarity() function sorting
//! - Arithmetic expression sorting (EPIC-042)
//! - Mixed type JSON value comparison with total ordering

use crate::collection::types::Collection;
use crate::error::Result;
use crate::point::SearchResult;
use crate::velesql::{ArithmeticExpr, ArithmeticOp};
use std::cmp::Ordering;

/// Looks up a (possibly dotted) field path inside a JSON payload.
///
/// A bare name (`"source"`) is a top-level lookup; a dotted path
/// (`"meta.source"`) walks nested objects segment by segment. Returns `None`
/// when any segment is missing or a non-object is traversed. Shared by ORDER BY
/// field comparison and score-variable payload resolution so both honor the same
/// nested-path semantics as projection/filter/aggregation.
fn get_nested_payload<'a>(
    payload: &'a serde_json::Value,
    path: &str,
) -> Option<&'a serde_json::Value> {
    let mut current = payload;
    for segment in path.split('.') {
        match current {
            serde_json::Value::Object(map) => current = map.get(segment)?,
            _ => return None,
        }
    }
    Some(current)
}

/// Compare two JSON values for sorting with total ordering.
///
/// Ordering priority (ascending): Null < Bool < Number < String < Array < Object
/// This ensures deterministic sorting even with mixed types.
#[must_use]
pub fn compare_json_values(
    a: Option<&serde_json::Value>,
    b: Option<&serde_json::Value>,
) -> Ordering {
    match (a, b) {
        (None, None) => Ordering::Equal,
        (None, Some(_)) => Ordering::Less,
        (Some(_), None) => Ordering::Greater,
        (Some(va), Some(vb)) => {
            // BUG FIX: Define total ordering for mixed JSON types
            // Type priority: Null(0) < Bool(1) < Number(2) < String(3) < Array(4) < Object(5)
            let type_rank = |v: &serde_json::Value| -> u8 {
                match v {
                    serde_json::Value::Null => 0,
                    serde_json::Value::Bool(_) => 1,
                    serde_json::Value::Number(_) => 2,
                    serde_json::Value::String(_) => 3,
                    serde_json::Value::Array(_) => 4,
                    serde_json::Value::Object(_) => 5,
                }
            };

            let rank_a = type_rank(va);
            let rank_b = type_rank(vb);

            // First compare by type rank
            if rank_a != rank_b {
                return rank_a.cmp(&rank_b);
            }

            // Same type: compare values
            match (va, vb) {
                (serde_json::Value::Number(na), serde_json::Value::Number(nb)) => {
                    let fa = na.as_f64().unwrap_or(0.0);
                    let fb = nb.as_f64().unwrap_or(0.0);
                    fa.total_cmp(&fb) // Use total_cmp for NaN safety
                }
                (serde_json::Value::String(sa), serde_json::Value::String(sb)) => sa.cmp(sb),
                (serde_json::Value::Bool(ba), serde_json::Value::Bool(bb)) => ba.cmp(bb),
                // Null vs Null, Array vs Array, Object vs Object: treat as equal
                // (comparing array/object contents would be complex and rarely needed)
                _ => Ordering::Equal,
            }
        }
    }
}

impl Collection {
    /// Apply ORDER BY clause to results.
    ///
    /// Supports multiple ORDER BY columns with stable sorting.
    /// Each column is compared in order; ties are broken by subsequent columns.
    ///
    /// # Examples
    ///
    /// ```sql
    /// SELECT * FROM collection ORDER BY category ASC, priority DESC
    /// SELECT * FROM collection ORDER BY similarity() DESC, timestamp ASC
    /// ```
    pub(crate) fn apply_order_by(
        &self,
        results: &mut [SearchResult],
        order_by: &[crate::velesql::SelectOrderBy],
        params: &std::collections::HashMap<String, serde_json::Value>,
    ) -> Result<()> {
        self.apply_order_by_with_let(results, order_by, params, &[])
    }

    /// Apply ORDER BY with pre-evaluated LET bindings (VelesQL v1.10 Phase 3).
    ///
    /// `let_bindings` contains per-result pre-evaluated `(name, value)` pairs
    /// stored as `[result_idx][(name, value)]`. If empty, behaves identically
    /// to [`apply_order_by`].
    pub(crate) fn apply_order_by_with_let(
        &self,
        results: &mut [SearchResult],
        order_by: &[crate::velesql::SelectOrderBy],
        params: &std::collections::HashMap<String, serde_json::Value>,
        per_result_let: &[Vec<(String, f32)>],
    ) -> Result<()> {
        if order_by.is_empty() {
            return Ok(());
        }

        let similarity_scores_map = self.precompute_similarity_scores(results, order_by, params)?;
        let higher_is_better = self.storage.config.read().metric.higher_is_better();

        let mut indices: Vec<usize> = (0..results.len()).collect();
        indices.sort_unstable_by(|&i, &j| {
            Self::compare_by_order_columns(
                i,
                j,
                results,
                order_by,
                &similarity_scores_map,
                higher_is_better,
                per_result_let,
            )
            // Deterministic tie-break: rows whose ORDER BY keys are all equal are
            // ordered by ascending point id, regardless of ASC/DESC on the keys.
            // This makes ORDER BY fully deterministic and — crucially — matches the
            // index-backed `ordered_ids` walk (ascending id within an equal-key
            // bucket), so the index fast path returns the same id sequence as this
            // full-scan sort (EPIC-081 phase 2).
            .then_with(|| results[i].point.id.cmp(&results[j].point.id))
        });

        // Pre-compute sorted scores before the permutation destroys the
        // original-position mapping (indices[i] = original slot for output i).
        let sorted_scores =
            extract_sorted_similarity_scores(order_by, &similarity_scores_map, &indices);

        // Apply the sort permutation in-place using cycle decomposition.
        // Avoids cloning every SearchResult (no heap allocation for vector data).
        apply_permutation_in_place(results, &mut indices);

        // Write back pre-computed scores.
        if let Some(scores) = sorted_scores {
            for (result, score) in results.iter_mut().zip(scores) {
                result.score = score;
            }
        }

        Ok(())
    }

    /// Pre-computes similarity scores for all ORDER BY similarity() columns.
    fn precompute_similarity_scores(
        &self,
        results: &[SearchResult],
        order_by: &[crate::velesql::SelectOrderBy],
        params: &std::collections::HashMap<String, serde_json::Value>,
    ) -> Result<std::collections::HashMap<usize, Vec<f32>>> {
        use crate::velesql::OrderByExpr;
        let mut map = std::collections::HashMap::new();
        for (idx, ob) in order_by.iter().enumerate() {
            match &ob.expr {
                OrderByExpr::Similarity(sim) => {
                    let order_vec = Self::resolve_vector(&sim.vector, params)?;
                    let scores =
                        self.similarity_scores_for_field(results, &sim.field, &order_vec)?;
                    map.insert(idx, scores);
                }
                OrderByExpr::SimilarityBare => {
                    // Zero-arg similarity(): use existing search scores (no recompute).
                    let scores: Vec<f32> = results.iter().map(|r| r.score).collect();
                    map.insert(idx, scores);
                }
                OrderByExpr::Field(_) | OrderByExpr::Aggregate(_) | OrderByExpr::Arithmetic(_) => {}
            }
        }
        Ok(map)
    }

    /// Scores every result row against `order_vec` using the vector selected by
    /// `field`: the default vector when `field == "vector"`, otherwise the named
    /// payload vector for that row's point id (multi-vector support, P1-A).
    ///
    /// A row whose named vector is missing, `None`, or length-mismatched scores
    /// as least similar (mirroring the MATCH-side `ORDER BY similarity()`): the
    /// metric-appropriate worst key — `NEG_INFINITY` for higher-is-better
    /// metrics, `INFINITY` (largest distance) for lower-is-better metrics — so it
    /// sorts last under DESC regardless of metric.
    fn similarity_scores_for_field(
        &self,
        results: &[SearchResult],
        field: &str,
        order_vec: &[f32],
    ) -> Result<Vec<f32>> {
        let worst = if self.storage.config.read().metric.higher_is_better() {
            f32::NEG_INFINITY
        } else {
            f32::INFINITY
        };
        results
            .iter()
            .map(|r| {
                let vec: std::borrow::Cow<[f32]> = if field == "vector" {
                    std::borrow::Cow::Borrowed(&r.point.vector)
                } else {
                    match self.get_vector_for_field(r.point.id, field)? {
                        Some(v) => std::borrow::Cow::Owned(v),
                        None => return Ok(worst),
                    }
                };
                if vec.len() != order_vec.len() || vec.is_empty() {
                    return Ok(worst);
                }
                Ok(self.compute_metric_score(&vec, order_vec))
            })
            .collect()
    }

    /// Compares two result indices across all ORDER BY columns.
    #[allow(clippy::too_many_arguments)]
    fn compare_by_order_columns(
        i: usize,
        j: usize,
        results: &[SearchResult],
        order_by: &[crate::velesql::SelectOrderBy],
        similarity_scores: &std::collections::HashMap<usize, Vec<f32>>,
        higher_is_better: bool,
        per_result_let: &[Vec<(String, f32)>],
    ) -> Ordering {
        use crate::velesql::OrderByExpr;
        for (idx, ob) in order_by.iter().enumerate() {
            let cmp = match &ob.expr {
                OrderByExpr::Similarity(_) | OrderByExpr::SimilarityBare => similarity_scores
                    .get(&idx)
                    .map_or(Ordering::Equal, |scores| scores[i].total_cmp(&scores[j])),
                OrderByExpr::Field(field_name) => {
                    Self::compare_field_expr(field_name, i, j, results, per_result_let)
                }
                OrderByExpr::Aggregate(_) => Ordering::Equal,
                // Design: Arithmetic ORDER BY uses direct numeric ordering without
                // distance-metric inversion. Users constructing custom formulas with
                // Euclidean/Hamming scores should account for lower-is-better semantics
                // in their expression (e.g., `ORDER BY -1 * vector_score + price ASC`).
                OrderByExpr::Arithmetic(expr) => {
                    Self::compare_arithmetic(expr, i, j, results, per_result_let)
                }
            };

            let is_similarity = matches!(
                &ob.expr,
                OrderByExpr::Similarity(_) | OrderByExpr::SimilarityBare
            );
            let directed_cmp =
                Self::apply_sort_direction(cmp, ob.descending, is_similarity, higher_is_better);
            if directed_cmp != Ordering::Equal {
                return directed_cmp;
            }
        }
        Ordering::Equal
    }

    /// Compares a payload field value between two results.
    fn compare_payload_field(
        field_name: &str,
        i: usize,
        j: usize,
        results: &[SearchResult],
    ) -> Ordering {
        let val_i = results[i]
            .point
            .payload
            .as_ref()
            .and_then(|p| get_nested_payload(p, field_name));
        let val_j = results[j]
            .point
            .payload
            .as_ref()
            .and_then(|p| get_nested_payload(p, field_name));
        compare_json_values(val_i, val_j)
    }

    /// Compares a field name, checking LET bindings first, then payload.
    fn compare_field_or_let(
        field_name: &str,
        i: usize,
        j: usize,
        results: &[SearchResult],
        per_result_let: &[Vec<(String, f32)>],
    ) -> Ordering {
        if let (Some(let_i), Some(let_j)) = (per_result_let.get(i), per_result_let.get(j)) {
            if let Some(vi) = let_i.iter().find(|(k, _)| k == field_name).map(|(_, v)| *v) {
                let vj = let_j
                    .iter()
                    .find(|(k, _)| k == field_name)
                    .map_or(0.0, |(_, v)| *v);
                return vi.total_cmp(&vj);
            }
        }
        Self::compare_payload_field(field_name, i, j, results)
    }

    /// Compares a bare `ORDER BY` field. A built-in score variable
    /// (`sparse_score`, `bm25_score`, …) resolves from each result's component
    /// scores — exactly as the arithmetic ORDER BY path does — so
    /// `ORDER BY sparse_score DESC` actually ranks instead of silently no-op'ing
    /// as an absent payload field. Anything else is a LET binding or payload.
    fn compare_field_expr(
        field_name: &str,
        i: usize,
        j: usize,
        results: &[SearchResult],
        per_result_let: &[Vec<(String, f32)>],
    ) -> Ordering {
        if is_builtin_score_variable(field_name) {
            Self::compare_score_variable(field_name, i, j, results, per_result_let)
        } else {
            Self::compare_field_or_let(field_name, i, j, results, per_result_let)
        }
    }

    /// Compares two results by a built-in score variable, resolving each via the
    /// same `ScoreContext::resolve_variable` the arithmetic path uses (LET first,
    /// then the component breakdown with the absent-component default).
    fn compare_score_variable(
        name: &str,
        i: usize,
        j: usize,
        results: &[SearchResult],
        per_result_let: &[Vec<(String, f32)>],
    ) -> Ordering {
        let (ctx_i, ctx_j) = Self::score_context_pair(i, j, results, per_result_let);
        ctx_i
            .resolve_variable(name)
            .total_cmp(&ctx_j.resolve_variable(name))
    }

    /// Compares two results by an arithmetic expression with full context.
    fn compare_arithmetic(
        expr: &crate::velesql::ArithmeticExpr,
        i: usize,
        j: usize,
        results: &[SearchResult],
        per_result_let: &[Vec<(String, f32)>],
    ) -> Ordering {
        let (ctx_i, ctx_j) = Self::score_context_pair(i, j, results, per_result_let);
        let val_i = evaluate_arithmetic(expr, &ctx_i);
        let val_j = evaluate_arithmetic(expr, &ctx_j);
        val_i.total_cmp(&val_j)
    }

    /// Builds the `ScoreContext` pair for comparing results `i` and `j`, each
    /// with its own LET bindings, score, payload, and component breakdown.
    fn score_context_pair<'a>(
        i: usize,
        j: usize,
        results: &'a [SearchResult],
        per_result_let: &'a [Vec<(String, f32)>],
    ) -> (ScoreContext<'a>, ScoreContext<'a>) {
        let ctx_i = ScoreContext::with_let_bindings(
            results[i].score,
            results[i].point.payload.as_ref(),
            results[i].component_scores.as_deref(),
            per_result_let.get(i).map(Vec::as_slice),
        );
        let ctx_j = ScoreContext::with_let_bindings(
            results[j].score,
            results[j].point.payload.as_ref(),
            results[j].component_scores.as_deref(),
            per_result_let.get(j).map(Vec::as_slice),
        );
        (ctx_i, ctx_j)
    }

    /// Applies ASC/DESC direction, accounting for distance metric inversion.
    fn apply_sort_direction(
        cmp: Ordering,
        descending: bool,
        is_similarity: bool,
        higher_is_better: bool,
    ) -> Ordering {
        if descending {
            if is_similarity && !higher_is_better {
                cmp
            } else {
                cmp.reverse()
            }
        } else if is_similarity && !higher_is_better {
            cmp.reverse()
        } else {
            cmp
        }
    }
}

/// Whether `name` is a built-in score variable resolved from the component-score
/// breakdown (rather than a payload field), matching the names
/// `ScoreContext::resolve_variable` recognizes. Bare such names in `ORDER BY`
/// route through score resolution so e.g. `ORDER BY sparse_score DESC` ranks.
fn is_builtin_score_variable(name: &str) -> bool {
    matches!(
        name,
        "vector_score" | "graph_score" | "bm25_score" | "sparse_score" | "fused_score"
    )
}

/// Applies a permutation to `slice` in-place using cycle decomposition.
///
/// `perm[i]` is the original index of the element that belongs at position `i`
/// in the output. Each cycle in the permutation is resolved by swapping elements
/// along the cycle, which requires only O(n) `usize` bookkeeping with no heap
/// allocation for the elements themselves.
///
/// `perm` is consumed (all entries set to their own index) as a side-effect.
fn apply_permutation_in_place<T>(slice: &mut [T], perm: &mut [usize]) {
    debug_assert_eq!(slice.len(), perm.len());
    for i in 0..perm.len() {
        let mut j = i;
        while perm[j] != i {
            let k = perm[j];
            slice.swap(j, k);
            perm[j] = j;
            j = k;
        }
        perm[j] = j;
    }
}

/// Returns sorted similarity scores in output order, or `None` when the query
/// has no similarity ORDER BY column.
///
/// Must be called before `apply_permutation_in_place` because that call
/// overwrites `indices` (the original-position mapping).
fn extract_sorted_similarity_scores(
    order_by: &[crate::velesql::SelectOrderBy],
    similarity_scores_map: &std::collections::HashMap<usize, Vec<f32>>,
    indices: &[usize],
) -> Option<Vec<f32>> {
    use crate::velesql::OrderByExpr;
    let sim_idx = order_by
        .iter()
        .enumerate()
        .find(|(_, ob)| {
            matches!(
                ob.expr,
                OrderByExpr::Similarity(_) | OrderByExpr::SimilarityBare
            )
        })
        .map(|(idx, _)| idx)?;

    let scores = similarity_scores_map.get(&sim_idx)?;
    Some(indices.iter().map(|&orig| scores[orig]).collect())
}

/// Context for evaluating arithmetic ORDER BY expressions (EPIC-042).
///
/// Holds the pre-computed search score, optional per-component score breakdown,
/// LET bindings (v1.10), and optional payload for variable resolution.
pub(crate) struct ScoreContext<'a> {
    /// Pre-computed search score (vector similarity or fused score).
    search_score: f32,
    /// Payload fields for variable resolution.
    payload: Option<&'a serde_json::Value>,
    /// Optional per-component scores from hybrid search (v1.10+).
    ///
    /// When present, built-in score variables (`vector_score`, `bm25_score`, etc.)
    /// resolve to their individual component values instead of the fused score.
    component_scores: Option<&'a [(&'static str, f32)]>,
    /// Pre-evaluated LET binding values (VelesQL v1.10 Phase 3).
    ///
    /// Resolution priority: LET bindings > component_scores > search_score > payload.
    let_bindings: Option<&'a [(String, f32)]>,
}

impl<'a> ScoreContext<'a> {
    /// Creates a new score context without component scores (backward compat).
    ///
    /// Used by tests and simple code paths where component scores are not available.
    #[allow(dead_code)] // Used by ordering_tests and component_scores_tests.
    pub(crate) fn new(search_score: f32, payload: Option<&'a serde_json::Value>) -> Self {
        Self {
            search_score,
            payload,
            component_scores: None,
            let_bindings: None,
        }
    }

    /// Creates a score context with per-component score breakdown.
    #[allow(dead_code)] // Used by component_scores_tests.
    pub(crate) fn with_components(
        search_score: f32,
        payload: Option<&'a serde_json::Value>,
        component_scores: Option<&'a [(&'static str, f32)]>,
    ) -> Self {
        Self {
            search_score,
            payload,
            component_scores,
            let_bindings: None,
        }
    }

    /// Creates a score context with LET bindings and component scores.
    pub(crate) fn with_let_bindings(
        search_score: f32,
        payload: Option<&'a serde_json::Value>,
        component_scores: Option<&'a [(&'static str, f32)]>,
        let_bindings: Option<&'a [(String, f32)]>,
    ) -> Self {
        Self {
            search_score,
            payload,
            component_scores,
            let_bindings,
        }
    }

    /// Resolves a variable name to a numeric value.
    ///
    /// Resolution priority:
    /// 1. LET bindings (highest — user-defined score aliases).
    /// 2. Built-in component scores (`vector_score`, `bm25_score`, etc.).
    /// 3. For a built-in absent from a *tagged* result (component_scores `Some`),
    ///    the component defaults to `0.0` (per VELESQL_SPEC); on an *untagged*
    ///    result (component_scores `None`) it falls back to `search_score`.
    /// 4. Payload fields for non-built-in names.
    ///
    /// `fused_score` and `similarity` always resolve to `search_score` (they
    /// represent the combined result, not an individual component).
    ///
    /// `graph_score` is reserved: it is never populated as a component, so it
    /// resolves to `0.0` on tagged results (and to `search_score` only on the
    /// legacy untagged path).
    fn resolve_variable(&self, name: &str) -> f32 {
        // Priority 1: LET bindings override everything.
        if let Some(val) = self.lookup_let_binding(name) {
            return val;
        }
        match name {
            // fused_score and similarity always use the primary fused score.
            "fused_score" | "similarity" => self.search_score,
            // Component-aware built-ins: a tagged result (component_scores Some)
            // defaults an absent component to 0.0; only the untagged legacy path
            // falls back to the primary search_score.
            "vector_score" | "graph_score" | "bm25_score" | "sparse_score" => self
                .lookup_component(name)
                .unwrap_or_else(|| self.absent_component_default()),
            _ => self.resolve_payload_variable(name),
        }
    }

    /// Default for a built-in component absent from this result: `0.0` when the
    /// result carries a component breakdown (`component_scores` is `Some`),
    /// otherwise the primary `search_score` (legacy untagged compatibility).
    fn absent_component_default(&self) -> f32 {
        if self.component_scores.is_some() {
            0.0
        } else {
            self.search_score
        }
    }

    /// Looks up a named LET binding, returning `None` if absent.
    fn lookup_let_binding(&self, name: &str) -> Option<f32> {
        self.let_bindings?
            .iter()
            .find(|(k, _)| k == name)
            .map(|(_, v)| *v)
    }

    /// Looks up a named component score, returning `None` if absent.
    fn lookup_component(&self, name: &str) -> Option<f32> {
        self.component_scores?
            .iter()
            .find(|(k, _)| *k == name)
            .map(|(_, v)| *v)
    }

    /// Resolves a variable name from the payload.
    fn resolve_payload_variable(&self, name: &str) -> f32 {
        self.payload
            .and_then(|p| get_nested_payload(p, name))
            .and_then(serde_json::Value::as_f64)
            .map_or(0.0, |v| {
                #[allow(clippy::cast_possible_truncation)]
                // Reason: payload values are user-defined scores; f64→f32 precision loss is acceptable.
                {
                    v as f32
                }
            })
    }
}

/// Evaluates all LET bindings in declaration order for a single result.
///
/// Each binding can reference earlier bindings, component scores, or the
/// search score. The returned vec contains `(name, value)` pairs in order.
pub(crate) fn evaluate_let_bindings(
    bindings: &[crate::velesql::LetBinding],
    search_score: f32,
    payload: Option<&serde_json::Value>,
    component_scores: Option<&[(&'static str, f32)]>,
) -> Vec<(String, f32)> {
    let mut evaluated: Vec<(String, f32)> = Vec::with_capacity(bindings.len());
    for binding in bindings {
        let ctx = ScoreContext::with_let_bindings(
            search_score,
            payload,
            component_scores,
            Some(&evaluated),
        );
        let value = evaluate_arithmetic(&binding.expr, &ctx);
        evaluated.push((binding.name.clone(), value));
    }
    evaluated
}

/// Maximum recursion depth for arithmetic expression evaluation.
/// Matches `DEFAULT_MAX_AST_DEPTH` (64) from validation.
const MAX_ARITHMETIC_DEPTH: u8 = 64;

/// Evaluates an arithmetic expression against a score context (EPIC-042).
///
/// Division by zero returns `0.0` (safe default for sorting).
/// Recursion depth is capped at [`MAX_ARITHMETIC_DEPTH`] to prevent stack overflow.
pub(crate) fn evaluate_arithmetic(expr: &ArithmeticExpr, ctx: &ScoreContext<'_>) -> f32 {
    evaluate_arithmetic_inner(expr, ctx, 0)
}

/// Inner recursive evaluator with depth tracking.
fn evaluate_arithmetic_inner(expr: &ArithmeticExpr, ctx: &ScoreContext<'_>, depth: u8) -> f32 {
    if depth >= MAX_ARITHMETIC_DEPTH {
        return 0.0;
    }
    match expr {
        ArithmeticExpr::Literal(v) => {
            #[allow(clippy::cast_possible_truncation)]
            // Reason: arithmetic literals are user-defined weights; f64→f32 precision loss is acceptable.
            {
                *v as f32
            }
        }
        ArithmeticExpr::Variable(name) => ctx.resolve_variable(name),
        // Only bare similarity() passes validation inside arithmetic expressions.
        // Parameterized similarity(field, $vec) is rejected at validation time (V008).
        ArithmeticExpr::Similarity(_) => ctx.search_score,
        ArithmeticExpr::BinaryOp { left, op, right } => {
            let l = evaluate_arithmetic_inner(left, ctx, depth + 1);
            let r = evaluate_arithmetic_inner(right, ctx, depth + 1);
            match op {
                ArithmeticOp::Add => l + r,
                ArithmeticOp::Sub => l - r,
                ArithmeticOp::Mul => l * r,
                ArithmeticOp::Div => {
                    if r.abs() > f32::EPSILON {
                        l / r
                    } else {
                        0.0
                    }
                }
            }
        }
    }
}