// 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 std::sync::Arc;

use arrow::array::{Array, BinaryBuilder, OffsetSizeTrait};
use arrow::datatypes::DataType;
use datafusion::common::{cast::as_generic_string_array, exec_err, DataFusionError, ScalarValue};
use datafusion::logical_expr::ColumnarValue;

/// Sentinel stored in `HEX_LUT` for bytes that are not valid hex digits. Its value `0xFF` has all
/// bits set, so OR-ing the two nibbles of a byte and comparing the result against `INVALID_HEX`
/// rejects the byte whenever either nibble is invalid: valid nibbles are `<= 0x0F`, so a valid
/// pair can never OR up to `0xFF`. This lets the decode loop test both nibbles with one comparison.
const INVALID_HEX: u8 = 0xFF;

/// Lookup table mapping each possible input byte to its hex value, or `INVALID_HEX` for bytes
/// that are not `0-9`, `A-F`, or `a-f`. Built at compile time so decoding is a single indexed
/// load per digit rather than a chain of range comparisons.
const HEX_LUT: [u8; 256] = build_hex_lut();

const fn build_hex_lut() -> [u8; 256] {
    let mut lut = [INVALID_HEX; 256];
    let mut i = 0;
    while i < 256 {
        lut[i] = match i as u8 {
            c @ b'0'..=b'9' => c - b'0',
            c @ b'A'..=b'F' => c - b'A' + 10,
            c @ b'a'..=b'f' => c - b'a' + 10,
            _ => INVALID_HEX,
        };
        i += 1;
    }
    lut
}

#[inline]
fn invalid_hex_digit() -> DataFusionError {
    DataFusionError::Execution("Input to unhex is not a valid hex digit".to_string())
}

/// Convert a hex string to binary and store the result in `result`. Returns an error if the input
/// is not a valid hex string.
fn unhex(hex_str: &str, result: &mut Vec<u8>) -> Result<(), DataFusionError> {
    let bytes = hex_str.as_bytes();
    // Each output byte consumes two hex digits (the optional leading nibble rounds up).
    result.reserve(bytes.len().div_ceil(2));

    let mut i = 0;

    if (bytes.len() & 0x01) != 0 {
        let v = HEX_LUT[bytes[0] as usize];
        if v == INVALID_HEX {
            return Err(invalid_hex_digit());
        }
        result.push(v);
        i += 1;
    }

    while i < bytes.len() {
        let first = HEX_LUT[bytes[i] as usize];
        let second = HEX_LUT[bytes[i + 1] as usize];
        if (first | second) == INVALID_HEX {
            return Err(invalid_hex_digit());
        }
        result.push((first << 4) | second);

        i += 2;
    }

    Ok(())
}

fn spark_unhex_inner<T: OffsetSizeTrait>(
    array: &ColumnarValue,
    fail_on_error: bool,
) -> Result<ColumnarValue, DataFusionError> {
    match array {
        ColumnarValue::Array(array) => {
            let string_array = as_generic_string_array::<T>(array)?;

            let mut encoded = Vec::new();
            // Every two input hex digits produce one output byte, so the decoded data is at most
            // half the total input length. Preallocating both the offset and value buffers avoids
            // the repeated doublings a fresh `BinaryBuilder` would incur across the whole column.
            let mut builder = BinaryBuilder::with_capacity(
                string_array.len(),
                string_array.value_data().len() / 2,
            );

            for item in string_array.iter() {
                if let Some(s) = item {
                    if unhex(s, &mut encoded).is_ok() {
                        builder.append_value(encoded.as_slice());
                    } else if fail_on_error {
                        return exec_err!("Input to unhex is not a valid hex string: {s}");
                    } else {
                        builder.append_null();
                    }
                    encoded.clear();
                } else {
                    builder.append_null();
                }
            }
            Ok(ColumnarValue::Array(Arc::new(builder.finish())))
        }
        ColumnarValue::Scalar(ScalarValue::Utf8(Some(string))) => {
            let mut encoded = Vec::new();

            if unhex(string, &mut encoded).is_ok() {
                Ok(ColumnarValue::Scalar(ScalarValue::Binary(Some(encoded))))
            } else if fail_on_error {
                exec_err!("Input to unhex is not a valid hex string: {string}")
            } else {
                Ok(ColumnarValue::Scalar(ScalarValue::Binary(None)))
            }
        }
        ColumnarValue::Scalar(ScalarValue::Utf8(None)) => {
            Ok(ColumnarValue::Scalar(ScalarValue::Binary(None)))
        }
        _ => {
            exec_err!(
                "The first argument must be a string scalar or array, but got: {:?}",
                array
            )
        }
    }
}

/// Spark-compatible `unhex` expression
pub fn spark_unhex(args: &[ColumnarValue]) -> Result<ColumnarValue, DataFusionError> {
    if args.len() > 2 {
        return exec_err!("unhex takes at most 2 arguments, but got: {}", args.len());
    }

    let val_to_unhex = &args[0];
    let fail_on_error = if args.len() == 2 {
        match &args[1] {
            ColumnarValue::Scalar(ScalarValue::Boolean(Some(fail_on_error))) => *fail_on_error,
            _ => {
                return exec_err!(
                    "The second argument must be boolean scalar, but got: {:?}",
                    args[1]
                );
            }
        }
    } else {
        false
    };

    match val_to_unhex.data_type() {
        DataType::Utf8 => spark_unhex_inner::<i32>(val_to_unhex, fail_on_error),
        DataType::LargeUtf8 => spark_unhex_inner::<i64>(val_to_unhex, fail_on_error),
        other => exec_err!(
            "The first argument must be a Utf8 or LargeUtf8: {:?}",
            other
        ),
    }
}

#[cfg(test)]
mod test {
    use std::sync::Arc;

    use arrow::array::make_array;
    use arrow::array::ArrayData;
    use arrow::array::{BinaryBuilder, StringBuilder};
    use datafusion::common::ScalarValue;
    use datafusion::logical_expr::ColumnarValue;

    use super::unhex;

    #[test]
    fn test_spark_unhex_null() -> Result<(), Box<dyn std::error::Error>> {
        let input = ArrayData::new_null(&arrow::datatypes::DataType::Utf8, 2);
        let output = ArrayData::new_null(&arrow::datatypes::DataType::Binary, 2);

        let input = ColumnarValue::Array(Arc::new(make_array(input)));
        let expected = ColumnarValue::Array(Arc::new(make_array(output)));

        let result = super::spark_unhex(&[input])?;

        match (result, expected) {
            (ColumnarValue::Array(result), ColumnarValue::Array(expected)) => {
                assert_eq!(*result, *expected);
                Ok(())
            }
            _ => Err("Unexpected result type".into()),
        }
    }

    #[test]
    fn test_partial_error() -> Result<(), Box<dyn std::error::Error>> {
        let mut input = StringBuilder::new();

        input.append_value("1CGG"); // 1C is ok, but GG is invalid
        input.append_value("537061726B2053514C"); // followed by valid

        let input = ColumnarValue::Array(Arc::new(input.finish()));
        let fail_on_error = ColumnarValue::Scalar(ScalarValue::Boolean(Some(false)));

        let result = super::spark_unhex(&[input, fail_on_error])?;

        let mut expected = BinaryBuilder::new();
        expected.append_null();
        expected.append_value("Spark SQL".as_bytes());

        match (result, ColumnarValue::Array(Arc::new(expected.finish()))) {
            (ColumnarValue::Array(result), ColumnarValue::Array(expected)) => {
                assert_eq!(*result, *expected);

                Ok(())
            }
            _ => Err("Unexpected result type".into()),
        }
    }

    #[test]
    fn test_unhex_valid() -> Result<(), Box<dyn std::error::Error>> {
        let mut result = Vec::new();

        unhex("537061726B2053514C", &mut result)?;
        let result_str = std::str::from_utf8(&result)?;
        assert_eq!(result_str, "Spark SQL");
        result.clear();

        unhex("1C", &mut result)?;
        assert_eq!(result, vec![28]);
        result.clear();

        unhex("737472696E67", &mut result)?;
        assert_eq!(result, "string".as_bytes());
        result.clear();

        unhex("1", &mut result)?;
        assert_eq!(result, vec![1]);
        result.clear();

        Ok(())
    }

    #[test]
    fn test_odd_length() -> Result<(), Box<dyn std::error::Error>> {
        let mut result = Vec::new();

        unhex("A1B", &mut result)?;
        assert_eq!(result, vec![10, 27]);
        result.clear();

        unhex("0A1B", &mut result)?;
        assert_eq!(result, vec![10, 27]);
        result.clear();

        Ok(())
    }

    #[test]
    fn test_unhex_empty() {
        let mut result = Vec::new();

        // Empty hex string
        unhex("", &mut result).unwrap();
        assert!(result.is_empty());
    }

    #[test]
    fn test_unhex_invalid() {
        let mut result = Vec::new();

        // Invalid hex strings
        assert!(unhex("##", &mut result).is_err());
        assert!(unhex("G123", &mut result).is_err());
        assert!(unhex("hello", &mut result).is_err());
        assert!(unhex("\0", &mut result).is_err());
    }
}