jmespath_extensions 0.9.0

//! Array manipulation functions.
//!
//! This module provides array functions for JMESPath queries.
//!
//! For complete function reference with signatures and examples, see the
//! [`functions`](crate::functions) module documentation or use `jpx --list-category array`.
//!
//! # Example
//!
//! ```rust
//! use jmespath::{Runtime, Variable};
//! use jmespath_extensions::array;
//!
//! let mut runtime = Runtime::new();
//! runtime.register_builtin_functions();
//! array::register(&mut runtime);
//! ```

use std::collections::HashSet;
use std::rc::Rc;

use crate::common::{
    ArgumentType, Context, ErrorReason, Function, JmespathError, Rcvar, Runtime, Variable,
};
use crate::define_function;
use crate::register_if_enabled;

/// Register all array functions with the runtime.
pub fn register(runtime: &mut Runtime) {
    runtime.register_function("unique", Box::new(UniqueFn::new()));
    runtime.register_function("zip", Box::new(ZipFn::new()));
    runtime.register_function("chunk", Box::new(ChunkFn::new()));
    runtime.register_function("take", Box::new(TakeFn::new()));
    runtime.register_function("drop", Box::new(DropFn::new()));
    runtime.register_function("flatten_deep", Box::new(FlattenDeepFn::new()));
    runtime.register_function("flatten", Box::new(FlattenFn::new()));
    runtime.register_function("compact", Box::new(CompactFn::new()));
    runtime.register_function("range", Box::new(RangeFn::new()));
    runtime.register_function("index_at", Box::new(IndexAtFn::new()));
    runtime.register_function("includes", Box::new(IncludesFn::new()));
    runtime.register_function("find_index", Box::new(FindIndexFn::new()));
    runtime.register_function("first", Box::new(FirstFn::new()));
    runtime.register_function("last", Box::new(LastFn::new()));
    runtime.register_function("group_by", Box::new(GroupByFn::new()));
    runtime.register_function("index_by", Box::new(IndexByFn::new()));
    runtime.register_function("nth", Box::new(NthFn::new()));
    runtime.register_function("interleave", Box::new(InterleaveFn::new()));
    runtime.register_function("rotate", Box::new(RotateFn::new()));
    runtime.register_function("partition", Box::new(PartitionFn::new()));
    runtime.register_function("difference", Box::new(DifferenceFn::new()));
    runtime.register_function("intersection", Box::new(IntersectionFn::new()));
    runtime.register_function("union", Box::new(UnionFn::new()));
    runtime.register_function("frequencies", Box::new(FrequenciesFn::new()));
    runtime.register_function("mode", Box::new(ModeFn::new()));
    runtime.register_function("cartesian", Box::new(CartesianFn::new()));
    runtime.register_function("initial", Box::new(InitialFn::new()));
    // Alias for initial (Clojure-style name)
    runtime.register_function("butlast", Box::new(InitialFn::new()));
    // Clojure-inspired functions
    runtime.register_function("interpose", Box::new(InterposeFn::new()));
    runtime.register_function("zipmap", Box::new(ZipmapFn::new()));
    runtime.register_function("partition_by", Box::new(PartitionByFn::new()));
    runtime.register_function("dedupe", Box::new(DedupeFn::new()));
    runtime.register_function("tail", Box::new(TailFn::new()));
    runtime.register_function("without", Box::new(WithoutFn::new()));
    runtime.register_function("xor", Box::new(XorFn::new()));
    runtime.register_function("fill", Box::new(FillFn::new()));
    runtime.register_function("pull_at", Box::new(PullAtFn::new()));
    runtime.register_function("window", Box::new(WindowFn::new()));
    runtime.register_function("combinations", Box::new(CombinationsFn::new()));
    runtime.register_function("transpose", Box::new(TransposeFn::new()));
    runtime.register_function("pairwise", Box::new(PairwiseFn::new()));
    // Alias for window (sliding_window is a common name)
    runtime.register_function("sliding_window", Box::new(WindowFn::new()));
    // jq-parity functions
    runtime.register_function("indices_array", Box::new(IndicesArrayFn::new()));
    runtime.register_function("inside_array", Box::new(InsideArrayFn::new()));
    runtime.register_function("bsearch", Box::new(BsearchFn::new()));
    // Clojure-inspired functions (Phase 3)
    runtime.register_function("repeat_array", Box::new(RepeatArrayFn::new()));
    runtime.register_function("cycle", Box::new(CycleFn::new()));
}

/// Register only the array functions that are in the enabled set.
pub fn register_filtered(runtime: &mut Runtime, enabled: &HashSet<&str>) {
    register_if_enabled!(runtime, enabled, "unique", Box::new(UniqueFn::new()));
    register_if_enabled!(runtime, enabled, "zip", Box::new(ZipFn::new()));
    register_if_enabled!(runtime, enabled, "chunk", Box::new(ChunkFn::new()));
    register_if_enabled!(runtime, enabled, "take", Box::new(TakeFn::new()));
    register_if_enabled!(runtime, enabled, "drop", Box::new(DropFn::new()));
    register_if_enabled!(
        runtime,
        enabled,
        "flatten_deep",
        Box::new(FlattenDeepFn::new())
    );
    register_if_enabled!(runtime, enabled, "flatten", Box::new(FlattenFn::new()));
    register_if_enabled!(runtime, enabled, "compact", Box::new(CompactFn::new()));
    register_if_enabled!(runtime, enabled, "range", Box::new(RangeFn::new()));
    register_if_enabled!(runtime, enabled, "index_at", Box::new(IndexAtFn::new()));
    register_if_enabled!(runtime, enabled, "includes", Box::new(IncludesFn::new()));
    register_if_enabled!(runtime, enabled, "find_index", Box::new(FindIndexFn::new()));
    register_if_enabled!(runtime, enabled, "first", Box::new(FirstFn::new()));
    register_if_enabled!(runtime, enabled, "last", Box::new(LastFn::new()));
    register_if_enabled!(runtime, enabled, "group_by", Box::new(GroupByFn::new()));
    register_if_enabled!(runtime, enabled, "index_by", Box::new(IndexByFn::new()));
    register_if_enabled!(runtime, enabled, "nth", Box::new(NthFn::new()));
    register_if_enabled!(
        runtime,
        enabled,
        "interleave",
        Box::new(InterleaveFn::new())
    );
    register_if_enabled!(runtime, enabled, "rotate", Box::new(RotateFn::new()));
    register_if_enabled!(runtime, enabled, "partition", Box::new(PartitionFn::new()));
    register_if_enabled!(
        runtime,
        enabled,
        "difference",
        Box::new(DifferenceFn::new())
    );
    register_if_enabled!(
        runtime,
        enabled,
        "intersection",
        Box::new(IntersectionFn::new())
    );
    register_if_enabled!(runtime, enabled, "union", Box::new(UnionFn::new()));
    register_if_enabled!(
        runtime,
        enabled,
        "frequencies",
        Box::new(FrequenciesFn::new())
    );
    register_if_enabled!(runtime, enabled, "mode", Box::new(ModeFn::new()));
    register_if_enabled!(runtime, enabled, "cartesian", Box::new(CartesianFn::new()));
    register_if_enabled!(runtime, enabled, "initial", Box::new(InitialFn::new()));
    // Alias for initial (Clojure-style name)
    register_if_enabled!(runtime, enabled, "butlast", Box::new(InitialFn::new()));
    // Clojure-inspired functions
    register_if_enabled!(runtime, enabled, "interpose", Box::new(InterposeFn::new()));
    register_if_enabled!(runtime, enabled, "zipmap", Box::new(ZipmapFn::new()));
    register_if_enabled!(
        runtime,
        enabled,
        "partition_by",
        Box::new(PartitionByFn::new())
    );
    register_if_enabled!(runtime, enabled, "dedupe", Box::new(DedupeFn::new()));
    register_if_enabled!(runtime, enabled, "tail", Box::new(TailFn::new()));
    register_if_enabled!(runtime, enabled, "without", Box::new(WithoutFn::new()));
    register_if_enabled!(runtime, enabled, "xor", Box::new(XorFn::new()));
    register_if_enabled!(runtime, enabled, "fill", Box::new(FillFn::new()));
    register_if_enabled!(runtime, enabled, "pull_at", Box::new(PullAtFn::new()));
    register_if_enabled!(runtime, enabled, "window", Box::new(WindowFn::new()));
    register_if_enabled!(
        runtime,
        enabled,
        "combinations",
        Box::new(CombinationsFn::new())
    );
    register_if_enabled!(runtime, enabled, "transpose", Box::new(TransposeFn::new()));
    register_if_enabled!(runtime, enabled, "pairwise", Box::new(PairwiseFn::new()));
    // Alias for window (sliding_window is a common name)
    register_if_enabled!(
        runtime,
        enabled,
        "sliding_window",
        Box::new(WindowFn::new())
    );
    // jq-parity functions
    register_if_enabled!(
        runtime,
        enabled,
        "indices_array",
        Box::new(IndicesArrayFn::new())
    );
    register_if_enabled!(
        runtime,
        enabled,
        "inside_array",
        Box::new(InsideArrayFn::new())
    );
    register_if_enabled!(runtime, enabled, "bsearch", Box::new(BsearchFn::new()));
    // Clojure-inspired functions (Phase 3)
    register_if_enabled!(
        runtime,
        enabled,
        "repeat_array",
        Box::new(RepeatArrayFn::new())
    );
    register_if_enabled!(runtime, enabled, "cycle", Box::new(CycleFn::new()));
}

// =============================================================================
// unique(array) -> array
// =============================================================================

define_function!(UniqueFn, vec![ArgumentType::Array], None);

impl Function for UniqueFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let arr = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        let mut seen = HashSet::new();
        let mut result = Vec::new();

        for item in arr {
            let key = serde_json::to_string(&**item).unwrap_or_default();
            if seen.insert(key) {
                result.push(item.clone());
            }
        }

        Ok(Rc::new(Variable::Array(result)))
    }
}

// =============================================================================
// zip(array1, array2) -> array of pairs
// =============================================================================

define_function!(ZipFn, vec![ArgumentType::Array, ArgumentType::Array], None);

impl Function for ZipFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let arr1 = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        let arr2 = args[1].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        let result: Vec<Rcvar> = arr1
            .iter()
            .zip(arr2.iter())
            .map(|(a, b)| Rc::new(Variable::Array(vec![a.clone(), b.clone()])) as Rcvar)
            .collect();

        Ok(Rc::new(Variable::Array(result)))
    }
}

// =============================================================================
// chunk(array, size) -> array of arrays
// =============================================================================

define_function!(
    ChunkFn,
    vec![ArgumentType::Array, ArgumentType::Number],
    None
);

impl Function for ChunkFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let arr = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        let size = args[1].as_number().map(|n| n as usize).ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected positive number for size".to_owned()),
            )
        })?;

        if size == 0 {
            return Ok(Rc::new(Variable::Array(vec![])));
        }

        let chunks: Vec<Rcvar> = arr
            .chunks(size)
            .map(|chunk| Rc::new(Variable::Array(chunk.to_vec())) as Rcvar)
            .collect();

        Ok(Rc::new(Variable::Array(chunks)))
    }
}

// =============================================================================
// take(array, n) -> array (first n elements)
// =============================================================================

define_function!(
    TakeFn,
    vec![ArgumentType::Array, ArgumentType::Number],
    None
);

impl Function for TakeFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let arr = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        let n = args[1].as_number().map(|n| n as usize).ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected positive number".to_owned()),
            )
        })?;

        let result: Vec<Rcvar> = arr.iter().take(n).cloned().collect();

        Ok(Rc::new(Variable::Array(result)))
    }
}

// =============================================================================
// drop(array, n) -> array (skip first n elements)
// =============================================================================

define_function!(
    DropFn,
    vec![ArgumentType::Array, ArgumentType::Number],
    None
);

impl Function for DropFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let arr = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        let n = args[1].as_number().map(|n| n as usize).ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected positive number".to_owned()),
            )
        })?;

        let result: Vec<Rcvar> = arr.iter().skip(n).cloned().collect();

        Ok(Rc::new(Variable::Array(result)))
    }
}

// =============================================================================
// flatten_deep(array) -> array (recursively flatten)
// =============================================================================

define_function!(FlattenDeepFn, vec![ArgumentType::Array], None);

fn flatten_recursive(arr: &[Rcvar]) -> Vec<Rcvar> {
    let mut result = Vec::new();
    for item in arr {
        if let Some(inner) = item.as_array() {
            result.extend(flatten_recursive(inner));
        } else {
            result.push(item.clone());
        }
    }
    result
}

impl Function for FlattenDeepFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let arr = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        Ok(Rc::new(Variable::Array(flatten_recursive(arr))))
    }
}

// =============================================================================
// flatten(array) -> array (single-level flatten)
// =============================================================================

define_function!(FlattenFn, vec![ArgumentType::Array], None);

impl Function for FlattenFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let arr = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        let mut result = Vec::new();
        for item in arr {
            if let Some(inner) = item.as_array() {
                result.extend(inner.iter().cloned());
            } else {
                result.push(item.clone());
            }
        }

        Ok(Rc::new(Variable::Array(result)))
    }
}

// =============================================================================
// compact(array) -> array (remove null/false values)
// =============================================================================

define_function!(CompactFn, vec![ArgumentType::Array], None);

impl Function for CompactFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let arr = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        let result: Vec<Rcvar> = arr
            .iter()
            .filter(|v| !v.is_null() && !matches!(&***v, Variable::Bool(false)))
            .cloned()
            .collect();

        Ok(Rc::new(Variable::Array(result)))
    }
}

// =============================================================================
// range(start, end, step?) -> array
// =============================================================================

define_function!(
    RangeFn,
    vec![ArgumentType::Number, ArgumentType::Number],
    Some(ArgumentType::Number)
);

impl Function for RangeFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let start = args[0].as_number().map(|n| n as i64).ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected start number".to_owned()),
            )
        })?;

        let end = args[1].as_number().map(|n| n as i64).ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected end number".to_owned()),
            )
        })?;

        let step = if args.len() > 2 {
            args[2].as_number().map(|n| n as i64).ok_or_else(|| {
                JmespathError::new(
                    ctx.expression,
                    0,
                    ErrorReason::Parse("Expected step number".to_owned()),
                )
            })?
        } else {
            1
        };

        if step == 0 {
            return Err(JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Step cannot be zero".to_owned()),
            ));
        }

        let mut result = Vec::new();
        let mut current = start;

        const MAX_RANGE: usize = 10000;

        if step > 0 {
            while current < end && result.len() < MAX_RANGE {
                result.push(Rc::new(Variable::Number(serde_json::Number::from(current))) as Rcvar);
                current += step;
            }
        } else {
            while current > end && result.len() < MAX_RANGE {
                result.push(Rc::new(Variable::Number(serde_json::Number::from(current))) as Rcvar);
                current += step;
            }
        }

        Ok(Rc::new(Variable::Array(result)))
    }
}

// =============================================================================
// index_at(array, index) -> element (supports negative index)
// =============================================================================

define_function!(
    IndexAtFn,
    vec![ArgumentType::Array, ArgumentType::Number],
    None
);

impl Function for IndexAtFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let arr = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        let index = args[1].as_number().map(|n| n as i64).ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected number for index".to_owned()),
            )
        })?;

        let len = arr.len() as i64;
        let actual_index = if index < 0 {
            (len + index) as usize
        } else {
            index as usize
        };

        if actual_index < arr.len() {
            Ok(arr[actual_index].clone())
        } else {
            Ok(Rc::new(Variable::Null))
        }
    }
}

// =============================================================================
// includes(array, value) -> boolean
// =============================================================================

define_function!(
    IncludesFn,
    vec![ArgumentType::Array, ArgumentType::Any],
    None
);

impl Function for IncludesFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let arr = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        let search_key = serde_json::to_string(&*args[1]).unwrap_or_default();

        let found = arr.iter().any(|item| {
            let item_key = serde_json::to_string(&**item).unwrap_or_default();
            item_key == search_key
        });

        Ok(Rc::new(Variable::Bool(found)))
    }
}

// =============================================================================
// find_index(array, value) -> number (-1 if not found)
// =============================================================================

define_function!(
    FindIndexFn,
    vec![ArgumentType::Array, ArgumentType::Any],
    None
);

impl Function for FindIndexFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let arr = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        let search_key = serde_json::to_string(&*args[1]).unwrap_or_default();

        let index = arr
            .iter()
            .position(|item| {
                let item_key = serde_json::to_string(&**item).unwrap_or_default();
                item_key == search_key
            })
            .map(|i| i as i64)
            .unwrap_or(-1);

        Ok(Rc::new(Variable::Number(serde_json::Number::from(index))))
    }
}

// =============================================================================
// first(array) -> any (first element or null)
// =============================================================================

define_function!(FirstFn, vec![ArgumentType::Array], None);

impl Function for FirstFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let arr = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        Ok(arr
            .first()
            .cloned()
            .unwrap_or_else(|| Rc::new(Variable::Null)))
    }
}

// =============================================================================
// last(array) -> any (last element or null)
// =============================================================================

define_function!(LastFn, vec![ArgumentType::Array], None);

impl Function for LastFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let arr = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        Ok(arr
            .last()
            .cloned()
            .unwrap_or_else(|| Rc::new(Variable::Null)))
    }
}

// =============================================================================
// group_by(array, field_name) -> object
// =============================================================================

define_function!(
    GroupByFn,
    vec![ArgumentType::Array, ArgumentType::String],
    None
);

impl Function for GroupByFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let arr = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        let field_name = args[1].as_string().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected field name string".to_owned()),
            )
        })?;

        let mut groups: std::collections::BTreeMap<String, Vec<Rcvar>> =
            std::collections::BTreeMap::new();

        for item in arr {
            let key = if let Some(obj) = item.as_object() {
                if let Some(field_value) = obj.get(field_name) {
                    match &**field_value {
                        Variable::String(s) => s.clone(),
                        Variable::Number(n) => n.to_string(),
                        Variable::Bool(b) => b.to_string(),
                        Variable::Null => "null".to_string(),
                        _ => continue,
                    }
                } else {
                    "null".to_string()
                }
            } else {
                continue;
            };
            groups.entry(key).or_default().push(item.clone());
        }

        let result: std::collections::BTreeMap<String, Rcvar> = groups
            .into_iter()
            .map(|(k, v)| (k, Rc::new(Variable::Array(v)) as Rcvar))
            .collect();

        Ok(Rc::new(Variable::Object(result)))
    }
}

// =============================================================================
// index_by(array, field_name) -> object (last value wins for duplicates)
// =============================================================================

define_function!(
    IndexByFn,
    vec![ArgumentType::Array, ArgumentType::String],
    None
);

impl Function for IndexByFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let arr = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        let field_name = args[1].as_string().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected field name string".to_owned()),
            )
        })?;

        let mut result: std::collections::BTreeMap<String, Rcvar> =
            std::collections::BTreeMap::new();

        for item in arr {
            let key = if let Some(obj) = item.as_object() {
                if let Some(field_value) = obj.get(field_name) {
                    match &**field_value {
                        Variable::String(s) => s.clone(),
                        Variable::Number(n) => n.to_string(),
                        Variable::Bool(b) => b.to_string(),
                        Variable::Null => "null".to_string(),
                        _ => continue,
                    }
                } else {
                    // Skip items without the key field
                    continue;
                }
            } else {
                continue;
            };
            // Last value wins for duplicate keys
            result.insert(key, item.clone());
        }

        Ok(Rc::new(Variable::Object(result)))
    }
}

// =============================================================================
// nth(array, n) -> array (every nth element)
// =============================================================================

define_function!(NthFn, vec![ArgumentType::Array, ArgumentType::Number], None);

impl Function for NthFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let arr = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        let n = args[1].as_number().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected number argument".to_owned()),
            )
        })? as usize;

        if n == 0 {
            return Ok(Rc::new(Variable::Null));
        }

        let result: Vec<Rcvar> = arr.iter().step_by(n).cloned().collect();
        Ok(Rc::new(Variable::Array(result)))
    }
}

// =============================================================================
// interleave(array1, array2) -> array (alternate elements)
// =============================================================================

define_function!(
    InterleaveFn,
    vec![ArgumentType::Array, ArgumentType::Array],
    None
);

impl Function for InterleaveFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let arr1 = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        let arr2 = args[1].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        let mut result = Vec::with_capacity(arr1.len() + arr2.len());
        let mut iter1 = arr1.iter();
        let mut iter2 = arr2.iter();

        loop {
            match (iter1.next(), iter2.next()) {
                (Some(a), Some(b)) => {
                    result.push(a.clone());
                    result.push(b.clone());
                }
                (Some(a), None) => {
                    result.push(a.clone());
                    result.extend(iter1.cloned());
                    break;
                }
                (None, Some(b)) => {
                    result.push(b.clone());
                    result.extend(iter2.cloned());
                    break;
                }
                (None, None) => break,
            }
        }

        Ok(Rc::new(Variable::Array(result)))
    }
}

// =============================================================================
// rotate(array, n) -> array (rotate elements by n positions)
// =============================================================================

define_function!(
    RotateFn,
    vec![ArgumentType::Array, ArgumentType::Number],
    None
);

impl Function for RotateFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let arr = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        if arr.is_empty() {
            return Ok(Rc::new(Variable::Array(vec![])));
        }

        let n = args[1].as_number().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected number argument".to_owned()),
            )
        })? as i64;

        let len = arr.len() as i64;
        let rotation = ((n % len) + len) % len;
        let rotation = rotation as usize;

        let mut result = Vec::with_capacity(arr.len());
        result.extend(arr[rotation..].iter().cloned());
        result.extend(arr[..rotation].iter().cloned());

        Ok(Rc::new(Variable::Array(result)))
    }
}

// =============================================================================
// partition(array, n) -> array (split into n equal parts)
// =============================================================================

define_function!(
    PartitionFn,
    vec![ArgumentType::Array, ArgumentType::Number],
    None
);

impl Function for PartitionFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let arr = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        let n = args[1].as_number().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected number argument".to_owned()),
            )
        })? as usize;

        if n == 0 {
            return Ok(Rc::new(Variable::Null));
        }

        let len = arr.len();
        let base_size = len / n;
        let remainder = len % n;

        let mut result = Vec::with_capacity(n);
        let mut start = 0;

        for i in 0..n {
            let size = base_size + if i < remainder { 1 } else { 0 };
            if size > 0 {
                result.push(Rc::new(Variable::Array(arr[start..start + size].to_vec())) as Rcvar);
            } else {
                result.push(Rc::new(Variable::Array(vec![])) as Rcvar);
            }
            start += size;
        }

        Ok(Rc::new(Variable::Array(result)))
    }
}

// =============================================================================
// difference(arr1, arr2) -> array (set difference)
// =============================================================================

define_function!(
    DifferenceFn,
    vec![ArgumentType::Array, ArgumentType::Array],
    None
);

impl Function for DifferenceFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let arr1 = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        let arr2 = args[1].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        let set2: HashSet<String> = arr2
            .iter()
            .map(|v| serde_json::to_string(&**v).unwrap_or_default())
            .collect();

        let result: Vec<Rcvar> = arr1
            .iter()
            .filter(|v| {
                let key = serde_json::to_string(&***v).unwrap_or_default();
                !set2.contains(&key)
            })
            .cloned()
            .collect();

        Ok(Rc::new(Variable::Array(result)))
    }
}

// =============================================================================
// intersection(arr1, arr2) -> array (set intersection)
// =============================================================================

define_function!(
    IntersectionFn,
    vec![ArgumentType::Array, ArgumentType::Array],
    None
);

impl Function for IntersectionFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let arr1 = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        let arr2 = args[1].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        let set2: HashSet<String> = arr2
            .iter()
            .map(|v| serde_json::to_string(&**v).unwrap_or_default())
            .collect();

        let mut seen: HashSet<String> = HashSet::new();
        let result: Vec<Rcvar> = arr1
            .iter()
            .filter(|v| {
                let key = serde_json::to_string(&***v).unwrap_or_default();
                set2.contains(&key) && seen.insert(key)
            })
            .cloned()
            .collect();

        Ok(Rc::new(Variable::Array(result)))
    }
}

// =============================================================================
// union(arr1, arr2) -> array (set union)
// =============================================================================

define_function!(
    UnionFn,
    vec![ArgumentType::Array, ArgumentType::Array],
    None
);

impl Function for UnionFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let arr1 = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        let arr2 = args[1].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        let mut seen: HashSet<String> = HashSet::new();
        let mut result: Vec<Rcvar> = Vec::new();

        for item in arr1.iter().chain(arr2.iter()) {
            let key = serde_json::to_string(&**item).unwrap_or_default();
            if seen.insert(key) {
                result.push(item.clone());
            }
        }

        Ok(Rc::new(Variable::Array(result)))
    }
}

// =============================================================================
// frequencies(array) -> object (count occurrences)
// =============================================================================

define_function!(FrequenciesFn, vec![ArgumentType::Array], None);

impl Function for FrequenciesFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let arr = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        let mut counts: std::collections::HashMap<String, i64> = std::collections::HashMap::new();

        for item in arr {
            let key = match &**item {
                Variable::String(s) => s.clone(),
                Variable::Number(n) => n.to_string(),
                Variable::Bool(b) => b.to_string(),
                Variable::Null => "null".to_string(),
                _ => serde_json::to_string(&**item).unwrap_or_else(|_| "null".to_string()),
            };
            *counts.entry(key).or_insert(0) += 1;
        }

        let result: std::collections::BTreeMap<String, Rcvar> = counts
            .into_iter()
            .map(|(k, v)| {
                (
                    k,
                    Rc::new(Variable::Number(serde_json::Number::from(v))) as Rcvar,
                )
            })
            .collect();

        Ok(Rc::new(Variable::Object(result)))
    }
}

// =============================================================================
// mode(array) -> any (most frequent value)
// =============================================================================

define_function!(ModeFn, vec![ArgumentType::Array], None);

impl Function for ModeFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let arr = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        if arr.is_empty() {
            return Ok(Rc::new(Variable::Null));
        }

        let mut counts: std::collections::HashMap<String, (i64, Rcvar)> =
            std::collections::HashMap::new();

        for item in arr {
            let key = serde_json::to_string(&**item).unwrap_or_default();
            counts
                .entry(key)
                .and_modify(|(count, _)| *count += 1)
                .or_insert((1, item.clone()));
        }

        let (_, (_, mode_value)) = counts
            .into_iter()
            .max_by_key(|(_, (count, _))| *count)
            .unwrap();

        Ok(mode_value)
    }
}

// =============================================================================
// cartesian(arr1, arr2) -> array (cartesian product of 2 arrays)
// cartesian(array_of_arrays) -> array (cartesian product of N arrays, jq parity)
// =============================================================================

define_function!(
    CartesianFn,
    vec![ArgumentType::Array],
    Some(ArgumentType::Array)
);

impl Function for CartesianFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let first = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        // Two modes:
        // 1. cartesian(arr1, arr2) - two separate arrays (original behavior)
        // 2. cartesian(array_of_arrays) - single array containing arrays (jq-style)

        if args.len() == 2 {
            // Original two-argument mode
            let arr2 = args[1].as_array().ok_or_else(|| {
                JmespathError::new(
                    ctx.expression,
                    0,
                    ErrorReason::Parse("Expected array argument".to_owned()),
                )
            })?;

            let mut result = Vec::with_capacity(first.len() * arr2.len());
            for a in first {
                for b in arr2 {
                    result.push(Rc::new(Variable::Array(vec![a.clone(), b.clone()])) as Rcvar);
                }
            }
            Ok(Rc::new(Variable::Array(result)))
        } else {
            // Single argument mode - check if it's an array of arrays
            // If all elements are arrays, do N-way cartesian product
            let arrays: Vec<&Vec<Rcvar>> =
                first.iter().filter_map(|item| item.as_array()).collect();

            if arrays.len() != first.len() || arrays.is_empty() {
                // Not all elements are arrays, or empty - return empty
                return Ok(Rc::new(Variable::Array(vec![])));
            }

            // N-way cartesian product
            let result = cartesian_product_n(&arrays);
            Ok(Rc::new(Variable::Array(result)))
        }
    }
}

/// Compute N-way cartesian product of arrays
fn cartesian_product_n(arrays: &[&Vec<Rcvar>]) -> Vec<Rcvar> {
    if arrays.is_empty() {
        return vec![];
    }

    if arrays.len() == 1 {
        // Single array - each element becomes a 1-element tuple
        return arrays[0]
            .iter()
            .map(|item| Rc::new(Variable::Array(vec![item.clone()])) as Rcvar)
            .collect();
    }

    // Calculate total size for pre-allocation
    let total_size: usize = arrays.iter().map(|a| a.len()).product();
    if total_size == 0 {
        return vec![];
    }

    let mut result = Vec::with_capacity(total_size);

    // Iterative cartesian product using indices
    let mut indices = vec![0usize; arrays.len()];

    loop {
        // Build current combination
        let combo: Vec<Rcvar> = indices
            .iter()
            .enumerate()
            .map(|(arr_idx, &elem_idx)| arrays[arr_idx][elem_idx].clone())
            .collect();
        result.push(Rc::new(Variable::Array(combo)) as Rcvar);

        // Increment indices (like counting in mixed radix)
        let mut carry = true;
        for i in (0..arrays.len()).rev() {
            if carry {
                indices[i] += 1;
                if indices[i] >= arrays[i].len() {
                    indices[i] = 0;
                } else {
                    carry = false;
                }
            }
        }

        // If we carried all the way through, we're done
        if carry {
            break;
        }
    }

    result
}

// =============================================================================
// initial(array) -> array (all elements except the last)
// =============================================================================

define_function!(InitialFn, vec![ArgumentType::Array], None);

impl Function for InitialFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let arr = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        if arr.is_empty() {
            return Ok(Rc::new(Variable::Array(vec![])));
        }

        let result: Vec<Rcvar> = arr[..arr.len() - 1].to_vec();
        Ok(Rc::new(Variable::Array(result)))
    }
}

// =============================================================================
// interpose(array, separator) -> array (insert separator between elements)
// =============================================================================

// Insert a separator value between each element of an array.
//
// # Arguments
// * `array` - The input array
// * `separator` - The value to insert between elements
//
// # Returns
// A new array with the separator inserted between each element.
//
// # Example
// interpose([1, 2, 3], 0) -> [1, 0, 2, 0, 3]
// interpose(["a", "b", "c"], "-") -> ["a", "-", "b", "-", "c"]
define_function!(
    InterposeFn,
    vec![ArgumentType::Array, ArgumentType::Any],
    None
);

impl Function for InterposeFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let arr = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        let separator = args[1].clone();

        if arr.is_empty() {
            return Ok(Rc::new(Variable::Array(vec![])));
        }

        if arr.len() == 1 {
            return Ok(Rc::new(Variable::Array(arr.clone())));
        }

        let mut result = Vec::with_capacity(arr.len() * 2 - 1);
        for (i, item) in arr.iter().enumerate() {
            if i > 0 {
                result.push(separator.clone());
            }
            result.push(item.clone());
        }

        Ok(Rc::new(Variable::Array(result)))
    }
}

// =============================================================================
// zipmap(keys, values) -> object (create object from parallel arrays)
// =============================================================================

// Create an object from two parallel arrays of keys and values.
//
// # Arguments
// * `keys` - An array of strings to use as object keys
// * `values` - An array of values
//
// # Returns
// An object mapping each key to its corresponding value.
// If arrays have different lengths, uses the shorter length.
//
// # Example
// zipmap(["a", "b", "c"], [1, 2, 3]) -> {"a": 1, "b": 2, "c": 3}
// zipmap(["x", "y"], [10, 20, 30]) -> {"x": 10, "y": 20}
define_function!(
    ZipmapFn,
    vec![ArgumentType::Array, ArgumentType::Array],
    None
);

impl Function for ZipmapFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let keys = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument for keys".to_owned()),
            )
        })?;

        let values = args[1].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument for values".to_owned()),
            )
        })?;

        let len = keys.len().min(values.len());
        let mut result = std::collections::BTreeMap::new();

        for i in 0..len {
            let key = keys[i].as_string().ok_or_else(|| {
                JmespathError::new(
                    ctx.expression,
                    0,
                    ErrorReason::Parse("Keys must be strings".to_owned()),
                )
            })?;
            result.insert(key.to_string(), values[i].clone());
        }

        Ok(Rc::new(Variable::Object(result)))
    }
}

// =============================================================================
// partition_by(array, field_name) -> array (split when field value changes)
// =============================================================================

// Split an array into partitions when the field value changes.
//
// Unlike `group_by` which groups all elements with the same key together,
// `partition_by` creates a new partition each time the field value
// changes, preserving the original order.
//
// # Arguments
// * `array` - The input array of objects
// * `field_name` - The field name to partition by
//
// # Returns
// An array of arrays, each containing consecutive elements with the same
// field value.
//
// # Example
// partition_by([{type: "a"}, {type: "a"}, {type: "b"}], "type")
//   -> [[{type: "a"}, {type: "a"}], [{type: "b"}]]
define_function!(
    PartitionByFn,
    vec![ArgumentType::Array, ArgumentType::String],
    None
);

impl Function for PartitionByFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let arr = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        let field_name = args[1].as_string().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected field name string".to_owned()),
            )
        })?;

        if arr.is_empty() {
            return Ok(Rc::new(Variable::Array(vec![])));
        }

        let mut result: Vec<Rcvar> = Vec::new();
        let mut current_partition: Vec<Rcvar> = Vec::new();
        let mut last_key: Option<String> = None;

        for item in arr {
            // Extract key from object field, or serialize the item itself for primitives
            let key = if let Some(obj) = item.as_object() {
                if let Some(field_value) = obj.get(field_name) {
                    serde_json::to_string(&**field_value).unwrap_or_default()
                } else {
                    "null".to_string()
                }
            } else {
                // For non-objects, use the value itself as the key
                serde_json::to_string(&**item).unwrap_or_default()
            };

            match &last_key {
                Some(prev_key) if *prev_key == key => {
                    current_partition.push(item.clone());
                }
                _ => {
                    if !current_partition.is_empty() {
                        result.push(Rc::new(Variable::Array(current_partition)));
                    }
                    current_partition = vec![item.clone()];
                    last_key = Some(key);
                }
            }
        }

        // Don't forget the last partition
        if !current_partition.is_empty() {
            result.push(Rc::new(Variable::Array(current_partition)));
        }

        Ok(Rc::new(Variable::Array(result)))
    }
}

// =============================================================================
// dedupe(array) -> array (remove consecutive duplicates)
// =============================================================================

// Remove consecutive duplicate values from an array.
//
// Unlike `unique` which removes all duplicates, `dedupe` only removes
// adjacent duplicates, preserving the relative order and allowing the
// same value to appear multiple times if separated by other values.
//
// # Arguments
// * `array` - The input array
//
// # Returns
// A new array with consecutive duplicates removed.
//
// # Example
// dedupe([1, 1, 2, 2, 1, 1]) -> [1, 2, 1]
// dedupe(["a", "a", "b", "a"]) -> ["a", "b", "a"]
define_function!(DedupeFn, vec![ArgumentType::Array], None);

impl Function for DedupeFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let arr = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        if arr.is_empty() {
            return Ok(Rc::new(Variable::Array(vec![])));
        }

        let mut result: Vec<Rcvar> = Vec::with_capacity(arr.len());
        let mut last_key: Option<String> = None;

        for item in arr {
            let key = serde_json::to_string(&**item).unwrap_or_default();
            match &last_key {
                Some(prev_key) if *prev_key == key => {
                    // Skip consecutive duplicate
                }
                _ => {
                    result.push(item.clone());
                    last_key = Some(key);
                }
            }
        }

        Ok(Rc::new(Variable::Array(result)))
    }
}

// =============================================================================
// tail(array) -> array (all elements except the first)
// =============================================================================

define_function!(TailFn, vec![ArgumentType::Array], None);

impl Function for TailFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let arr = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        if arr.is_empty() {
            return Ok(Rc::new(Variable::Array(vec![])));
        }

        let result: Vec<Rcvar> = arr[1..].to_vec();
        Ok(Rc::new(Variable::Array(result)))
    }
}

// =============================================================================
// without(array, values_array) -> array (remove specified values)
// =============================================================================

define_function!(
    WithoutFn,
    vec![ArgumentType::Array, ArgumentType::Array],
    None
);

impl Function for WithoutFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let arr = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        let exclude = args[1].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument for values to exclude".to_owned()),
            )
        })?;

        // Create a set of serialized values to exclude for efficient lookup
        let exclude_set: HashSet<String> = exclude
            .iter()
            .map(|v| serde_json::to_string(&**v).unwrap_or_default())
            .collect();

        let result: Vec<Rcvar> = arr
            .iter()
            .filter(|item| {
                let key = serde_json::to_string(&***item).unwrap_or_default();
                !exclude_set.contains(&key)
            })
            .cloned()
            .collect();

        Ok(Rc::new(Variable::Array(result)))
    }
}

// =============================================================================
// xor(array1, array2) -> array (symmetric difference)
// =============================================================================

define_function!(XorFn, vec![ArgumentType::Array, ArgumentType::Array], None);

impl Function for XorFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let arr1 = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        let arr2 = args[1].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        // Create sets of serialized values
        let set1: HashSet<String> = arr1
            .iter()
            .map(|v| serde_json::to_string(&**v).unwrap_or_default())
            .collect();

        let set2: HashSet<String> = arr2
            .iter()
            .map(|v| serde_json::to_string(&**v).unwrap_or_default())
            .collect();

        let mut result = Vec::new();

        // Add elements from arr1 that are not in arr2
        for item in arr1 {
            let key = serde_json::to_string(&**item).unwrap_or_default();
            if !set2.contains(&key) {
                result.push(item.clone());
            }
        }

        // Add elements from arr2 that are not in arr1
        for item in arr2 {
            let key = serde_json::to_string(&**item).unwrap_or_default();
            if !set1.contains(&key) {
                result.push(item.clone());
            }
        }

        Ok(Rc::new(Variable::Array(result)))
    }
}

// =============================================================================
// window(array, size, step?) -> array (sliding window)
// =============================================================================

define_function!(
    WindowFn,
    vec![ArgumentType::Array, ArgumentType::Number],
    Some(ArgumentType::Number)
);

impl Function for WindowFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let arr = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        let size = args[1].as_number().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected number for window size".to_owned()),
            )
        })? as usize;

        if size == 0 {
            return Ok(Rc::new(Variable::Array(vec![])));
        }

        // Default step is 1
        let step = if args.len() > 2 {
            args[2].as_number().ok_or_else(|| {
                JmespathError::new(
                    ctx.expression,
                    0,
                    ErrorReason::Parse("Expected number for step".to_owned()),
                )
            })? as usize
        } else {
            1
        };

        if step == 0 {
            return Err(JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Step cannot be zero".to_owned()),
            ));
        }

        let len = arr.len();
        if len < size {
            return Ok(Rc::new(Variable::Array(vec![])));
        }

        let mut result = Vec::new();
        let mut i = 0;

        while i + size <= len {
            let window: Vec<Rcvar> = arr[i..i + size].to_vec();
            result.push(Rc::new(Variable::Array(window)) as Rcvar);
            i += step;
        }

        Ok(Rc::new(Variable::Array(result)))
    }
}

// =============================================================================
// combinations(array, k) -> array (k-combinations of array)
// =============================================================================

define_function!(
    CombinationsFn,
    vec![ArgumentType::Array, ArgumentType::Number],
    None
);

fn generate_combinations(arr: &[Rcvar], k: usize) -> Vec<Vec<Rcvar>> {
    if k == 0 {
        return vec![vec![]];
    }
    if arr.len() < k {
        return vec![];
    }

    let mut result = Vec::new();

    // Include first element in combination
    let first = arr[0].clone();
    let rest = &arr[1..];
    for mut combo in generate_combinations(rest, k - 1) {
        let mut new_combo = vec![first.clone()];
        new_combo.append(&mut combo);
        result.push(new_combo);
    }

    // Exclude first element
    result.extend(generate_combinations(rest, k));

    result
}

impl Function for CombinationsFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let arr = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        let k = args[1].as_number().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected number for k".to_owned()),
            )
        })? as usize;

        // Limit to prevent excessive computation
        const MAX_COMBINATIONS: usize = 10000;

        // Quick check: if C(n, k) would be too large, return error
        let n = arr.len();
        if n > 20 && k > 3 && k < n - 3 {
            // For large arrays with mid-range k, combinations could be huge
            // Rough estimate: C(n, k) grows quickly
            return Err(JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Combination size too large".to_owned()),
            ));
        }

        let combinations = generate_combinations(arr, k);

        if combinations.len() > MAX_COMBINATIONS {
            return Err(JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Too many combinations generated".to_owned()),
            ));
        }

        let result: Vec<Rcvar> = combinations
            .into_iter()
            .map(|combo| Rc::new(Variable::Array(combo)) as Rcvar)
            .collect();

        Ok(Rc::new(Variable::Array(result)))
    }
}

// =============================================================================
// fill(array, value, start?, end?) -> array (fill range with value)
// =============================================================================

define_function!(
    FillFn,
    vec![ArgumentType::Array, ArgumentType::Any],
    Some(ArgumentType::Number)
);

impl Function for FillFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let arr = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        let fill_value = args[1].clone();

        let len = arr.len();
        if len == 0 {
            return Ok(Rc::new(Variable::Array(vec![])));
        }

        // Default start is 0, default end is array length
        let start = if args.len() > 2 {
            let s = args[2].as_number().ok_or_else(|| {
                JmespathError::new(
                    ctx.expression,
                    0,
                    ErrorReason::Parse("Expected number for start index".to_owned()),
                )
            })? as i64;
            // Handle negative indices
            if s < 0 {
                (len as i64 + s).max(0) as usize
            } else {
                (s as usize).min(len)
            }
        } else {
            0
        };

        let end = if args.len() > 3 {
            let e = args[3].as_number().ok_or_else(|| {
                JmespathError::new(
                    ctx.expression,
                    0,
                    ErrorReason::Parse("Expected number for end index".to_owned()),
                )
            })? as i64;
            // Handle negative indices
            if e < 0 {
                (len as i64 + e).max(0) as usize
            } else {
                (e as usize).min(len)
            }
        } else {
            len
        };

        let mut result: Vec<Rcvar> = arr.clone();

        for item in result.iter_mut().take(end.min(len)).skip(start) {
            *item = fill_value.clone();
        }

        Ok(Rc::new(Variable::Array(result)))
    }
}

// =============================================================================
// pull_at(array, indices_array) -> array (get elements at specified indices)
// =============================================================================

define_function!(
    PullAtFn,
    vec![ArgumentType::Array, ArgumentType::Array],
    None
);

impl Function for PullAtFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let arr = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        let indices = args[1].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array of indices".to_owned()),
            )
        })?;

        let len = arr.len();
        let mut result = Vec::new();

        for idx_var in indices {
            let idx = idx_var.as_number().ok_or_else(|| {
                JmespathError::new(
                    ctx.expression,
                    0,
                    ErrorReason::Parse("Expected number in indices array".to_owned()),
                )
            })? as i64;

            // Handle negative indices
            let actual_idx = if idx < 0 {
                (len as i64 + idx).max(0) as usize
            } else {
                idx as usize
            };

            if actual_idx < len {
                result.push(arr[actual_idx].clone());
            }
        }

        Ok(Rc::new(Variable::Array(result)))
    }
}

// =============================================================================
// transpose(array) -> array
// =============================================================================

// Transpose a 2D array (swap rows and columns).
//
// # Arguments
// * `array` - A 2D array (array of arrays)
//
// # Returns
// A new 2D array with rows and columns swapped.
// The result has as many rows as the shortest inner array.
//
// # Example
// transpose([[1, 2, 3], [4, 5, 6]]) -> [[1, 4], [2, 5], [3, 6]]
// transpose([[1, 2], [3, 4], [5, 6]]) -> [[1, 3, 5], [2, 4, 6]]
define_function!(TransposeFn, vec![ArgumentType::Array], None);

impl Function for TransposeFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let arr = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        if arr.is_empty() {
            return Ok(Rc::new(Variable::Array(vec![])));
        }

        // Get all inner arrays and find the minimum length
        let mut inner_arrays: Vec<&Vec<Rcvar>> = Vec::new();
        let mut min_len = usize::MAX;

        for item in arr {
            if let Some(inner) = item.as_array() {
                min_len = min_len.min(inner.len());
                inner_arrays.push(inner);
            } else {
                // If any element is not an array, return empty
                return Ok(Rc::new(Variable::Array(vec![])));
            }
        }

        if inner_arrays.is_empty() || min_len == 0 {
            return Ok(Rc::new(Variable::Array(vec![])));
        }

        // Transpose: create new arrays where each contains the i-th element from each inner array
        let mut result = Vec::with_capacity(min_len);
        for i in 0..min_len {
            let mut row = Vec::with_capacity(inner_arrays.len());
            for inner in &inner_arrays {
                row.push(inner[i].clone());
            }
            result.push(Rc::new(Variable::Array(row)));
        }

        Ok(Rc::new(Variable::Array(result)))
    }
}

// =============================================================================
// pairwise(array) -> array
// =============================================================================

// Return adjacent pairs from an array.
//
// This is equivalent to `window(array, 2)` but provided as a convenience.
//
// # Arguments
// * `array` - The input array
//
// # Returns
// An array of 2-element arrays, each containing adjacent elements.
//
// # Example
// pairwise([1, 2, 3, 4]) -> [[1, 2], [2, 3], [3, 4]]
// pairwise([10, 15, 13, 20]) -> [[10, 15], [15, 13], [13, 20]]
define_function!(PairwiseFn, vec![ArgumentType::Array], None);

impl Function for PairwiseFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let arr = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        if arr.len() < 2 {
            return Ok(Rc::new(Variable::Array(vec![])));
        }

        let mut result = Vec::with_capacity(arr.len() - 1);
        for i in 0..arr.len() - 1 {
            let pair = vec![arr[i].clone(), arr[i + 1].clone()];
            result.push(Rc::new(Variable::Array(pair)));
        }

        Ok(Rc::new(Variable::Array(result)))
    }
}

// =============================================================================
// indices_array(array, value) -> array of indices
// =============================================================================

// Find all indices where a value appears in an array.
//
// # Arguments
// * `array` - The input array to search
// * `value` - The value to find
//
// # Returns
// An array of indices (numbers) where the value appears.
//
// # Example
// indices_array([1, 2, 3, 2, 4, 2], 2) -> [1, 3, 5]
// indices_array(["a", "b", "a", "c"], "a") -> [0, 2]
define_function!(
    IndicesArrayFn,
    vec![ArgumentType::Array, ArgumentType::Any],
    None
);

impl Function for IndicesArrayFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let arr = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        let search_key = serde_json::to_string(&*args[1]).unwrap_or_default();

        let mut indices: Vec<Rcvar> = Vec::new();
        for (i, item) in arr.iter().enumerate() {
            let item_key = serde_json::to_string(&**item).unwrap_or_default();
            if item_key == search_key {
                indices.push(Rc::new(Variable::Number(serde_json::Number::from(
                    i as i64,
                ))));
            }
        }

        Ok(Rc::new(Variable::Array(indices)))
    }
}

// =============================================================================
// inside_array(needle, haystack) -> boolean
// =============================================================================

// Check if all elements of the first array are contained in the second array.
// This is the inverse of contains for arrays.
//
// # Arguments
// * `needle` - The array whose elements to look for
// * `haystack` - The array to search in
//
// # Returns
// `true` if all elements of needle appear in haystack, `false` otherwise.
//
// # Example
// inside_array([1, 2], [1, 2, 3, 4]) -> true
// inside_array([1, 5], [1, 2, 3, 4]) -> false
define_function!(
    InsideArrayFn,
    vec![ArgumentType::Array, ArgumentType::Array],
    None
);

impl Function for InsideArrayFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let needle = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        let haystack = args[1].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        // Build a set of serialized haystack values for efficient lookup
        let haystack_set: HashSet<String> = haystack
            .iter()
            .map(|item| serde_json::to_string(&**item).unwrap_or_default())
            .collect();

        // Check if all needle elements are in the haystack
        let result = needle.iter().all(|item| {
            let item_key = serde_json::to_string(&**item).unwrap_or_default();
            haystack_set.contains(&item_key)
        });

        Ok(Rc::new(Variable::Bool(result)))
    }
}

// =============================================================================
// bsearch(sorted_array, value) -> number
// =============================================================================

// Binary search for a value in a sorted array.
//
// Returns the index if found, or the negative insertion point minus 1 if not found.
// The insertion point is the index where the value would be inserted to maintain
// sorted order. This follows jq's bsearch semantics.
//
// # Arguments
// * `sorted_array` - A sorted array to search in
// * `value` - The value to find
//
// # Returns
// - If found: the index of the element (0-based)
// - If not found: -(insertion_point) - 1
//
// # Example
// bsearch([1, 3, 5, 7, 9], 5) -> 2
// bsearch([1, 3, 5, 7, 9], 4) -> -3 (would be inserted at index 2)
// bsearch([1, 3, 5, 7, 9], 0) -> -1 (would be inserted at index 0)
define_function!(
    BsearchFn,
    vec![ArgumentType::Array, ArgumentType::Any],
    None
);

impl Function for BsearchFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let arr = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        let target = &args[1];

        if arr.is_empty() {
            return Ok(Rc::new(Variable::Number(serde_json::Number::from(-1))));
        }

        // Helper to compare two Variable values
        // Returns Ordering based on type-aware comparison
        fn compare_values(a: &Variable, b: &Variable) -> std::cmp::Ordering {
            match (a, b) {
                // Numbers: compare numerically
                (Variable::Number(n1), Variable::Number(n2)) => {
                    let f1 = n1.as_f64().unwrap_or(0.0);
                    let f2 = n2.as_f64().unwrap_or(0.0);
                    f1.partial_cmp(&f2).unwrap_or(std::cmp::Ordering::Equal)
                }
                // Strings: compare lexicographically
                (Variable::String(s1), Variable::String(s2)) => s1.cmp(s2),
                // Booleans: false < true
                (Variable::Bool(b1), Variable::Bool(b2)) => b1.cmp(b2),
                // Different types or complex types: fall back to JSON string comparison
                _ => {
                    let s1 = serde_json::to_string(a).unwrap_or_default();
                    let s2 = serde_json::to_string(b).unwrap_or_default();
                    s1.cmp(&s2)
                }
            }
        }

        let mut left = 0i64;
        let mut right = arr.len() as i64 - 1;

        while left <= right {
            let mid = left + (right - left) / 2;

            match compare_values(&arr[mid as usize], target) {
                std::cmp::Ordering::Equal => {
                    return Ok(Rc::new(Variable::Number(serde_json::Number::from(mid))));
                }
                std::cmp::Ordering::Less => {
                    left = mid + 1;
                }
                std::cmp::Ordering::Greater => {
                    right = mid - 1;
                }
            }
        }

        // Not found - return -(insertion_point) - 1
        // At this point, left is the insertion point
        let result = -(left) - 1;
        Ok(Rc::new(Variable::Number(serde_json::Number::from(result))))
    }
}

// =============================================================================
// repeat_array(value, n) -> array (create array with value repeated n times)
// =============================================================================

// Creates an array containing the given value repeated n times.
//
// # Arguments
// * `value` - The value to repeat
// * `n` - Number of times to repeat (must be non-negative)
//
// # Returns
// An array containing the value repeated n times.
//
// # Example
// repeat_array(`1`, `3`) -> [1, 1, 1]
// repeat_array("x", `4`) -> ["x", "x", "x", "x"]
// repeat_array(`0`, `0`) -> []
define_function!(
    RepeatArrayFn,
    vec![ArgumentType::Any, ArgumentType::Number],
    None
);

impl Function for RepeatArrayFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let value = &args[0];
        let n = args[1].as_number().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected number for count argument".to_owned()),
            )
        })? as i64;

        if n < 0 {
            return Err(JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Count must be non-negative".to_owned()),
            ));
        }

        let result: Vec<Rcvar> = (0..n).map(|_| value.clone()).collect();
        Ok(Rc::new(Variable::Array(result)))
    }
}

// =============================================================================
// cycle(array, n) -> array (cycle through array n times)
// =============================================================================

// Creates an array by cycling through the input array n times.
//
// # Arguments
// * `array` - The array to cycle through
// * `n` - Number of times to cycle (must be non-negative)
//
// # Returns
// An array containing the elements of the input array repeated n times.
//
// # Example
// cycle([1, 2, 3], `2`) -> [1, 2, 3, 1, 2, 3]
// cycle(["a", "b"], `3`) -> ["a", "b", "a", "b", "a", "b"]
// cycle([1, 2], `0`) -> []
define_function!(
    CycleFn,
    vec![ArgumentType::Array, ArgumentType::Number],
    None
);

impl Function for CycleFn {
    fn evaluate(&self, args: &[Rcvar], ctx: &mut Context<'_>) -> Result<Rcvar, JmespathError> {
        self.signature.validate(args, ctx)?;

        let arr = args[0].as_array().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected array argument".to_owned()),
            )
        })?;

        let n = args[1].as_number().ok_or_else(|| {
            JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Expected number for count argument".to_owned()),
            )
        })? as i64;

        if n < 0 {
            return Err(JmespathError::new(
                ctx.expression,
                0,
                ErrorReason::Parse("Count must be non-negative".to_owned()),
            ));
        }

        if arr.is_empty() || n == 0 {
            return Ok(Rc::new(Variable::Array(vec![])));
        }

        let mut result: Vec<Rcvar> = Vec::with_capacity(arr.len() * n as usize);
        for _ in 0..n {
            result.extend(arr.iter().cloned());
        }
        Ok(Rc::new(Variable::Array(result)))
    }
}

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

    fn setup_runtime() -> Runtime {
        let mut runtime = Runtime::new();
        runtime.register_builtin_functions();
        register(&mut runtime);
        runtime
    }

    #[test]
    fn test_unique() {
        let runtime = setup_runtime();
        let expr = runtime.compile("unique(@)").unwrap();
        let data = Variable::Array(vec![
            Rc::new(Variable::Number(serde_json::Number::from(1))),
            Rc::new(Variable::Number(serde_json::Number::from(2))),
            Rc::new(Variable::Number(serde_json::Number::from(1))),
        ]);
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 2);
    }

    #[test]
    fn test_first() {
        let runtime = setup_runtime();
        let expr = runtime.compile("first(@)").unwrap();
        let data = Variable::Array(vec![
            Rc::new(Variable::Number(serde_json::Number::from(1))),
            Rc::new(Variable::Number(serde_json::Number::from(2))),
        ]);
        let result = expr.search(&data).unwrap();
        assert_eq!(result.as_number().unwrap() as i64, 1);
    }

    #[test]
    fn test_last() {
        let runtime = setup_runtime();
        let expr = runtime.compile("last(@)").unwrap();
        let data = Variable::Array(vec![
            Rc::new(Variable::Number(serde_json::Number::from(1))),
            Rc::new(Variable::Number(serde_json::Number::from(2))),
        ]);
        let result = expr.search(&data).unwrap();
        assert_eq!(result.as_number().unwrap() as i64, 2);
    }

    #[test]
    fn test_range() {
        let runtime = setup_runtime();
        let expr = runtime.compile("range(`0`, `5`)").unwrap();
        let data = Variable::Null;
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 5);
    }

    #[test]
    fn test_initial() {
        let runtime = setup_runtime();
        let expr = runtime.compile("initial(@)").unwrap();
        let data = Variable::Array(vec![
            Rc::new(Variable::Number(serde_json::Number::from(1))),
            Rc::new(Variable::Number(serde_json::Number::from(2))),
            Rc::new(Variable::Number(serde_json::Number::from(3))),
        ]);
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 2);
        assert_eq!(arr[0].as_number().unwrap() as i64, 1);
        assert_eq!(arr[1].as_number().unwrap() as i64, 2);
    }

    #[test]
    fn test_initial_empty() {
        let runtime = setup_runtime();
        let expr = runtime.compile("initial(@)").unwrap();
        let data = Variable::Array(vec![]);
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 0);
    }

    #[test]
    fn test_butlast() {
        // butlast is an alias for initial
        let runtime = setup_runtime();
        let expr = runtime.compile("butlast(@)").unwrap();
        let data = Variable::Array(vec![
            Rc::new(Variable::Number(serde_json::Number::from(1))),
            Rc::new(Variable::Number(serde_json::Number::from(2))),
            Rc::new(Variable::Number(serde_json::Number::from(3))),
        ]);
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 2);
        assert_eq!(arr[0].as_number().unwrap() as i64, 1);
        assert_eq!(arr[1].as_number().unwrap() as i64, 2);
    }

    #[test]
    fn test_interpose() {
        let runtime = setup_runtime();
        let expr = runtime.compile("interpose(@, `0`)").unwrap();
        let data = Variable::Array(vec![
            Rc::new(Variable::Number(serde_json::Number::from(1))),
            Rc::new(Variable::Number(serde_json::Number::from(2))),
            Rc::new(Variable::Number(serde_json::Number::from(3))),
        ]);
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 5);
        assert_eq!(arr[0].as_number().unwrap() as i64, 1);
        assert_eq!(arr[1].as_number().unwrap() as i64, 0);
        assert_eq!(arr[2].as_number().unwrap() as i64, 2);
        assert_eq!(arr[3].as_number().unwrap() as i64, 0);
        assert_eq!(arr[4].as_number().unwrap() as i64, 3);
    }

    #[test]
    fn test_interpose_empty() {
        let runtime = setup_runtime();
        let expr = runtime.compile("interpose(@, `0`)").unwrap();
        let data = Variable::Array(vec![]);
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 0);
    }

    #[test]
    fn test_interpose_single() {
        let runtime = setup_runtime();
        let expr = runtime.compile("interpose(@, `0`)").unwrap();
        let data = Variable::Array(vec![Rc::new(Variable::Number(serde_json::Number::from(1)))]);
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 1);
        assert_eq!(arr[0].as_number().unwrap() as i64, 1);
    }

    #[test]
    fn test_interpose_string_separator() {
        let runtime = setup_runtime();
        let expr = runtime.compile("interpose(@, `\"-\"`)").unwrap();
        let data = Variable::Array(vec![
            Rc::new(Variable::String("a".to_string())),
            Rc::new(Variable::String("b".to_string())),
            Rc::new(Variable::String("c".to_string())),
        ]);
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 5);
        assert_eq!(arr[0].as_string().unwrap(), "a");
        assert_eq!(arr[1].as_string().unwrap(), "-");
        assert_eq!(arr[2].as_string().unwrap(), "b");
        assert_eq!(arr[3].as_string().unwrap(), "-");
        assert_eq!(arr[4].as_string().unwrap(), "c");
    }

    #[test]
    fn test_zipmap() {
        let runtime = setup_runtime();
        let expr = runtime
            .compile("zipmap(`[\"a\", \"b\", \"c\"]`, `[1, 2, 3]`)")
            .unwrap();
        let data = Variable::Null;
        let result = expr.search(&data).unwrap();
        let obj = result.as_object().unwrap();
        assert_eq!(obj.len(), 3);
        assert_eq!(obj.get("a").unwrap().as_number().unwrap() as i64, 1);
        assert_eq!(obj.get("b").unwrap().as_number().unwrap() as i64, 2);
        assert_eq!(obj.get("c").unwrap().as_number().unwrap() as i64, 3);
    }

    #[test]
    fn test_zipmap_unequal_lengths() {
        let runtime = setup_runtime();
        // Keys shorter than values
        let expr = runtime
            .compile("zipmap(`[\"x\", \"y\"]`, `[10, 20, 30]`)")
            .unwrap();
        let data = Variable::Null;
        let result = expr.search(&data).unwrap();
        let obj = result.as_object().unwrap();
        assert_eq!(obj.len(), 2);
        assert_eq!(obj.get("x").unwrap().as_number().unwrap() as i64, 10);
        assert_eq!(obj.get("y").unwrap().as_number().unwrap() as i64, 20);
    }

    #[test]
    fn test_zipmap_empty() {
        let runtime = setup_runtime();
        let expr = runtime.compile("zipmap(`[]`, `[]`)").unwrap();
        let data = Variable::Null;
        let result = expr.search(&data).unwrap();
        let obj = result.as_object().unwrap();
        assert_eq!(obj.len(), 0);
    }

    #[test]
    fn test_partition_by() {
        let runtime = setup_runtime();
        // Test with objects - partition by "type" field
        let expr = runtime.compile(r#"partition_by(@, `"type"`)"#).unwrap();
        let data: Variable = serde_json::from_str(
            r#"[{"type": "a", "v": 1}, {"type": "a", "v": 2}, {"type": "b", "v": 3}, {"type": "a", "v": 4}]"#,
        )
        .unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 3);

        let partition1 = arr[0].as_array().unwrap();
        assert_eq!(partition1.len(), 2); // Two "a" items

        let partition2 = arr[1].as_array().unwrap();
        assert_eq!(partition2.len(), 1); // One "b" item

        let partition3 = arr[2].as_array().unwrap();
        assert_eq!(partition3.len(), 1); // One more "a" item
    }

    #[test]
    fn test_partition_by_primitives() {
        let runtime = setup_runtime();
        // For primitives, use any field name - it will use the value itself
        let expr = runtime.compile(r#"partition_by(@, `"_"`)"#).unwrap();
        let data = Variable::Array(vec![
            Rc::new(Variable::Number(serde_json::Number::from(1))),
            Rc::new(Variable::Number(serde_json::Number::from(1))),
            Rc::new(Variable::Number(serde_json::Number::from(2))),
            Rc::new(Variable::Number(serde_json::Number::from(2))),
            Rc::new(Variable::Number(serde_json::Number::from(1))),
            Rc::new(Variable::Number(serde_json::Number::from(1))),
        ]);
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 3);

        let partition1 = arr[0].as_array().unwrap();
        assert_eq!(partition1.len(), 2);
        assert_eq!(partition1[0].as_number().unwrap() as i64, 1);

        let partition2 = arr[1].as_array().unwrap();
        assert_eq!(partition2.len(), 2);
        assert_eq!(partition2[0].as_number().unwrap() as i64, 2);

        let partition3 = arr[2].as_array().unwrap();
        assert_eq!(partition3.len(), 2);
        assert_eq!(partition3[0].as_number().unwrap() as i64, 1);
    }

    #[test]
    fn test_partition_by_empty() {
        let runtime = setup_runtime();
        let expr = runtime.compile(r#"partition_by(@, `"type"`)"#).unwrap();
        let data = Variable::Array(vec![]);
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 0);
    }

    #[test]
    fn test_partition_by_single() {
        let runtime = setup_runtime();
        let expr = runtime.compile(r#"partition_by(@, `"type"`)"#).unwrap();
        let data: Variable = serde_json::from_str(r#"[{"type": "a"}]"#).unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 1);
        let partition1 = arr[0].as_array().unwrap();
        assert_eq!(partition1.len(), 1);
    }

    #[test]
    fn test_dedupe() {
        let runtime = setup_runtime();
        let expr = runtime.compile("dedupe(@)").unwrap();
        let data = Variable::Array(vec![
            Rc::new(Variable::Number(serde_json::Number::from(1))),
            Rc::new(Variable::Number(serde_json::Number::from(1))),
            Rc::new(Variable::Number(serde_json::Number::from(2))),
            Rc::new(Variable::Number(serde_json::Number::from(2))),
            Rc::new(Variable::Number(serde_json::Number::from(1))),
            Rc::new(Variable::Number(serde_json::Number::from(1))),
        ]);
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 3);
        assert_eq!(arr[0].as_number().unwrap() as i64, 1);
        assert_eq!(arr[1].as_number().unwrap() as i64, 2);
        assert_eq!(arr[2].as_number().unwrap() as i64, 1);
    }

    #[test]
    fn test_dedupe_empty() {
        let runtime = setup_runtime();
        let expr = runtime.compile("dedupe(@)").unwrap();
        let data = Variable::Array(vec![]);
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 0);
    }

    #[test]
    fn test_dedupe_no_consecutive() {
        let runtime = setup_runtime();
        let expr = runtime.compile("dedupe(@)").unwrap();
        let data = Variable::Array(vec![
            Rc::new(Variable::Number(serde_json::Number::from(1))),
            Rc::new(Variable::Number(serde_json::Number::from(2))),
            Rc::new(Variable::Number(serde_json::Number::from(3))),
        ]);
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 3);
    }

    #[test]
    fn test_dedupe_strings() {
        let runtime = setup_runtime();
        let expr = runtime.compile("dedupe(@)").unwrap();
        let data = Variable::Array(vec![
            Rc::new(Variable::String("a".to_string())),
            Rc::new(Variable::String("a".to_string())),
            Rc::new(Variable::String("b".to_string())),
            Rc::new(Variable::String("a".to_string())),
        ]);
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 3);
        assert_eq!(arr[0].as_string().unwrap(), "a");
        assert_eq!(arr[1].as_string().unwrap(), "b");
        assert_eq!(arr[2].as_string().unwrap(), "a");
    }

    #[test]
    fn test_dedupe_objects() {
        let runtime = setup_runtime();
        let expr = runtime.compile("dedupe(@)").unwrap();
        let data: Variable =
            serde_json::from_str(r#"[{"x": 1}, {"x": 1}, {"x": 2}, {"x": 1}]"#).unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 3);
    }

    #[test]
    fn test_dedupe_all_same() {
        let runtime = setup_runtime();
        let expr = runtime.compile("dedupe(@)").unwrap();
        let data = Variable::Array(vec![
            Rc::new(Variable::Number(serde_json::Number::from(1))),
            Rc::new(Variable::Number(serde_json::Number::from(1))),
            Rc::new(Variable::Number(serde_json::Number::from(1))),
        ]);
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 1);
        assert_eq!(arr[0].as_number().unwrap() as i64, 1);
    }

    #[test]
    fn test_zipmap_duplicate_keys() {
        // Later key should overwrite earlier one
        let runtime = setup_runtime();
        let expr = runtime
            .compile("zipmap(`[\"a\", \"b\", \"a\"]`, `[1, 2, 3]`)")
            .unwrap();
        let data = Variable::Null;
        let result = expr.search(&data).unwrap();
        let obj = result.as_object().unwrap();
        assert_eq!(obj.len(), 2);
        assert_eq!(obj.get("a").unwrap().as_number().unwrap() as i64, 3); // Last value wins
        assert_eq!(obj.get("b").unwrap().as_number().unwrap() as i64, 2);
    }

    #[test]
    fn test_zipmap_values_longer() {
        let runtime = setup_runtime();
        let expr = runtime.compile("zipmap(`[\"a\"]`, `[1, 2, 3]`)").unwrap();
        let data = Variable::Null;
        let result = expr.search(&data).unwrap();
        let obj = result.as_object().unwrap();
        assert_eq!(obj.len(), 1);
        assert_eq!(obj.get("a").unwrap().as_number().unwrap() as i64, 1);
    }

    #[test]
    fn test_partition_by_missing_field() {
        let runtime = setup_runtime();
        let expr = runtime.compile(r#"partition_by(@, `"type"`)"#).unwrap();
        // Some objects have type, some don't
        let data: Variable =
            serde_json::from_str(r#"[{"type": "a"}, {"name": "no-type"}, {"type": "a"}]"#).unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        // Should partition into 3: [type:a], [no type -> null], [type:a]
        assert_eq!(arr.len(), 3);
    }

    #[test]
    fn test_partition_by_all_same() {
        let runtime = setup_runtime();
        let expr = runtime.compile(r#"partition_by(@, `"type"`)"#).unwrap();
        let data: Variable =
            serde_json::from_str(r#"[{"type": "a"}, {"type": "a"}, {"type": "a"}]"#).unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 1);
        let partition = arr[0].as_array().unwrap();
        assert_eq!(partition.len(), 3);
    }

    #[test]
    fn test_interpose_null_separator() {
        let runtime = setup_runtime();
        let expr = runtime.compile("interpose(@, `null`)").unwrap();
        let data = Variable::Array(vec![
            Rc::new(Variable::Number(serde_json::Number::from(1))),
            Rc::new(Variable::Number(serde_json::Number::from(2))),
        ]);
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 3);
        assert_eq!(arr[0].as_number().unwrap() as i64, 1);
        assert!(arr[1].is_null());
        assert_eq!(arr[2].as_number().unwrap() as i64, 2);
    }

    #[test]
    fn test_interpose_object_separator() {
        let runtime = setup_runtime();
        let expr = runtime.compile(r#"interpose(@, `{"sep": true}`)"#).unwrap();
        let data = Variable::Array(vec![
            Rc::new(Variable::Number(serde_json::Number::from(1))),
            Rc::new(Variable::Number(serde_json::Number::from(2))),
        ]);
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 3);
        assert!(arr[1].as_object().is_some());
    }

    #[test]
    fn test_tail() {
        let runtime = setup_runtime();
        let expr = runtime.compile("tail(@)").unwrap();
        let data = Variable::Array(vec![
            Rc::new(Variable::Number(serde_json::Number::from(1))),
            Rc::new(Variable::Number(serde_json::Number::from(2))),
            Rc::new(Variable::Number(serde_json::Number::from(3))),
        ]);
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 2);
        assert_eq!(arr[0].as_number().unwrap() as i64, 2);
        assert_eq!(arr[1].as_number().unwrap() as i64, 3);
    }

    #[test]
    fn test_tail_empty() {
        let runtime = setup_runtime();
        let expr = runtime.compile("tail(@)").unwrap();
        let data = Variable::Array(vec![]);
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 0);
    }

    #[test]
    fn test_without() {
        let runtime = setup_runtime();
        let expr = runtime.compile("without(@, `[2, 4]`)").unwrap();
        let data = Variable::Array(vec![
            Rc::new(Variable::Number(serde_json::Number::from(1))),
            Rc::new(Variable::Number(serde_json::Number::from(2))),
            Rc::new(Variable::Number(serde_json::Number::from(3))),
            Rc::new(Variable::Number(serde_json::Number::from(4))),
            Rc::new(Variable::Number(serde_json::Number::from(5))),
        ]);
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 3);
        assert_eq!(arr[0].as_number().unwrap() as i64, 1);
        assert_eq!(arr[1].as_number().unwrap() as i64, 3);
        assert_eq!(arr[2].as_number().unwrap() as i64, 5);
    }

    #[test]
    fn test_xor() {
        let runtime = setup_runtime();
        let expr = runtime.compile("xor(`[1, 2, 3]`, `[2, 3, 4]`)").unwrap();
        let data = Variable::Null;
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 2);
        assert_eq!(arr[0].as_number().unwrap() as i64, 1);
        assert_eq!(arr[1].as_number().unwrap() as i64, 4);
    }

    #[test]
    fn test_fill() {
        let runtime = setup_runtime();
        let expr = runtime.compile("fill(@, `0`)").unwrap();
        let data = Variable::Array(vec![
            Rc::new(Variable::Number(serde_json::Number::from(1))),
            Rc::new(Variable::Number(serde_json::Number::from(2))),
            Rc::new(Variable::Number(serde_json::Number::from(3))),
        ]);
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 3);
        assert_eq!(arr[0].as_number().unwrap() as i64, 0);
        assert_eq!(arr[1].as_number().unwrap() as i64, 0);
        assert_eq!(arr[2].as_number().unwrap() as i64, 0);
    }

    #[test]
    fn test_fill_with_range() {
        let runtime = setup_runtime();
        let expr = runtime.compile("fill(@, `0`, `1`, `3`)").unwrap();
        let data = Variable::Array(vec![
            Rc::new(Variable::Number(serde_json::Number::from(1))),
            Rc::new(Variable::Number(serde_json::Number::from(2))),
            Rc::new(Variable::Number(serde_json::Number::from(3))),
            Rc::new(Variable::Number(serde_json::Number::from(4))),
        ]);
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 4);
        assert_eq!(arr[0].as_number().unwrap() as i64, 1);
        assert_eq!(arr[1].as_number().unwrap() as i64, 0);
        assert_eq!(arr[2].as_number().unwrap() as i64, 0);
        assert_eq!(arr[3].as_number().unwrap() as i64, 4);
    }

    #[test]
    fn test_pull_at() {
        let runtime = setup_runtime();
        let expr = runtime.compile("pull_at(@, `[0, 2]`)").unwrap();
        let data = Variable::Array(vec![
            Rc::new(Variable::String("a".to_string())),
            Rc::new(Variable::String("b".to_string())),
            Rc::new(Variable::String("c".to_string())),
            Rc::new(Variable::String("d".to_string())),
        ]);
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 2);
        assert_eq!(arr[0].as_string().unwrap(), "a");
        assert_eq!(arr[1].as_string().unwrap(), "c");
    }

    #[test]
    fn test_pull_at_negative_index() {
        let runtime = setup_runtime();
        let expr = runtime.compile("pull_at(@, `[-1, -2]`)").unwrap();
        let data = Variable::Array(vec![
            Rc::new(Variable::String("a".to_string())),
            Rc::new(Variable::String("b".to_string())),
            Rc::new(Variable::String("c".to_string())),
            Rc::new(Variable::String("d".to_string())),
        ]);
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 2);
        assert_eq!(arr[0].as_string().unwrap(), "d");
        assert_eq!(arr[1].as_string().unwrap(), "c");
    }

    #[test]
    fn test_window() {
        let runtime = setup_runtime();
        let expr = runtime.compile("window(@, `3`)").unwrap();
        let data = Variable::Array(vec![
            Rc::new(Variable::Number(serde_json::Number::from(1))),
            Rc::new(Variable::Number(serde_json::Number::from(2))),
            Rc::new(Variable::Number(serde_json::Number::from(3))),
            Rc::new(Variable::Number(serde_json::Number::from(4))),
            Rc::new(Variable::Number(serde_json::Number::from(5))),
        ]);
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        // [1,2,3], [2,3,4], [3,4,5]
        assert_eq!(arr.len(), 3);
        let first = arr[0].as_array().unwrap();
        assert_eq!(first.len(), 3);
        assert_eq!(first[0].as_number().unwrap() as i64, 1);
        assert_eq!(first[1].as_number().unwrap() as i64, 2);
        assert_eq!(first[2].as_number().unwrap() as i64, 3);
    }

    #[test]
    fn test_window_with_step() {
        let runtime = setup_runtime();
        let expr = runtime.compile("window(@, `2`, `2`)").unwrap();
        let data = Variable::Array(vec![
            Rc::new(Variable::Number(serde_json::Number::from(1))),
            Rc::new(Variable::Number(serde_json::Number::from(2))),
            Rc::new(Variable::Number(serde_json::Number::from(3))),
            Rc::new(Variable::Number(serde_json::Number::from(4))),
            Rc::new(Variable::Number(serde_json::Number::from(5))),
        ]);
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        // [1,2], [3,4]
        assert_eq!(arr.len(), 2);
    }

    #[test]
    fn test_window_empty_result() {
        let runtime = setup_runtime();
        let expr = runtime.compile("window(@, `5`)").unwrap();
        let data = Variable::Array(vec![
            Rc::new(Variable::Number(serde_json::Number::from(1))),
            Rc::new(Variable::Number(serde_json::Number::from(2))),
        ]);
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 0);
    }

    #[test]
    fn test_combinations() {
        let runtime = setup_runtime();
        let expr = runtime.compile("combinations(@, `2`)").unwrap();
        let data = Variable::Array(vec![
            Rc::new(Variable::Number(serde_json::Number::from(1))),
            Rc::new(Variable::Number(serde_json::Number::from(2))),
            Rc::new(Variable::Number(serde_json::Number::from(3))),
        ]);
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        // C(3,2) = 3: [1,2], [1,3], [2,3]
        assert_eq!(arr.len(), 3);
    }

    #[test]
    fn test_combinations_k_zero() {
        let runtime = setup_runtime();
        let expr = runtime.compile("combinations(@, `0`)").unwrap();
        let data = Variable::Array(vec![
            Rc::new(Variable::Number(serde_json::Number::from(1))),
            Rc::new(Variable::Number(serde_json::Number::from(2))),
        ]);
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        // C(n,0) = 1 (the empty set)
        assert_eq!(arr.len(), 1);
        assert_eq!(arr[0].as_array().unwrap().len(), 0);
    }

    #[test]
    fn test_combinations_k_equals_n() {
        let runtime = setup_runtime();
        let expr = runtime.compile("combinations(@, `3`)").unwrap();
        let data = Variable::Array(vec![
            Rc::new(Variable::Number(serde_json::Number::from(1))),
            Rc::new(Variable::Number(serde_json::Number::from(2))),
            Rc::new(Variable::Number(serde_json::Number::from(3))),
        ]);
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        // C(3,3) = 1: [1,2,3]
        assert_eq!(arr.len(), 1);
        assert_eq!(arr[0].as_array().unwrap().len(), 3);
    }

    #[test]
    fn test_combinations_k_greater_than_n() {
        let runtime = setup_runtime();
        let expr = runtime.compile("combinations(@, `5`)").unwrap();
        let data = Variable::Array(vec![
            Rc::new(Variable::Number(serde_json::Number::from(1))),
            Rc::new(Variable::Number(serde_json::Number::from(2))),
        ]);
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        // C(2,5) = 0
        assert_eq!(arr.len(), 0);
    }

    // =========================================================================
    // zip tests
    // =========================================================================

    #[test]
    fn test_zip_basic() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"{"a": [1, 2, 3], "b": ["x", "y", "z"]}"#).unwrap();
        let expr = runtime.compile("zip(a, b)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 3);
        assert_eq!(arr[0].as_array().unwrap()[0].as_number().unwrap() as i64, 1);
        assert_eq!(arr[0].as_array().unwrap()[1].as_string().unwrap(), "x");
    }

    #[test]
    fn test_zip_unequal_lengths() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"{"a": [1, 2], "b": ["x", "y", "z"]}"#).unwrap();
        let expr = runtime.compile("zip(a, b)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        // Stops at shorter array
        assert_eq!(arr.len(), 2);
    }

    #[test]
    fn test_zip_empty_array() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"{"a": [], "b": [1, 2, 3]}"#).unwrap();
        let expr = runtime.compile("zip(a, b)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 0);
    }

    #[test]
    fn test_zip_with_objects() {
        let runtime = setup_runtime();
        let data =
            Variable::from_json(r#"{"names": ["Alice", "Bob"], "scores": [95, 87]}"#).unwrap();
        let expr = runtime.compile("zip(names, scores)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 2);
        assert_eq!(arr[0].as_array().unwrap()[0].as_string().unwrap(), "Alice");
        assert_eq!(
            arr[0].as_array().unwrap()[1].as_number().unwrap() as i64,
            95
        );
    }

    // =========================================================================
    // chunk tests
    // =========================================================================

    #[test]
    fn test_chunk_basic() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[1, 2, 3, 4, 5]"#).unwrap();
        let expr = runtime.compile("chunk(@, `2`)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 3); // [1,2], [3,4], [5]
        assert_eq!(arr[0].as_array().unwrap().len(), 2);
        assert_eq!(arr[2].as_array().unwrap().len(), 1);
    }

    #[test]
    fn test_chunk_exact_fit() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[1, 2, 3, 4, 5, 6]"#).unwrap();
        let expr = runtime.compile("chunk(@, `3`)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 2);
        assert_eq!(arr[0].as_array().unwrap().len(), 3);
        assert_eq!(arr[1].as_array().unwrap().len(), 3);
    }

    #[test]
    fn test_chunk_size_larger_than_array() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[1, 2, 3]"#).unwrap();
        let expr = runtime.compile("chunk(@, `10`)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 1);
        assert_eq!(arr[0].as_array().unwrap().len(), 3);
    }

    #[test]
    fn test_chunk_size_one() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[1, 2, 3]"#).unwrap();
        let expr = runtime.compile("chunk(@, `1`)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 3);
    }

    #[test]
    fn test_chunk_and_process_pipeline() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]"#).unwrap();
        let expr = runtime.compile("chunk(@, `3`)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        // [1,2,3], [4,5,6], [7,8,9], [10]
        assert_eq!(arr.len(), 4);
    }

    // =========================================================================
    // take tests
    // =========================================================================

    #[test]
    fn test_take_basic() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[1, 2, 3, 4, 5]"#).unwrap();
        let expr = runtime.compile("take(@, `3`)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 3);
        assert_eq!(arr[0].as_number().unwrap() as i64, 1);
        assert_eq!(arr[2].as_number().unwrap() as i64, 3);
    }

    #[test]
    fn test_take_more_than_length() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[1, 2]"#).unwrap();
        let expr = runtime.compile("take(@, `10`)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 2);
    }

    #[test]
    fn test_take_zero() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[1, 2, 3]"#).unwrap();
        let expr = runtime.compile("take(@, `0`)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 0);
    }

    // =========================================================================
    // drop tests
    // =========================================================================

    #[test]
    fn test_drop_basic() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[1, 2, 3, 4, 5]"#).unwrap();
        let expr = runtime.compile("drop(@, `2`)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 3);
        assert_eq!(arr[0].as_number().unwrap() as i64, 3);
    }

    #[test]
    fn test_drop_more_than_length() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[1, 2]"#).unwrap();
        let expr = runtime.compile("drop(@, `10`)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 0);
    }

    #[test]
    fn test_drop_zero() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[1, 2, 3]"#).unwrap();
        let expr = runtime.compile("drop(@, `0`)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 3);
    }

    // =========================================================================
    // flatten_deep tests
    // =========================================================================

    #[test]
    fn test_flatten_deep_basic() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[[1, 2], [3, 4]]"#).unwrap();
        let expr = runtime.compile("flatten_deep(@)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 4);
    }

    #[test]
    fn test_flatten_deep_nested() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[1, [2, [3, [4, [5]]]]]"#).unwrap();
        let expr = runtime.compile("flatten_deep(@)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 5);
        assert_eq!(arr[4].as_number().unwrap() as i64, 5);
    }

    #[test]
    fn test_flatten_deep_already_flat() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[1, 2, 3]"#).unwrap();
        let expr = runtime.compile("flatten_deep(@)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 3);
    }

    #[test]
    fn test_flatten_deep_mixed() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[1, [2, 3], [[4]], [[[5, 6]]]]"#).unwrap();
        let expr = runtime.compile("flatten_deep(@)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 6);
    }

    // =========================================================================
    // flatten tests (single-level)
    // =========================================================================

    #[test]
    fn test_flatten_basic() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[[1, 2], [3, 4]]"#).unwrap();
        let expr = runtime.compile("flatten(@)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 4);
    }

    #[test]
    fn test_flatten_single_level_only() {
        let runtime = setup_runtime();
        // flatten should only go one level deep
        let data = Variable::from_json(r#"[1, [2, [3, 4]]]"#).unwrap();
        let expr = runtime.compile("flatten(@)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        // Should be [1, 2, [3, 4]] - 3 elements, not 4
        assert_eq!(arr.len(), 3);
        // The third element should still be an array
        assert!(arr[2].as_array().is_some());
    }

    #[test]
    fn test_flatten_already_flat() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[1, 2, 3]"#).unwrap();
        let expr = runtime.compile("flatten(@)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 3);
    }

    #[test]
    fn test_flatten_mixed_nesting() {
        let runtime = setup_runtime();
        // [[1], [[2]], [[[3]]]] should become [1, [2], [[3]]]
        let data = Variable::from_json(r#"[[1], [[2]], [[[3]]]]"#).unwrap();
        let expr = runtime.compile("flatten(@)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 3);
        // First element is a number
        assert!(arr[0].as_number().is_some());
        // Second element is [2]
        assert!(arr[1].as_array().is_some());
        // Third element is [[3]]
        assert!(arr[2].as_array().is_some());
    }

    #[test]
    fn test_flatten_empty() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[]"#).unwrap();
        let expr = runtime.compile("flatten(@)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 0);
    }

    // =========================================================================
    // compact tests
    // =========================================================================

    #[test]
    fn test_compact_basic() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[1, null, 2, false, 3]"#).unwrap();
        let expr = runtime.compile("compact(@)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 3);
    }

    #[test]
    fn test_compact_keeps_zero_and_empty_string() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[0, "", null, true]"#).unwrap();
        let expr = runtime.compile("compact(@)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 3); // 0, "", true
    }

    #[test]
    fn test_compact_all_falsy() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[null, false, null]"#).unwrap();
        let expr = runtime.compile("compact(@)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 0);
    }

    // =========================================================================
    // index_at tests
    // =========================================================================

    #[test]
    fn test_index_at_positive() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"["a", "b", "c", "d"]"#).unwrap();
        let expr = runtime.compile("index_at(@, `2`)").unwrap();
        let result = expr.search(&data).unwrap();
        assert_eq!(result.as_string().unwrap(), "c");
    }

    #[test]
    fn test_index_at_negative() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"["a", "b", "c", "d"]"#).unwrap();
        let expr = runtime.compile("index_at(@, `-1`)").unwrap();
        let result = expr.search(&data).unwrap();
        assert_eq!(result.as_string().unwrap(), "d");
    }

    #[test]
    fn test_index_at_negative_second() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"["a", "b", "c", "d"]"#).unwrap();
        let expr = runtime.compile("index_at(@, `-2`)").unwrap();
        let result = expr.search(&data).unwrap();
        assert_eq!(result.as_string().unwrap(), "c");
    }

    #[test]
    fn test_index_at_out_of_bounds() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"["a", "b", "c"]"#).unwrap();
        let expr = runtime.compile("index_at(@, `10`)").unwrap();
        let result = expr.search(&data).unwrap();
        assert!(result.is_null());
    }

    // =========================================================================
    // includes tests
    // =========================================================================

    #[test]
    fn test_includes_number() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[1, 2, 3, 4, 5]"#).unwrap();
        let expr = runtime.compile("includes(@, `3`)").unwrap();
        let result = expr.search(&data).unwrap();
        assert!(result.as_boolean().unwrap());
    }

    #[test]
    fn test_includes_not_found() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[1, 2, 3]"#).unwrap();
        let expr = runtime.compile("includes(@, `10`)").unwrap();
        let result = expr.search(&data).unwrap();
        assert!(!result.as_boolean().unwrap());
    }

    #[test]
    fn test_includes_string() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"["apple", "banana", "cherry"]"#).unwrap();
        let expr = runtime.compile(r#"includes(@, `"banana"`)"#).unwrap();
        let result = expr.search(&data).unwrap();
        assert!(result.as_boolean().unwrap());
    }

    #[test]
    fn test_includes_object() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[{"a": 1}, {"b": 2}]"#).unwrap();
        let expr = runtime.compile(r#"includes(@, `{"a": 1}`)"#).unwrap();
        let result = expr.search(&data).unwrap();
        assert!(result.as_boolean().unwrap());
    }

    // =========================================================================
    // find_index tests
    // =========================================================================

    #[test]
    fn test_find_index_found() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"["a", "b", "c", "d"]"#).unwrap();
        let expr = runtime.compile(r#"find_index(@, `"c"`)"#).unwrap();
        let result = expr.search(&data).unwrap();
        assert_eq!(result.as_number().unwrap() as i64, 2);
    }

    #[test]
    fn test_find_index_not_found() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"["a", "b", "c"]"#).unwrap();
        let expr = runtime.compile(r#"find_index(@, `"z"`)"#).unwrap();
        let result = expr.search(&data).unwrap();
        assert_eq!(result.as_number().unwrap() as i64, -1);
    }

    // =========================================================================
    // group_by tests
    // =========================================================================

    #[test]
    fn test_group_by_basic() {
        let runtime = setup_runtime();
        let data = Variable::from_json(
            r#"[{"type": "a", "v": 1}, {"type": "b", "v": 2}, {"type": "a", "v": 3}]"#,
        )
        .unwrap();
        let expr = runtime.compile(r#"group_by(@, `"type"`)"#).unwrap();
        let result = expr.search(&data).unwrap();
        let obj = result.as_object().unwrap();
        assert_eq!(obj.get("a").unwrap().as_array().unwrap().len(), 2);
        assert_eq!(obj.get("b").unwrap().as_array().unwrap().len(), 1);
    }

    // =========================================================================
    // index_by tests
    // =========================================================================

    #[test]
    fn test_index_by_basic() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[{"id": 1, "name": "alice"}, {"id": 2, "name": "bob"}]"#)
            .unwrap();
        let expr = runtime.compile(r#"index_by(@, `"id"`)"#).unwrap();
        let result = expr.search(&data).unwrap();
        let obj = result.as_object().unwrap();
        assert_eq!(obj.len(), 2);
        // Keys are string versions of the id
        let alice = obj.get("1").unwrap().as_object().unwrap();
        assert_eq!(alice.get("name").unwrap().as_string().unwrap(), "alice");
        let bob = obj.get("2").unwrap().as_object().unwrap();
        assert_eq!(bob.get("name").unwrap().as_string().unwrap(), "bob");
    }

    #[test]
    fn test_index_by_string_key() {
        let runtime = setup_runtime();
        let data = Variable::from_json(
            r#"[{"code": "US", "name": "United States"}, {"code": "UK", "name": "United Kingdom"}]"#,
        )
        .unwrap();
        let expr = runtime.compile(r#"index_by(@, `"code"`)"#).unwrap();
        let result = expr.search(&data).unwrap();
        let obj = result.as_object().unwrap();
        assert_eq!(obj.len(), 2);
        let us = obj.get("US").unwrap().as_object().unwrap();
        assert_eq!(
            us.get("name").unwrap().as_string().unwrap(),
            "United States"
        );
    }

    #[test]
    fn test_index_by_duplicate_keys() {
        // Last value wins for duplicate keys
        let runtime = setup_runtime();
        let data = Variable::from_json(
            r#"[{"type": "a", "v": 1}, {"type": "a", "v": 2}, {"type": "a", "v": 3}]"#,
        )
        .unwrap();
        let expr = runtime.compile(r#"index_by(@, `"type"`)"#).unwrap();
        let result = expr.search(&data).unwrap();
        let obj = result.as_object().unwrap();
        assert_eq!(obj.len(), 1);
        // Last value (v=3) wins
        let a = obj.get("a").unwrap().as_object().unwrap();
        assert_eq!(a.get("v").unwrap().as_number().unwrap() as i64, 3);
    }

    #[test]
    fn test_index_by_missing_key() {
        // Items without the key field are skipped
        let runtime = setup_runtime();
        let data = Variable::from_json(
            r#"[{"id": 1, "name": "alice"}, {"name": "bob"}, {"id": 3, "name": "charlie"}]"#,
        )
        .unwrap();
        let expr = runtime.compile(r#"index_by(@, `"id"`)"#).unwrap();
        let result = expr.search(&data).unwrap();
        let obj = result.as_object().unwrap();
        assert_eq!(obj.len(), 2);
        assert!(obj.contains_key("1"));
        assert!(obj.contains_key("3"));
        assert!(!obj.contains_key("2")); // bob was skipped
    }

    #[test]
    fn test_index_by_empty_array() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[]"#).unwrap();
        let expr = runtime.compile(r#"index_by(@, `"id"`)"#).unwrap();
        let result = expr.search(&data).unwrap();
        let obj = result.as_object().unwrap();
        assert!(obj.is_empty());
    }

    // =========================================================================
    // nth tests
    // =========================================================================

    #[test]
    fn test_nth_every_second() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[1, 2, 3, 4, 5, 6]"#).unwrap();
        let expr = runtime.compile("nth(@, `2`)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 3); // 1, 3, 5
        assert_eq!(arr[0].as_number().unwrap() as i64, 1);
        assert_eq!(arr[1].as_number().unwrap() as i64, 3);
        assert_eq!(arr[2].as_number().unwrap() as i64, 5);
    }

    #[test]
    fn test_nth_every_third() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[1, 2, 3, 4, 5, 6, 7, 8, 9]"#).unwrap();
        let expr = runtime.compile("nth(@, `3`)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 3); // 1, 4, 7
    }

    // =========================================================================
    // interleave tests
    // =========================================================================

    #[test]
    fn test_interleave_equal() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"{"a": [1, 2, 3], "b": ["a", "b", "c"]}"#).unwrap();
        let expr = runtime.compile("interleave(a, b)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 6);
        assert_eq!(arr[0].as_number().unwrap() as i64, 1);
        assert_eq!(arr[1].as_string().unwrap(), "a");
        assert_eq!(arr[2].as_number().unwrap() as i64, 2);
        assert_eq!(arr[3].as_string().unwrap(), "b");
    }

    #[test]
    fn test_interleave_unequal() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"{"a": [1, 2], "b": ["a", "b", "c"]}"#).unwrap();
        let expr = runtime.compile("interleave(a, b)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 5); // 1, a, 2, b, c
    }

    // =========================================================================
    // rotate tests
    // =========================================================================

    #[test]
    fn test_rotate_left() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[1, 2, 3, 4, 5]"#).unwrap();
        let expr = runtime.compile("rotate(@, `2`)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr[0].as_number().unwrap() as i64, 3);
        assert_eq!(arr[4].as_number().unwrap() as i64, 2);
    }

    #[test]
    fn test_rotate_right() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[1, 2, 3, 4, 5]"#).unwrap();
        let expr = runtime.compile("rotate(@, `-1`)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr[0].as_number().unwrap() as i64, 5);
        assert_eq!(arr[1].as_number().unwrap() as i64, 1);
    }

    // =========================================================================
    // partition tests
    // =========================================================================

    #[test]
    fn test_partition_even() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[1, 2, 3, 4, 5, 6]"#).unwrap();
        let expr = runtime.compile("partition(@, `2`)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 2);
        assert_eq!(arr[0].as_array().unwrap().len(), 3);
        assert_eq!(arr[1].as_array().unwrap().len(), 3);
    }

    #[test]
    fn test_partition_uneven() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[1, 2, 3, 4, 5]"#).unwrap();
        let expr = runtime.compile("partition(@, `3`)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 3);
    }

    // =========================================================================
    // set operations tests
    // =========================================================================

    #[test]
    fn test_difference() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"{"a": [1, 2, 3, 4], "b": [2, 4]}"#).unwrap();
        let expr = runtime.compile("difference(a, b)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 2); // 1, 3
    }

    #[test]
    fn test_intersection() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"{"a": [1, 2, 3], "b": [2, 3, 4]}"#).unwrap();
        let expr = runtime.compile("intersection(a, b)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 2); // 2, 3
    }

    #[test]
    fn test_union() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"{"a": [1, 2], "b": [2, 3]}"#).unwrap();
        let expr = runtime.compile("union(a, b)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 3); // 1, 2, 3
    }

    // =========================================================================
    // frequencies tests
    // =========================================================================

    #[test]
    fn test_frequencies_basic() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"["a", "b", "a", "c", "a", "b"]"#).unwrap();
        let expr = runtime.compile("frequencies(@)").unwrap();
        let result = expr.search(&data).unwrap();
        let obj = result.as_object().unwrap();
        assert_eq!(obj.get("a").unwrap().as_number().unwrap() as i64, 3);
        assert_eq!(obj.get("b").unwrap().as_number().unwrap() as i64, 2);
        assert_eq!(obj.get("c").unwrap().as_number().unwrap() as i64, 1);
    }

    #[test]
    fn test_frequencies_numbers() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[1, 2, 1, 1, 2, 3]"#).unwrap();
        let expr = runtime.compile("frequencies(@)").unwrap();
        let result = expr.search(&data).unwrap();
        let obj = result.as_object().unwrap();
        assert_eq!(obj.get("1").unwrap().as_number().unwrap() as i64, 3);
        assert_eq!(obj.get("2").unwrap().as_number().unwrap() as i64, 2);
    }

    // =========================================================================
    // mode tests
    // =========================================================================

    #[test]
    fn test_mode_basic() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[1, 2, 2, 3, 2, 4]"#).unwrap();
        let expr = runtime.compile("mode(@)").unwrap();
        let result = expr.search(&data).unwrap();
        assert_eq!(result.as_number().unwrap() as i64, 2);
    }

    #[test]
    fn test_mode_empty() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[]"#).unwrap();
        let expr = runtime.compile("mode(@)").unwrap();
        let result = expr.search(&data).unwrap();
        assert!(result.is_null());
    }

    // =========================================================================
    // cartesian tests
    // =========================================================================

    #[test]
    fn test_cartesian_basic() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"{"a": [1, 2], "b": ["x", "y"]}"#).unwrap();
        let expr = runtime.compile("cartesian(a, b)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 4); // [1,x], [1,y], [2,x], [2,y]
    }

    #[test]
    fn test_cartesian_empty() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"{"a": [], "b": [1, 2]}"#).unwrap();
        let expr = runtime.compile("cartesian(a, b)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 0);
    }

    #[test]
    fn test_cartesian_n_way_two_arrays() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[[1, 2], ["a", "b"]]"#).unwrap();
        let expr = runtime.compile("cartesian(@)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 4); // [1,a], [1,b], [2,a], [2,b]
    }

    #[test]
    fn test_cartesian_n_way_three_arrays() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[[1, 2], ["a", "b"], [true, false]]"#).unwrap();
        let expr = runtime.compile("cartesian(@)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 8); // 2 * 2 * 2 = 8 combinations
    }

    #[test]
    fn test_cartesian_n_way_single_array() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[[1, 2, 3]]"#).unwrap();
        let expr = runtime.compile("cartesian(@)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 3); // [1], [2], [3]
    }

    #[test]
    fn test_cartesian_n_way_empty() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[]"#).unwrap();
        let expr = runtime.compile("cartesian(@)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 0);
    }

    // =========================================================================
    // Edge cases
    // =========================================================================

    #[test]
    fn test_first_empty_array() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[]"#).unwrap();
        let expr = runtime.compile("first(@)").unwrap();
        let result = expr.search(&data).unwrap();
        assert!(result.is_null());
    }

    #[test]
    fn test_last_empty_array() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[]"#).unwrap();
        let expr = runtime.compile("last(@)").unwrap();
        let result = expr.search(&data).unwrap();
        assert!(result.is_null());
    }

    #[test]
    fn test_unique_preserves_order() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"["c", "a", "b", "a", "c"]"#).unwrap();
        let expr = runtime.compile("unique(@)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 3);
        assert_eq!(arr[0].as_string().unwrap(), "c");
        assert_eq!(arr[1].as_string().unwrap(), "a");
        assert_eq!(arr[2].as_string().unwrap(), "b");
    }

    #[test]
    fn test_unique_different_types() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[1, "1", 1, "1"]"#).unwrap();
        let expr = runtime.compile("unique(@)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 2); // 1 and "1" are different
    }

    #[test]
    fn test_range_with_step() {
        let runtime = setup_runtime();
        let data = Variable::Null;
        let expr = runtime.compile("range(`1`, `10`, `2`)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 5); // 1, 3, 5, 7, 9
        assert_eq!(arr[0].as_number().unwrap() as i64, 1);
        assert_eq!(arr[4].as_number().unwrap() as i64, 9);
    }

    #[test]
    fn test_range_descending() {
        let runtime = setup_runtime();
        let data = Variable::Null;
        let expr = runtime.compile("range(`5`, `0`, `-1`)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 5); // 5, 4, 3, 2, 1
        assert_eq!(arr[0].as_number().unwrap() as i64, 5);
        assert_eq!(arr[4].as_number().unwrap() as i64, 1);
    }

    // =========================================================================
    // Pipeline patterns with arrays
    // =========================================================================

    #[test]
    fn test_pipeline_unique_sort() {
        let runtime = setup_runtime();
        let data =
            Variable::from_json(r#"["redis", "database", "redis", "nosql", "database"]"#).unwrap();
        let expr = runtime.compile("unique(@) | sort(@)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 3);
        assert_eq!(arr[0].as_string().unwrap(), "database");
        assert_eq!(arr[1].as_string().unwrap(), "nosql");
        assert_eq!(arr[2].as_string().unwrap(), "redis");
    }

    #[test]
    fn test_pipeline_filter_take() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]"#).unwrap();
        let expr = runtime.compile("[?@ > `3`] | take(@, `3`)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 3);
        assert_eq!(arr[0].as_number().unwrap() as i64, 4);
        assert_eq!(arr[1].as_number().unwrap() as i64, 5);
        assert_eq!(arr[2].as_number().unwrap() as i64, 6);
    }

    #[test]
    fn test_pipeline_flatten_unique() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[[1, 2], [2, 3], [3, 4]]"#).unwrap();
        let expr = runtime.compile("flatten_deep(@) | unique(@)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 4); // 1, 2, 3, 4
    }

    #[test]
    fn test_large_array_processing() {
        let runtime = setup_runtime();
        // Create array with 1000 elements
        let items: Vec<i32> = (1..=1000).collect();
        let json = serde_json::to_string(&items).unwrap();
        let data = Variable::from_json(&json).unwrap();

        let expr = runtime.compile("length(@)").unwrap();
        let result = expr.search(&data).unwrap();
        assert_eq!(result.as_number().unwrap() as i64, 1000);
    }

    #[test]
    fn test_transpose_basic() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[[1, 2, 3], [4, 5, 6]]"#).unwrap();
        let expr = runtime.compile("transpose(@)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 3);
        // First column: [1, 4]
        let col0 = arr[0].as_array().unwrap();
        assert_eq!(col0[0].as_number().unwrap() as i64, 1);
        assert_eq!(col0[1].as_number().unwrap() as i64, 4);
        // Second column: [2, 5]
        let col1 = arr[1].as_array().unwrap();
        assert_eq!(col1[0].as_number().unwrap() as i64, 2);
        assert_eq!(col1[1].as_number().unwrap() as i64, 5);
    }

    #[test]
    fn test_transpose_empty() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[]"#).unwrap();
        let expr = runtime.compile("transpose(@)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 0);
    }

    #[test]
    fn test_transpose_unequal_rows() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[[1, 2], [3, 4, 5], [6, 7]]"#).unwrap();
        let expr = runtime.compile("transpose(@)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        // Should use minimum length (2)
        assert_eq!(arr.len(), 2);
    }

    #[test]
    fn test_pairwise_basic() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[1, 2, 3, 4]"#).unwrap();
        let expr = runtime.compile("pairwise(@)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 3);
        // First pair: [1, 2]
        let pair0 = arr[0].as_array().unwrap();
        assert_eq!(pair0[0].as_number().unwrap() as i64, 1);
        assert_eq!(pair0[1].as_number().unwrap() as i64, 2);
        // Second pair: [2, 3]
        let pair1 = arr[1].as_array().unwrap();
        assert_eq!(pair1[0].as_number().unwrap() as i64, 2);
        assert_eq!(pair1[1].as_number().unwrap() as i64, 3);
    }

    #[test]
    fn test_pairwise_short_array() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[1]"#).unwrap();
        let expr = runtime.compile("pairwise(@)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 0);
    }

    #[test]
    fn test_sliding_window_alias() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[1, 2, 3, 4, 5]"#).unwrap();
        let expr = runtime.compile("sliding_window(@, `3`)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 3);
        // First window: [1, 2, 3]
        let win0 = arr[0].as_array().unwrap();
        assert_eq!(win0.len(), 3);
        assert_eq!(win0[0].as_number().unwrap() as i64, 1);
    }

    // indices_array tests
    #[test]
    fn test_indices_array_found() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[1, 2, 3, 2, 4, 2]"#).unwrap();
        let expr = runtime.compile("indices_array(@, `2`)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 3);
        assert_eq!(arr[0].as_number().unwrap() as i64, 1);
        assert_eq!(arr[1].as_number().unwrap() as i64, 3);
        assert_eq!(arr[2].as_number().unwrap() as i64, 5);
    }

    #[test]
    fn test_indices_array_not_found() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[1, 2, 3]"#).unwrap();
        let expr = runtime.compile("indices_array(@, `5`)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 0);
    }

    #[test]
    fn test_indices_array_strings() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"["a", "b", "a", "c", "a"]"#).unwrap();
        let expr = runtime.compile(r#"indices_array(@, `"a"`)"#).unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 3);
        assert_eq!(arr[0].as_number().unwrap() as i64, 0);
        assert_eq!(arr[1].as_number().unwrap() as i64, 2);
        assert_eq!(arr[2].as_number().unwrap() as i64, 4);
    }

    // inside_array tests
    #[test]
    fn test_inside_array_true() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"{"a": [1, 2], "b": [1, 2, 3, 4]}"#).unwrap();
        let expr = runtime.compile("inside_array(a, b)").unwrap();
        let result = expr.search(&data).unwrap();
        assert!(result.as_boolean().unwrap());
    }

    #[test]
    fn test_inside_array_false() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"{"a": [1, 5], "b": [1, 2, 3, 4]}"#).unwrap();
        let expr = runtime.compile("inside_array(a, b)").unwrap();
        let result = expr.search(&data).unwrap();
        assert!(!result.as_boolean().unwrap());
    }

    #[test]
    fn test_inside_array_empty() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"{"a": [], "b": [1, 2, 3]}"#).unwrap();
        let expr = runtime.compile("inside_array(a, b)").unwrap();
        let result = expr.search(&data).unwrap();
        assert!(result.as_boolean().unwrap());
    }

    // bsearch tests
    #[test]
    fn test_bsearch_found() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[1, 3, 5, 7, 9]"#).unwrap();
        let expr = runtime.compile("bsearch(@, `5`)").unwrap();
        let result = expr.search(&data).unwrap();
        assert_eq!(result.as_number().unwrap() as i64, 2);
    }

    #[test]
    fn test_bsearch_not_found_middle() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[1, 3, 5, 7, 9]"#).unwrap();
        let expr = runtime.compile("bsearch(@, `4`)").unwrap();
        let result = expr.search(&data).unwrap();
        assert_eq!(result.as_number().unwrap() as i64, -3);
    }

    #[test]
    fn test_bsearch_not_found_start() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[1, 3, 5, 7, 9]"#).unwrap();
        let expr = runtime.compile("bsearch(@, `0`)").unwrap();
        let result = expr.search(&data).unwrap();
        assert_eq!(result.as_number().unwrap() as i64, -1);
    }

    #[test]
    fn test_bsearch_not_found_end() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[1, 3, 5, 7, 9]"#).unwrap();
        let expr = runtime.compile("bsearch(@, `10`)").unwrap();
        let result = expr.search(&data).unwrap();
        assert_eq!(result.as_number().unwrap() as i64, -6);
    }

    #[test]
    fn test_bsearch_empty_array() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[]"#).unwrap();
        let expr = runtime.compile("bsearch(@, `5`)").unwrap();
        let result = expr.search(&data).unwrap();
        assert_eq!(result.as_number().unwrap() as i64, -1);
    }

    // repeat_array tests
    #[test]
    fn test_repeat_array_basic() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"null"#).unwrap();
        let expr = runtime.compile("repeat_array(`1`, `3`)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 3);
        assert_eq!(arr[0].as_number().unwrap() as i64, 1);
        assert_eq!(arr[1].as_number().unwrap() as i64, 1);
        assert_eq!(arr[2].as_number().unwrap() as i64, 1);
    }

    #[test]
    fn test_repeat_array_string() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"null"#).unwrap();
        let expr = runtime.compile(r#"repeat_array(`"x"`, `4`)"#).unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 4);
        assert_eq!(arr[0].as_string().unwrap(), "x");
        assert_eq!(arr[3].as_string().unwrap(), "x");
    }

    #[test]
    fn test_repeat_array_zero() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"null"#).unwrap();
        let expr = runtime.compile("repeat_array(`1`, `0`)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 0);
    }

    #[test]
    fn test_repeat_array_object() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"null"#).unwrap();
        let expr = runtime.compile(r#"repeat_array(`{"a": 1}`, `2`)"#).unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 2);
        assert_eq!(
            arr[0]
                .as_object()
                .unwrap()
                .get("a")
                .unwrap()
                .as_number()
                .unwrap() as i64,
            1
        );
    }

    // cycle tests
    #[test]
    fn test_cycle_basic() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[1, 2, 3]"#).unwrap();
        let expr = runtime.compile("cycle(@, `2`)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 6);
        assert_eq!(arr[0].as_number().unwrap() as i64, 1);
        assert_eq!(arr[1].as_number().unwrap() as i64, 2);
        assert_eq!(arr[2].as_number().unwrap() as i64, 3);
        assert_eq!(arr[3].as_number().unwrap() as i64, 1);
        assert_eq!(arr[4].as_number().unwrap() as i64, 2);
        assert_eq!(arr[5].as_number().unwrap() as i64, 3);
    }

    #[test]
    fn test_cycle_strings() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"["a", "b"]"#).unwrap();
        let expr = runtime.compile("cycle(@, `3`)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 6);
        assert_eq!(arr[0].as_string().unwrap(), "a");
        assert_eq!(arr[1].as_string().unwrap(), "b");
        assert_eq!(arr[2].as_string().unwrap(), "a");
    }

    #[test]
    fn test_cycle_zero() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[1, 2, 3]"#).unwrap();
        let expr = runtime.compile("cycle(@, `0`)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 0);
    }

    #[test]
    fn test_cycle_empty_array() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[]"#).unwrap();
        let expr = runtime.compile("cycle(@, `5`)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 0);
    }

    #[test]
    fn test_cycle_once() {
        let runtime = setup_runtime();
        let data = Variable::from_json(r#"[1, 2]"#).unwrap();
        let expr = runtime.compile("cycle(@, `1`)").unwrap();
        let result = expr.search(&data).unwrap();
        let arr = result.as_array().unwrap();
        assert_eq!(arr.len(), 2);
        assert_eq!(arr[0].as_number().unwrap() as i64, 1);
        assert_eq!(arr[1].as_number().unwrap() as i64, 2);
    }
}