hamelin_lib 0.21.7

use std::any::TypeId;
use std::collections::VecDeque;
use std::fmt::{Debug, Display, Formatter};
use std::sync::Arc;

use anyhow::{anyhow, bail};
use derive_more::From;
use itertools::Itertools;
use vecmap::VecMap;

pub use crate::func::err::{
    FunctionBindFailure, FunctionParameterBindingFailure, FunctionParameterBindingFailures,
    FunctionTranslationFailure, MatchTestFailure, NoBindingError,
};
use crate::translation::ExpressionTranslation;
use crate::tree::ast::expression::Expression;
use crate::tree::typed_ast::context::ExpressionTranslationContext;
use crate::tree::typed_ast::expression::TypedExpression;
use crate::types::matcher::Matcher;
use crate::types::Type;

pub trait HasType {
    fn typ(&self) -> &Type;
}

/// Trait that abstracts over parameter bindings to allow return type functions
/// to work with both ParameterBinding<TypedExpression> and ParameterBinding<ExpressionTranslation>
pub trait ParameterBindingProvider {
    fn get_by_name(&self, name: &str) -> anyhow::Result<&dyn HasType>;
    fn get_by_index(&self, index: usize) -> anyhow::Result<&dyn HasType>;
    fn iter(&self) -> Box<dyn Iterator<Item = &dyn HasType> + '_>;
    fn len(&self) -> usize;
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum SpecialPosition {
    Agg,
    Window,
    Match,
}

impl Display for SpecialPosition {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            SpecialPosition::Agg => write!(f, "agg function"),
            SpecialPosition::Window => write!(f, "window function"),
            SpecialPosition::Match => write!(f, "match function"),
        }
    }
}

use crate::sql::expression::OrderByExpression;
use crate::sql::query::window::WindowReference;

#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct FunctionTranslationContext {
    pub window: Option<WindowReference>,
    pub specials_allowed: Vec<SpecialPosition>,
    pub order_by: Vec<OrderByExpression>,
}

impl FunctionTranslationContext {
    pub fn with_window(mut self, window: WindowReference) -> Self {
        self.window = Some(window);
        self
    }

    pub fn with_special_allowed(mut self, special: SpecialPosition) -> Self {
        self.specials_allowed.push(special);
        self
    }

    pub fn with_order_by(mut self, mut order_by: Vec<OrderByExpression>) -> Self {
        self.order_by.append(&mut order_by);
        self
    }

    pub fn add_order_by(&mut self, mut order_by: Vec<OrderByExpression>) {
        self.order_by.append(&mut order_by);
    }
}

/// FunctionDef trait - the core abstraction for function definitions.
///
/// Each function is a separate type implementing this trait. Translation and evaluation
/// are handled by separate registries (TranslationRegistry, EvalRegistry) that look up
/// implementations by TypeId.
pub trait FunctionDef: Send + Sync + 'static {
    fn name(&self) -> &'static str;

    /// Parameter definitions with type matchers
    fn parameters(&self) -> Parameters {
        Parameters::new()
    }

    /// Custom binding logic for functions that need to control typing order.
    ///
    /// This is used by functions like `transform` that need to infer lambda parameter
    /// types from other arguments. The function controls the entire binding process.
    ///
    /// Return values:
    /// - `None`: Use standard binding path (default)
    /// - `Some(Ok(binding))`: Custom bind succeeded, use this binding
    /// - `Some(Err(e))`: Custom bind failed, add to error list
    ///
    /// Contract for implementers:
    /// 1. Type and validate non-lambda args first
    /// 2. If non-lambda args don't match, return `Some(Err(...))`
    /// 3. Only type lambdas after non-lambda validation passes
    /// 4. Once you type a lambda, you own all errors for this call
    fn custom_bind(
        &self,
        _ast_binding: &ParameterBinding<Arc<Expression>>,
        _ctx: &mut ExpressionTranslationContext,
    ) -> anyhow::Result<Option<ParameterBinding<Arc<TypedExpression>>>> {
        Ok(None)
    }

    /// Refine already-typed bindings before computing return type.
    ///
    /// This hook is called after standard binding has typed all arguments and
    /// matchers have passed. Use this to validate literal values or apply
    /// coercions (e.g., merging array element types) without re-typing
    /// arguments.
    ///
    /// The default implementation returns the binding unchanged.
    ///
    /// Unlike `custom_bind`, this receives already-typed expressions and should
    /// only wrap them with casts or other transformations, never re-type them.
    fn refine_binding(
        &self,
        binding: ParameterBinding<Arc<TypedExpression>>,
    ) -> anyhow::Result<ParameterBinding<Arc<TypedExpression>>> {
        Ok(binding)
    }

    /// Compute return type given bound parameters (may fail for type mismatches)
    fn return_type(&self, bindings: &dyn ParameterBindingProvider) -> anyhow::Result<Type>;

    /// TypeId for registry lookups
    fn type_id(&self) -> TypeId {
        TypeId::of::<Self>()
    }

    /// Special position semantics (agg, window, match function)
    fn special_position(&self) -> Option<SpecialPosition> {
        None
    }

    /// Whether this function is deterministic (same inputs always produce same output)
    fn is_deterministic(&self) -> bool {
        true
    }

    /// Whether this function accepts sorted input (for aggregates)
    fn sortable_input(&self) -> bool {
        false
    }

    /// When true, SQL translation attaches any `OVER` clause itself.
    ///
    /// Use this for aggregates whose SQL form is a reshape around an inner
    /// aggregate call. The registry must not wrap the outer expression.
    fn manages_window_clause(&self) -> bool {
        false
    }
}

impl Debug for dyn FunctionDef {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("FunctionDef")
            .field("name", &self.name())
            .field("parameters", &self.parameters())
            .finish()
    }
}

impl Display for dyn FunctionDef {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}({})", self.name(), self.parameters())
    }
}

/// Direct function application - arguments match the function signature directly.
pub struct DirectResolution<T: HasType> {
    pub function_def: Arc<dyn FunctionDef>,
    pub binding: ParameterBinding<T>,
    pub typ: Type,
}

/// Broadcast application - one argument is an array, apply the function element-wise.
/// During lowering, this becomes `transform(array, x -> function(x, other_args...))`.
pub struct BroadcastResolution<T: HasType> {
    pub function_def: Arc<dyn FunctionDef>,
    pub binding: ParameterBinding<T>,
    pub typ: Type,
    /// Which positional argument position contains the array being broadcast over.
    pub broadcast_position: usize,
}

/// The result of successfully binding and type-checking a function call.
#[derive(From)]
pub enum FunctionResolution<T: HasType> {
    Direct(DirectResolution<T>),
    Broadcast(BroadcastResolution<T>),
}

#[derive(Debug, Default)]
pub struct Parameters {
    parameters: Vec<Parameter>,
    var_args: Option<Box<dyn Matcher + Send + Sync>>,
}

impl Parameters {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn with<M: Matcher + Send + Sync + 'static>(mut self, name: &str, typ: M) -> Self {
        self.parameters.push(Parameter::of(name, typ));
        self
    }

    pub fn with_default<M: Matcher + Send + Sync + 'static>(
        mut self,
        name: &str,
        typ: M,
        default: Expression,
    ) -> Self {
        self.parameters
            .push(Parameter::with_default(name, typ, default));
        self
    }

    pub fn with_var_args<M: Matcher + Send + Sync + 'static>(mut self, typ: M) -> Self {
        self.var_args = Some(Box::new(typ));
        self
    }

    pub fn add(mut self, parameter: Parameter) -> Self {
        self.parameters.push(parameter);
        self
    }

    pub fn add_var_args<M: Matcher + Send + Sync + 'static>(mut self, typ: M) -> Self {
        self.var_args = Some(Box::new(typ));
        self
    }

    /// Bind arguments to parameters by position and name, then type check.
    ///
    /// This is equivalent to calling `bind` followed by `check`.
    pub fn bind_and_check<T>(
        &self,
        positional: Vec<T>,
        named: impl IntoIterator<Item = (String, T)>,
    ) -> Result<ParameterBinding<T>, Vec<FunctionParameterBindingFailure>>
    where
        T: HasType + Clone + TryFrom<Expression>,
        <T as TryFrom<Expression>>::Error: std::fmt::Debug,
    {
        let binding = self.bind(positional, named)?;
        self.check(&binding)?;
        Ok(binding)
    }

    /// Bind arguments to parameters by position and name, without type checking.
    ///
    /// This handles:
    /// - Matching positional arguments to parameters in order
    /// - Matching named arguments to parameters by name
    /// - Applying default values for missing optional parameters
    /// - Collecting extra positional arguments into varargs
    ///
    /// Does NOT check that argument types match parameter type matchers.
    /// Call `check` afterwards if you need type validation.
    pub fn bind<T>(
        &self,
        positional: Vec<T>,
        named: impl IntoIterator<Item = (String, T)>,
    ) -> Result<ParameterBinding<T>, Vec<FunctionParameterBindingFailure>>
    where
        T: Clone + TryFrom<Expression>,
        <T as TryFrom<Expression>>::Error: std::fmt::Debug,
    {
        let mut arguments = VecMap::new();
        let mut variable_arguments = VecDeque::new();
        let mut errors = Vec::new();

        // Initialize the returned structure with default values for all parameters
        // Use None if no default has been configured.
        // (We will check this later and fail if any are still None)
        for param in &self.parameters {
            let default_value = param.default_value.as_ref().map(|v| {
                v.clone()
                    .try_into()
                    .expect("default value conversion should never fail")
            });
            arguments.insert(param.name.clone(), default_value);
        }

        // Update the returned arguments with the expressions of those that were passed by name.
        // Fail to bind if any named arguments were passed that were not defined as parameters.
        for (name, arg) in named.into_iter() {
            if !arguments.contains_key(name.as_str()) {
                errors.push(FunctionParameterBindingFailure::UnknownNamedParameter(name));
            } else {
                arguments.insert(name, Some(arg));
            }
        }

        // Update the returned arguments with the expressions of those that were passed
        // positionally.
        let mut param_iter = self.parameters.iter();
        let num_positional = positional.len();
        let mut positional_args_iter = positional.into_iter();
        while let Some(param) = param_iter.next() {
            if let Some(arg) = positional_args_iter.next() {
                arguments.insert(param.name.clone(), Some(arg));
            }
        }

        // If they passed more positional arguments than we have parameters:
        //   - use those for varargs, if possible,
        //   - otherwise fail to bind.
        while let Some(arg) = positional_args_iter.next() {
            if let Some(_) = &self.var_args {
                variable_arguments.push_back(arg);
            } else {
                errors.push(
                    FunctionParameterBindingFailure::TooManyPositionalArguments {
                        expected: self.parameters.len(),
                        got: num_positional,
                    },
                );
                break;
            }
        }

        // If there are any parameters not bound at this point, fail. They forgot to specify
        // something required. Remove the option in the process.
        let mut final_arguments = VecMap::new();
        for (name, arg) in arguments.into_iter() {
            match (name, arg) {
                (name, None) => {
                    errors.push(FunctionParameterBindingFailure::MissingRequiredArgument(
                        name,
                    ));
                }
                (name, Some(arg)) => {
                    final_arguments.insert(name, arg);
                }
            }
        }

        if !errors.is_empty() {
            return Err(errors);
        }

        Ok(ParameterBinding {
            arguments: final_arguments,
            var_args: variable_arguments,
        })
    }

    /// Check that bound argument types match parameter type matchers.
    ///
    /// Call this after `bind` to validate types.
    pub fn check<T: HasType>(
        &self,
        binding: &ParameterBinding<T>,
    ) -> Result<(), Vec<FunctionParameterBindingFailure>> {
        let mut errors = Vec::new();

        // Type check named parameters
        for ((name, arg), param) in binding.arguments.iter().zip(self.parameters.iter()) {
            if !param.typ.matches(arg.typ()) {
                errors.push(FunctionParameterBindingFailure::TypeCheckFailure {
                    name: name.clone(),
                    expected: param.typ.to_string(),
                    got: arg.typ().clone(),
                });
            }
        }

        // Type check varargs
        if let Some(var_args_matcher) = &self.var_args {
            for arg in &binding.var_args {
                if !var_args_matcher.matches(arg.typ()) {
                    errors.push(
                        FunctionParameterBindingFailure::TypeCheckFailureForVarargs {
                            expected: var_args_matcher.to_string(),
                            got: arg.typ().clone(),
                        },
                    );
                }
            }
        }

        if !errors.is_empty() {
            return Err(errors);
        }

        Ok(())
    }

    pub fn autocomplete_snippet(&self) -> String {
        self.parameters
            .iter()
            .enumerate()
            .map(|(i, p)| format!("${{{}:{}}}", i + 1, p.name.clone()))
            .join(", ")
    }
}

impl Display for Parameters {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        let mut first = true;
        for param in &self.parameters {
            if first {
                first = false;
            } else {
                write!(f, ", ")?;
            }
            write!(f, "{}", param)?;
        }
        if let Some(var_args) = &self.var_args {
            if !first {
                write!(f, ", ")?;
            }
            write!(f, " ...: {}", var_args)?;
        }
        Ok(())
    }
}

pub struct Parameter {
    name: String,
    typ: Box<dyn Matcher + Send + Sync>,
    /// Default value as an AST expression (must be a literal).
    default_value: Option<Expression>,
}

impl Parameter {
    pub fn of<M: Matcher + Send + Sync + 'static>(name: &str, typ: M) -> Self {
        debug_assert_eq!(name, name.to_lowercase());
        Self {
            name: name.to_string(),
            typ: Box::new(typ),
            default_value: None,
        }
    }

    /// Create a parameter with a default value expression.
    /// The expression should be a literal (e.g., `BooleanLiteral { value: true }.into()`).
    pub fn with_default<M: Matcher + Send + Sync + 'static>(
        name: &str,
        typ: M,
        default: Expression,
    ) -> Self {
        debug_assert_eq!(name, name.to_lowercase());
        Self {
            name: name.to_string(),
            typ: Box::new(typ),
            default_value: Some(default),
        }
    }

    /// Returns the default value expression if one is set.
    pub fn default_value(&self) -> Option<&Expression> {
        self.default_value.as_ref()
    }
}

impl Display for Parameter {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}={}", self.name, self.typ)
    }
}

impl Debug for Parameter {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Parameter")
            .field("name", &self.name)
            .field("typ", &self.typ)
            .field("default_value", &self.default_value)
            .finish()
    }
}

#[derive(Debug, Clone)]
pub struct ParameterBinding<T> {
    arguments: VecMap<String, T>,
    var_args: VecDeque<T>,
}

impl<T> ParameterBinding<T> {
    pub fn get_by_name(&self, name: &str) -> anyhow::Result<&T> {
        self.arguments.get(name).ok_or_else(|| {
            anyhow!(NoBindingError {
                name: name.to_string(),
            })
        })
    }

    pub fn get_by_index(&self, index: usize) -> anyhow::Result<&T> {
        if index < self.arguments.len() {
            self.arguments.values().nth(index).ok_or_else(|| {
                anyhow!(NoBindingError {
                    name: format!("index {}", index),
                })
            })
        } else {
            let var_index = index - self.arguments.len();
            self.var_args.get(var_index).ok_or_else(|| {
                anyhow!(NoBindingError {
                    name: format!("index {}", index),
                })
            })
        }
    }

    pub fn take_by_name(&mut self, name: &str) -> anyhow::Result<T> {
        self.arguments.swap_remove(name).ok_or_else(|| {
            anyhow!(NoBindingError {
                name: name.to_string(),
            })
        })
    }

    pub fn take(&mut self) -> anyhow::Result<T> {
        if !self.arguments.is_empty() {
            Ok(self.arguments.remove_index(0).1)
        } else {
            self.var_args
                .pop_front()
                .ok_or_else(|| anyhow!("No more arguments to take"))
        }
    }

    pub fn iter(&self) -> impl Iterator<Item = &T> {
        self.arguments.values().chain(self.var_args.iter())
    }

    pub fn into_iter(self) -> impl Iterator<Item = T> {
        self.arguments
            .into_iter()
            .map(|(_, v)| v)
            .chain(self.var_args.into_iter())
    }

    pub fn len(&self) -> usize {
        self.arguments.len() + self.var_args.len()
    }

    pub fn is_empty(&self) -> bool {
        self.arguments.is_empty() && self.var_args.is_empty()
    }

    /// Map a function over all parameters, returning a new ParameterBinding.
    pub fn map<U, F: FnMut(T) -> U>(self, mut f: F) -> ParameterBinding<U> {
        ParameterBinding {
            arguments: self.arguments.into_iter().map(|(k, v)| (k, f(v))).collect(),
            var_args: self.var_args.into_iter().map(&mut f).collect(),
        }
    }

    /// Try to map a fallible function over all parameters.
    pub fn try_map<U, E, F: FnMut(T) -> Result<U, E>>(
        self,
        mut f: F,
    ) -> Result<ParameterBinding<U>, E> {
        let mut arguments = VecMap::new();
        for (k, v) in self.arguments {
            arguments.insert(k, f(v)?);
        }
        let mut var_args = VecDeque::new();
        for v in self.var_args {
            var_args.push_back(f(v)?);
        }
        Ok(ParameterBinding {
            arguments,
            var_args,
        })
    }

    /// Replace a bound argument by parameter name.
    /// Returns an error if the parameter doesn't exist.
    pub fn replace_by_name(&mut self, name: &str, value: T) -> anyhow::Result<()> {
        if self.arguments.contains_key(name) {
            self.arguments.insert(name.to_string(), value);
            Ok(())
        } else {
            Err(anyhow!(NoBindingError {
                name: name.to_string(),
            }))
        }
    }

    /// Replace a bound argument by positional index.
    /// Returns a new ParameterBinding with the value replaced.
    pub fn replace_by_index(mut self, index: usize, value: T) -> anyhow::Result<Self> {
        if index < self.arguments.len() {
            // Get the key at this index
            let key = self
                .arguments
                .keys()
                .nth(index)
                .ok_or_else(|| {
                    anyhow!(NoBindingError {
                        name: format!("index {}", index),
                    })
                })?
                .clone();
            self.arguments.insert(key, value);
            Ok(self)
        } else {
            let var_index = index - self.arguments.len();
            if var_index < self.var_args.len() {
                self.var_args[var_index] = value;
                Ok(self)
            } else {
                Err(anyhow!(NoBindingError {
                    name: format!("index {}", index),
                }))
            }
        }
    }

    /// Create an empty binding.
    pub fn empty() -> Self {
        ParameterBinding {
            arguments: VecMap::new(),
            var_args: VecDeque::new(),
        }
    }

    /// Create a binding from an iterator of named pairs.
    pub fn from_named(pairs: impl IntoIterator<Item = (String, T)>) -> Self {
        ParameterBinding {
            arguments: pairs.into_iter().collect(),
            var_args: VecDeque::new(),
        }
    }

    /// Insert a named argument into the binding.
    pub fn insert(&mut self, name: String, value: T) {
        self.arguments.insert(name, value);
    }
}

impl<T> ParameterBinding<Option<T>> {
    /// Get the single defined parameter for reverse evaluation.
    /// Returns the one parameter that is Some, or an error if there are zero or two+ defined values.
    pub fn take_first_defined(self) -> anyhow::Result<T> {
        let mut found: Option<T> = None;

        for opt in self.into_iter() {
            if let Some(val) = opt {
                if found.is_some() {
                    bail!("Expected exactly one variable, found two or more (multiple constants)");
                }
                found = Some(val);
            }
        }

        found.ok_or_else(|| anyhow!("Expected exactly one variable, found zero (no constants)"))
    }
}

impl HasType for ExpressionTranslation {
    fn typ(&self) -> &Type {
        &self.typ
    }
}

impl HasType for Arc<TypedExpression> {
    fn typ(&self) -> &Type {
        &self.resolved_type
    }
}

impl<T: HasType> ParameterBindingProvider for ParameterBinding<T> {
    fn get_by_name(&self, name: &str) -> anyhow::Result<&dyn HasType> {
        self.get_by_name(name).map(|v| v as &dyn HasType)
    }

    fn get_by_index(&self, index: usize) -> anyhow::Result<&dyn HasType> {
        self.get_by_index(index).map(|v| v as &dyn HasType)
    }

    fn len(&self) -> usize {
        self.len()
    }

    fn iter(&self) -> Box<dyn Iterator<Item = &dyn HasType> + '_> {
        Box::new(self.iter().map(|v| v as &dyn HasType))
    }
}