// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you 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.

use arrow::array::{ArrayRef, AsArray, PrimitiveArray};
use arrow::compute::try_binary;
use arrow::datatypes::{DataType, Int64Type};
use arrow::error::ArrowError;
use std::mem::swap;
use std::sync::Arc;

use datafusion_common::{Result, ScalarValue, exec_err, internal_datafusion_err};
use datafusion_expr::{
    ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
    Volatility,
};
use datafusion_macros::user_doc;

#[user_doc(
    doc_section(label = "Math Functions"),
    description = "Returns the greatest common divisor of `expression_x` and `expression_y`. Returns 0 if both inputs are zero.",
    syntax_example = "gcd(expression_x, expression_y)",
    sql_example = r#"```sql
> SELECT gcd(48, 18);
+------------+
| gcd(48,18) |
+------------+
| 6          |
+------------+
```"#,
    standard_argument(name = "expression_x", prefix = "First numeric"),
    standard_argument(name = "expression_y", prefix = "Second numeric")
)]
#[derive(Debug, PartialEq, Eq, Hash)]
pub struct GcdFunc {
    signature: Signature,
}

impl Default for GcdFunc {
    fn default() -> Self {
        Self::new()
    }
}

impl GcdFunc {
    pub fn new() -> Self {
        Self {
            signature: Signature::uniform(
                2,
                vec![DataType::Int64],
                Volatility::Immutable,
            ),
        }
    }
}

impl ScalarUDFImpl for GcdFunc {
    fn name(&self) -> &str {
        "gcd"
    }

    fn signature(&self) -> &Signature {
        &self.signature
    }

    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
        Ok(DataType::Int64)
    }

    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
        let args: [ColumnarValue; 2] = args.args.try_into().map_err(|_| {
            internal_datafusion_err!("Expected 2 arguments for function gcd")
        })?;

        match args {
            [ColumnarValue::Array(a), ColumnarValue::Array(b)] => {
                compute_gcd_for_arrays(&a, &b)
            }
            [
                ColumnarValue::Scalar(ScalarValue::Int64(a)),
                ColumnarValue::Scalar(ScalarValue::Int64(b)),
            ] => match (a, b) {
                (Some(a), Some(b)) => Ok(ColumnarValue::Scalar(ScalarValue::Int64(
                    Some(compute_gcd(a, b)?),
                ))),
                _ => Ok(ColumnarValue::Scalar(ScalarValue::Int64(None))),
            },
            [
                ColumnarValue::Array(a),
                ColumnarValue::Scalar(ScalarValue::Int64(b)),
            ] => compute_gcd_with_scalar(&a, b),
            [
                ColumnarValue::Scalar(ScalarValue::Int64(a)),
                ColumnarValue::Array(b),
            ] => compute_gcd_with_scalar(&b, a),
            _ => exec_err!("Unsupported argument types for function gcd"),
        }
    }

    fn documentation(&self) -> Option<&Documentation> {
        self.doc()
    }
}

fn compute_gcd_for_arrays(a: &ArrayRef, b: &ArrayRef) -> Result<ColumnarValue> {
    let a = a.as_primitive::<Int64Type>();
    let b = b.as_primitive::<Int64Type>();
    try_binary(a, b, compute_gcd)
        .map(|arr: PrimitiveArray<Int64Type>| {
            ColumnarValue::Array(Arc::new(arr) as ArrayRef)
        })
        .map_err(Into::into) // convert ArrowError to DataFusionError
}

fn compute_gcd_with_scalar(arr: &ArrayRef, scalar: Option<i64>) -> Result<ColumnarValue> {
    let prim = arr.as_primitive::<Int64Type>();
    match scalar {
        Some(scalar_value) if scalar_value != 0 && scalar_value != i64::MIN => {
            // The gcd result divides both inputs' absolute values. When the
            // scalar is neither 0 nor i64::MIN, the gcd's absolute value fits
            // in i64, so the cast to i64 below cannot overflow. This allows us
            // to use `unary` instead of `try_unary`, which allows LLVM to
            // vectorize more effectively.
            let sv = scalar_value.unsigned_abs();
            let result: PrimitiveArray<Int64Type> =
                prim.unary(|val| unsigned_gcd(val.unsigned_abs(), sv) as i64);
            Ok(ColumnarValue::Array(Arc::new(result) as ArrayRef))
        }
        Some(scalar_value) => {
            let result: PrimitiveArray<Int64Type> =
                prim.try_unary(|val| compute_gcd(val, scalar_value))?;
            Ok(ColumnarValue::Array(Arc::new(result) as ArrayRef))
        }
        None => Ok(ColumnarValue::Scalar(ScalarValue::Int64(None))),
    }
}

/// Computes gcd of two unsigned integers using Binary GCD algorithm.
pub(super) fn unsigned_gcd(mut a: u64, mut b: u64) -> u64 {
    if a == 0 {
        return b;
    }
    if b == 0 {
        return a;
    }

    let shift = (a | b).trailing_zeros();
    a >>= a.trailing_zeros();
    loop {
        b >>= b.trailing_zeros();
        if a > b {
            swap(&mut a, &mut b);
        }
        b -= a;
        if b == 0 {
            return a << shift;
        }
    }
}

/// Computes greatest common divisor using Binary GCD algorithm.
pub fn compute_gcd(x: i64, y: i64) -> Result<i64, ArrowError> {
    let a = x.unsigned_abs();
    let b = y.unsigned_abs();
    let r = unsigned_gcd(a, b);
    // The result can be up to 2^63 (e.g. gcd(i64::MIN, 0) or
    // gcd(i64::MIN, i64::MIN)), which does not fit into i64.
    r.try_into().map_err(|_| {
        ArrowError::ComputeError(format!("Signed integer overflow in GCD({x}, {y})"))
    })
}
