hamelin_eval 0.10.13

//! Reverse evaluation for typed expressions
//!
//! This module implements reverse evaluation, which takes an output constraint
//! and an expression AST, then calculates the domain of input values that would
//! produce outputs satisfying that constraint.

use std::sync::Arc;

use crate::eval::environment::Environment;
use crate::eval::error::{EvalError, EvalResult};
use crate::eval::timestamp::next_truncation_boundary;
use crate::reverse_eval::domain::Constraint;
use crate::value::Value;
use hamelin_lib::func::def::ParameterBinding;
use hamelin_lib::tree::ast::expression::{Expression, ExpressionKind};
use hamelin_lib::tree::typed_ast::expression::{
    TypedApply, TypedExpression, TypedExpressionKind, TypedTsTrunc,
};
/// Main entry point for reverse evaluation
///
/// Takes a typed expression and an output constraint, then calculates the constraint
/// on the input variable that would produce outputs satisfying that constraint.
///
/// This function only supports expressions with a single variable. Expressions with
/// multiple variables will return an error.
///
/// # Arguments
/// * `expr` - The typed expression to reverse evaluate
/// * `output_constraint` - The constraint that the output must satisfy
/// * `env` - Environment with constant values (for evaluating constant subexpressions)
///
/// # Returns
/// * `Some(Constraint)` - The constraint on the single variable in the expression
/// * `None` - The expression contains no variables (e.g., it's a literal)
/// * `Err` - The constraints are unsatisfiable, multiple variables exist, or the operation is not supported
pub fn reverse_eval(
    expr: &TypedExpression,
    output_constraint: Constraint,
    env: &Environment,
) -> EvalResult<Option<Constraint>> {
    // Handle empty constraint early - nothing can satisfy it
    if let Constraint::Empty = output_constraint {
        return Ok(None);
    }

    match &expr.kind {
        TypedExpressionKind::Leaf => reverse_eval_leaf(&expr.ast, output_constraint),

        TypedExpressionKind::FieldReference(_) => {
            // Column reference - return the output constraint as the input constraint
            Ok(Some(output_constraint))
        }

        TypedExpressionKind::Apply(apply) => reverse_eval_apply(apply, output_constraint, env),

        TypedExpressionKind::VariantIndexAccess(_) => Err(EvalError::NotImplemented {
            feature: "Reverse evaluation for VariantIndexAccess".to_string(),
        }),

        TypedExpressionKind::Cast(_) => Err(EvalError::NotImplemented {
            feature: "Reverse evaluation for Cast expressions".to_string(),
        }),

        TypedExpressionKind::TsTrunc(ts_trunc) => {
            reverse_eval_ts_trunc(ts_trunc, output_constraint, env)
        }

        _ => Err(EvalError::NotImplemented {
            feature: format!("Reverse evaluation for {:?}", expr.kind),
        }),
    }
}

/// Handle leaf expressions (literals and column references)
fn reverse_eval_leaf(
    expr: &Arc<Expression>,
    output_constraint: Constraint,
) -> EvalResult<Option<Constraint>> {
    match &expr.kind {
        ExpressionKind::FieldReference(_col_ref) => {
            // Base case: variable gets the output constraint
            Ok(Some(output_constraint))
        }

        ExpressionKind::IntLiteral(lit) => {
            // Check if constant satisfies constraint
            let value = Value::Int(lit.int);
            if output_constraint.is_satisfied_by(&value) {
                Ok(None) // No variables to constrain
            } else {
                Err(EvalError::execution("Constraint is unsatisfiable"))
            }
        }

        ExpressionKind::BooleanLiteral(lit) => {
            let value = Value::Boolean(lit.value);
            if output_constraint.is_satisfied_by(&value) {
                Ok(None)
            } else {
                Err(EvalError::execution("Constraint is unsatisfiable"))
            }
        }

        ExpressionKind::StringLiteral(lit) => {
            let value = Value::String(lit.value.clone());
            if output_constraint.is_satisfied_by(&value) {
                Ok(None)
            } else {
                Err(EvalError::execution("Constraint is unsatisfiable"))
            }
        }

        ExpressionKind::ScientificLiteral(lit) => {
            let value = Value::Double(lit.value);
            if output_constraint.is_satisfied_by(&value) {
                Ok(None)
            } else {
                Err(EvalError::execution("Constraint is unsatisfiable"))
            }
        }

        ExpressionKind::DoubleLiteral(lit) => {
            let value = Value::Double(lit.value);
            if output_constraint.is_satisfied_by(&value) {
                Ok(None)
            } else {
                Err(EvalError::execution("Constraint is unsatisfiable"))
            }
        }

        ExpressionKind::NullLiteral(_) => {
            let value = Value::Null;
            if output_constraint.is_satisfied_by(&value) {
                Ok(None)
            } else {
                Err(EvalError::execution("Constraint is unsatisfiable"))
            }
        }

        _ => Err(EvalError::NotImplemented {
            feature: format!("Reverse evaluation for leaf expression: {:?}", expr.kind),
        }),
    }
}

/// Reverse evaluate function applications (operators, functions, array/map indexing)
fn reverse_eval_apply(
    apply: &TypedApply,
    output_constraint: Constraint,
    env: &Environment,
) -> EvalResult<Option<Constraint>> {
    let func_def = &apply.function_def;

    // Evaluate constant parameters and find the variable path
    let (eval_params, variable_param) = evaluate_constant_params(&apply.parameter_binding, env)?;

    // Handle based on whether we have variables
    match variable_param {
        Some(variable_expr) => {
            // Look up the function's reverse implementation in the registry
            let type_id = func_def.type_id();
            let input_constraint = env
                .registry()
                .reverse(type_id, &output_constraint, eval_params)
                .ok_or_else(|| {
                    EvalError::execution(format!(
                        "Function '{}' has no reverse evaluation implementation",
                        func_def.name()
                    ))
                })??;
            // Recurse into the parameter with variables
            reverse_eval(variable_expr, input_constraint, env)
        }
        None => Ok(None), // All parameters were constants - no variables to constrain
    }
}

/// Reverse evaluate timestamp truncation
///
/// Given an output constraint (a truncated timestamp value), determine what
/// input timestamps would produce that output when truncated.
fn reverse_eval_ts_trunc(
    ts_trunc: &TypedTsTrunc,
    output_constraint: Constraint,
    env: &Environment,
) -> EvalResult<Option<Constraint>> {
    match output_constraint {
        Constraint::Empty => Ok(None),
        Constraint::Universal => {
            // If output can be any timestamp, input can also be any timestamp
            reverse_eval(&ts_trunc.expression, Constraint::Universal, env)
        }
        Constraint::Equals(Value::Timestamp(truncated_ts)) => {
            // The output is a specific truncated timestamp
            // Find the range of timestamps that would truncate to this value

            // The input must be in the range [truncated_ts, next_boundary)
            let next =
                next_truncation_boundary(&truncated_ts, &ts_trunc.unit, ts_trunc.multiplier)?;

            // Create a range constraint for the input
            let input_constraint = Constraint::Range {
                min: Some(truncated_ts.into()),
                max: Some(next.into()),
            };

            reverse_eval(&ts_trunc.expression, input_constraint, env)
        }
        Constraint::Equals(_) => {
            // Output must be a timestamp, not another type
            Err(EvalError::execution(
                "TsTrunc output must be a timestamp value",
            ))
        }
        Constraint::Range { min, max } => {
            // Handle range constraints on the truncated output
            // This is more complex as we need to expand both boundaries

            // For simplicity, we'll handle the case where both min and max are timestamps
            match (min, max) {
                (Some(Value::Timestamp(min_ts)), Some(Value::Timestamp(max_ts))) => {
                    // The minimum input is min_ts (already truncated)
                    // The maximum input is just before the boundary after max_ts
                    let max_next =
                        next_truncation_boundary(&max_ts, &ts_trunc.unit, ts_trunc.multiplier)?;

                    let input_constraint = Constraint::Range {
                        min: Some(min_ts.clone().into()),
                        max: Some(max_next.into()),
                    };

                    reverse_eval(&ts_trunc.expression, input_constraint, env)
                }
                (Some(Value::Timestamp(min_ts)), None) => {
                    // Only minimum bound
                    let input_constraint = Constraint::Range {
                        min: Some(min_ts.clone().into()),
                        max: None,
                    };

                    reverse_eval(&ts_trunc.expression, input_constraint, env)
                }
                (None, Some(Value::Timestamp(max_ts))) => {
                    // Only maximum bound - need to find the next boundary
                    let max_next =
                        next_truncation_boundary(&max_ts, &ts_trunc.unit, ts_trunc.multiplier)?;

                    let input_constraint = Constraint::Range {
                        min: None,
                        max: Some(max_next.into()),
                    };

                    reverse_eval(&ts_trunc.expression, input_constraint, env)
                }
                _ => {
                    // Other cases or non-timestamp bounds
                    Err(EvalError::execution(
                        "TsTrunc range bounds must be timestamps",
                    ))
                }
            }
        }
    }
}

/// Helper function to evaluate constant parameters and identify the variable path
///
/// Returns a tuple of:
/// - ParameterBinding with Option<Value> for each parameter (Some if constant, None if has variables)
/// - The expression that contains variables (if any)
///
/// Errors if more than one parameter contains variables (not supported)
fn evaluate_constant_params<'a>(
    params: &'a ParameterBinding<Arc<TypedExpression>>,
    env: &Environment,
) -> EvalResult<(ParameterBinding<Option<Value>>, Option<&'a TypedExpression>)> {
    use crate::eval::evaluator::eval;

    // Collect all variable expressions (those that can't be evaluated to constants)
    let variable_exprs: Vec<&TypedExpression> = params
        .iter()
        .filter(|expr| eval(expr, env).is_err())
        .map(|expr| expr.as_ref())
        .collect();

    // Check if we have more than one variable expression
    if variable_exprs.len() > 1 {
        return Err(EvalError::execution(
            "Multiple parameters with variables not supported in reverse evaluation",
        ));
    }

    // Create the parameter binding with evaluated values (Some for constants, None for variables)
    let eval_params = params.clone().map(|expr| eval(&expr, env).ok());

    Ok((eval_params, variable_exprs.first().copied()))
}