//! URL encode/decode UDFs for DataFusion.
//!
//! `application/x-www-form-urlencoded` style encoding (alphanumeric and `.-*_` preserved,
//! space as `+`, other bytes as `%XX` uppercase), matching Hamelin’s `url_encode` / `url_decode`.

use std::any::Any;
use std::fmt::Write;
use std::sync::Arc;

use datafusion::arrow::array::{Array, ArrayRef, AsArray, GenericStringBuilder};
use datafusion::arrow::datatypes::DataType;
use datafusion::common::{exec_err, Result, ScalarValue};
use datafusion::error::DataFusionError;
use datafusion::logical_expr::{
    ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, TypeSignature,
    Volatility,
};
use percent_encoding::{percent_decode_str, percent_encode_byte};

use super::string_utils::{scalar_to_str, STRING_TYPES};

/// Form-style URL encoding (UTF-8 bytes, uppercase hex) into `fmt::Write`.
fn write_form_urlencode<W: Write>(w: &mut W, input: &str) -> std::fmt::Result {
    for &b in input.as_bytes() {
        match b {
            b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'.' | b'-' | b'*' | b'_' => {
                w.write_char(b as char)?;
            }
            // Space is `+` in application/x-www-form-urlencoded (not `%20`); see decode below.
            b' ' => w.write_str("+")?,
            _ => w.write_str(percent_encode_byte(b))?,
        }
    }
    Ok(())
}

fn form_urldecode(s: &str) -> Result<String> {
    // `percent_decode_str` does not treat `+` as space; form-urlencoded uses `+` for ASCII space.
    let with_spaces = s.replace('+', " ");
    percent_decode_str(&with_spaces)
        .decode_utf8()
        .map(|c| c.into_owned())
        .map_err(|e| DataFusionError::Execution(format!("url_decode: {e}")))
}

// ============================================================================
// url_encode: String -> String
// ============================================================================

#[derive(Debug, PartialEq, Eq, Hash)]
pub struct UrlEncodeUdf {
    signature: Signature,
}

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

impl UrlEncodeUdf {
    pub fn new() -> Self {
        let sigs: Vec<TypeSignature> = STRING_TYPES
            .iter()
            .map(|t| TypeSignature::Exact(vec![t.clone()]))
            .collect();
        Self {
            signature: Signature::new(TypeSignature::OneOf(sigs), Volatility::Immutable),
        }
    }
}

impl ScalarUDFImpl for UrlEncodeUdf {
    fn as_any(&self) -> &dyn Any {
        self
    }

    fn name(&self) -> &str {
        "hamelin_url_encode"
    }

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

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

    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
        let args = args.args;
        if args.len() != 1 {
            return exec_err!("url_encode expects 1 argument, got {}", args.len());
        }

        match &args[0] {
            ColumnarValue::Scalar(scalar) => {
                let result = scalar_to_str(scalar)?
                    .map(|s| -> Result<String, DataFusionError> {
                        let mut out = String::with_capacity(s.len().saturating_mul(2));
                        write_form_urlencode(&mut out, s)
                            .map_err(|e| DataFusionError::Execution(format!("url_encode: {e}")))?;
                        Ok(out)
                    })
                    .transpose()?;
                Ok(ColumnarValue::Scalar(ScalarValue::Utf8(result)))
            }
            ColumnarValue::Array(array) => {
                let result = match array.data_type() {
                    DataType::Utf8 => map_string_array_url_encode(array.as_string::<i32>())?,
                    DataType::LargeUtf8 => map_string_array_url_encode(array.as_string::<i64>())?,
                    DataType::Utf8View => map_string_array_url_encode(array.as_string_view())?,
                    other => return exec_err!("url_encode expects string array, got {}", other),
                };
                Ok(ColumnarValue::Array(result))
            }
        }
    }
}

pub fn url_encode_udf() -> ScalarUDF {
    ScalarUDF::new_from_impl(UrlEncodeUdf::new())
}

// ============================================================================
// url_decode: String -> String
// ============================================================================

#[derive(Debug, PartialEq, Eq, Hash)]
pub struct UrlDecodeUdf {
    signature: Signature,
}

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

impl UrlDecodeUdf {
    pub fn new() -> Self {
        let sigs: Vec<TypeSignature> = STRING_TYPES
            .iter()
            .map(|t| TypeSignature::Exact(vec![t.clone()]))
            .collect();
        Self {
            signature: Signature::new(TypeSignature::OneOf(sigs), Volatility::Immutable),
        }
    }
}

impl ScalarUDFImpl for UrlDecodeUdf {
    fn as_any(&self) -> &dyn Any {
        self
    }

    fn name(&self) -> &str {
        "hamelin_url_decode"
    }

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

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

    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
        let args = args.args;
        if args.len() != 1 {
            return exec_err!("url_decode expects 1 argument, got {}", args.len());
        }

        match &args[0] {
            ColumnarValue::Scalar(scalar) => {
                let s = scalar_to_str(scalar)?;
                let result = match s {
                    None => None,
                    Some(v) => Some(form_urldecode(v)?),
                };
                Ok(ColumnarValue::Scalar(ScalarValue::Utf8(result)))
            }
            ColumnarValue::Array(array) => {
                let result = match array.data_type() {
                    DataType::Utf8 => map_string_array_url_decode(array.as_string::<i32>())?,
                    DataType::LargeUtf8 => map_string_array_url_decode(array.as_string::<i64>())?,
                    DataType::Utf8View => map_string_array_url_decode(array.as_string_view())?,
                    other => return exec_err!("url_decode expects string array, got {}", other),
                };
                Ok(ColumnarValue::Array(result))
            }
        }
    }
}

pub fn url_decode_udf() -> ScalarUDF {
    ScalarUDF::new_from_impl(UrlDecodeUdf::new())
}

fn map_string_array_url_encode<T>(array: &T) -> Result<ArrayRef>
where
    T: Array + 'static,
    for<'a> &'a T: IntoIterator<Item = Option<&'a str>>,
{
    let len = array.len();
    let mut builder = GenericStringBuilder::<i32>::with_capacity(len, len.saturating_mul(16));
    for opt in array.into_iter() {
        match opt {
            None => builder.append_null(),
            Some(s) => {
                let mut encoded = String::with_capacity(s.len().saturating_mul(2));
                write_form_urlencode(&mut encoded, s)
                    .map_err(|e| DataFusionError::Execution(format!("url_encode: {e}")))?;
                builder.append_value(&encoded);
            }
        }
    }
    Ok(Arc::new(builder.finish()))
}

fn map_string_array_url_decode<T>(array: &T) -> Result<ArrayRef>
where
    T: Array + 'static,
    for<'a> &'a T: IntoIterator<Item = Option<&'a str>>,
{
    let len = array.len();
    let mut builder = GenericStringBuilder::<i32>::with_capacity(len, len.saturating_mul(16));
    for opt in array.into_iter() {
        match opt {
            None => builder.append_null(),
            Some(s) => builder.append_value(form_urldecode(s)?),
        }
    }
    Ok(Arc::new(builder.finish()))
}
