pqb 0.1.2

// Copyright 2025 FastLabs Developers
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! SQL built-in functions.

use std::borrow::Cow;

use crate::expr::BinaryOp;
use crate::expr::Expr;
use crate::expr::write_expr;
use crate::types::IntoColumnRef;
use crate::types::IntoIden;
use crate::writer::SqlWriter;

#[derive(Debug, Clone, PartialEq, Eq)]
enum Func {
    Max,
    Min,
    Sum,
    Avg,
    Cast,
    Count,
    Coalesce,
    Lower,
    Upper,
    Custom(Cow<'static, str>),
}

/// A function call expression.
#[derive(Debug, Clone, PartialEq)]
pub struct FunctionCall {
    func: Func,
    args: Vec<Expr>,
}

impl FunctionCall {
    /// Create a new MAX function call.
    pub fn max<T>(expr: T) -> Self
    where
        T: Into<Expr>,
    {
        Self {
            func: Func::Max,
            args: vec![expr.into()],
        }
    }

    /// Create a new MIN function call.
    pub fn min<T>(expr: T) -> Self
    where
        T: Into<Expr>,
    {
        Self {
            func: Func::Min,
            args: vec![expr.into()],
        }
    }

    /// Create a new SUM function call.
    pub fn sum<T>(expr: T) -> Self
    where
        T: Into<Expr>,
    {
        Self {
            func: Func::Sum,
            args: vec![expr.into()],
        }
    }

    /// Create a new AVG function call.
    pub fn avg<T>(expr: T) -> Self
    where
        T: Into<Expr>,
    {
        Self {
            func: Func::Avg,
            args: vec![expr.into()],
        }
    }

    /// Create a new COUNT function call.
    pub fn count<T>(expr: T) -> Self
    where
        T: Into<Expr>,
    {
        Self {
            func: Func::Count,
            args: vec![expr.into()],
        }
    }

    /// Create a COUNT(*) function call.
    pub fn count_all() -> Self {
        Self {
            func: Func::Count,
            args: vec![Expr::Asterisk],
        }
    }

    /// Create a COALESCE function call.
    pub fn coalesce<A, B>(a: A, b: B) -> Self
    where
        A: Into<Expr>,
        B: Into<Expr>,
    {
        Self {
            func: Func::Coalesce,
            args: vec![a.into(), b.into()],
        }
    }

    /// Create an LOWER function call.
    pub fn lower<T>(expr: T) -> FunctionCall
    where
        T: Into<Expr>,
    {
        FunctionCall {
            func: Func::Lower,
            args: vec![expr.into()],
        }
    }

    /// Create an UPPER function call.
    pub fn upper<T>(expr: T) -> FunctionCall
    where
        T: Into<Expr>,
    {
        FunctionCall {
            func: Func::Upper,
            args: vec![expr.into()],
        }
    }

    /// Call `CAST` function with a custom type.
    pub fn cast_as<V, I>(expr: V, iden: I) -> FunctionCall
    where
        V: Into<Expr>,
        I: IntoIden,
    {
        let expr = expr.into();
        FunctionCall {
            func: Func::Cast,
            args: vec![expr.binary(BinaryOp::As, Expr::custom(iden.into_iden().into_inner()))],
        }
    }

    /// Create a function call with custom name and arguments.
    pub fn custom<N, T, I>(name: N, args: I) -> Self
    where
        N: Into<Cow<'static, str>>,
        I: IntoIterator<Item = T>,
        T: Into<Expr>,
    {
        Self {
            func: Func::Custom(name.into()),
            args: args.into_iter().map(|arg| arg.into()).collect(),
        }
    }
}

impl From<FunctionCall> for Expr {
    fn from(call: FunctionCall) -> Self {
        Expr::FunctionCall(call)
    }
}

pub(crate) fn write_function_call<W: SqlWriter>(w: &mut W, call: &FunctionCall) {
    match &call.func {
        Func::Max => w.push_str("MAX"),
        Func::Min => w.push_str("MIN"),
        Func::Sum => w.push_str("SUM"),
        Func::Avg => w.push_str("AVG"),
        Func::Cast => w.push_str("CAST"),
        Func::Count => w.push_str("COUNT"),
        Func::Coalesce => w.push_str("COALESCE"),
        Func::Lower => w.push_str("LOWER"),
        Func::Upper => w.push_str("UPPER"),
        Func::Custom(name) => w.push_str(name),
    }
    w.push_char('(');
    for (i, arg) in call.args.iter().enumerate() {
        if i > 0 {
            w.push_str(", ");
        }
        write_expr(w, arg);
    }
    w.push_char(')');
}

/// Express a column reference for use in aggregate functions.
pub fn col<T>(col: T) -> Expr
where
    T: IntoColumnRef,
{
    Expr::column(col)
}