//! Custom UDFs (User Defined Functions) for AiFastDb SQL queries.
//!
//! Provides vector-aware and AI-specific functions that can be used in SQL:
//!
//! - `vector_distance(vec1, vec2, metric)` — compute distance between two vectors
//! - `vector_norm(vec)` — compute L2 norm of a vector
//! - `cosine_similarity(vec1, vec2)` — compute cosine similarity

use std::sync::Arc;

use arrow_array::{Array, Float32Array, FixedSizeListArray};
use arrow_schema::DataType;
use datafusion::error::DataFusionError;
use datafusion::execution::context::SessionContext;
use datafusion::logical_expr::{
    ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility,
};

use super::context::SqlError;

/// Register all custom UDFs into the given SessionContext.
pub fn register_all_udfs(ctx: &SessionContext) -> Result<(), SqlError> {
    ctx.register_udf(ScalarUDF::from(VectorDistanceUdf::new()));
    ctx.register_udf(ScalarUDF::from(VectorNormUdf::new()));
    ctx.register_udf(ScalarUDF::from(CosineSimilarityUdf::new()));
    ctx.register_udf(ScalarUDF::from(DotProductUdf::new()));
    ctx.register_udf(ScalarUDF::from(NormalizeVectorUdf::new()));
    Ok(())
}

// ============================================================================
// vector_distance(vec1, vec2, metric)
// ============================================================================

/// UDF: `vector_distance(vec1, vec2, metric)`
///
/// Computes the distance between two vectors using the specified metric.
///
/// # Arguments (in SQL)
/// - `vec1`: FixedSizeList<Float32> — first vector
/// - `vec2`: FixedSizeList<Float32> — second vector
/// - `metric`: Utf8 — one of 'cosine', 'l2', 'dot'
///
/// # Returns
/// Float32 — the computed distance
///
/// # Example SQL
/// ```sql
/// SELECT id, vector_distance(embedding, ARRAY[0.1, 0.2, 0.3], 'cosine') AS dist
/// FROM memories
/// ORDER BY dist ASC
/// LIMIT 10;
/// ```
#[derive(Debug, PartialEq, Eq, Hash)]
pub struct VectorDistanceUdf {
    signature: Signature,
}

impl VectorDistanceUdf {
    pub fn new() -> Self {
        Self {
            signature: Signature::any(3, Volatility::Immutable),
        }
    }
}

impl ScalarUDFImpl for VectorDistanceUdf {
    fn as_any(&self) -> &dyn std::any::Any {
        self
    }

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

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

    fn return_type(&self, _args: &[DataType]) -> Result<DataType, DataFusionError> {
        Ok(DataType::Float32)
    }

    fn invoke_with_args(
        &self,
        args: ScalarFunctionArgs,
    ) -> Result<ColumnarValue, DataFusionError> {
        let args = &args.args;
        if args.len() < 2 {
            return Err(DataFusionError::Execution(
                "vector_distance requires at least 2 arguments (vec1, vec2[, metric])".to_string(),
            ));
        }

        // Extract metric (default: "l2")
        let metric = if args.len() >= 3 {
            match &args[2] {
                ColumnarValue::Scalar(scalar) => scalar.to_string(),
                _ => "l2".to_string(),
            }
        } else {
            "l2".to_string()
        };

        // Process based on whether inputs are scalar or array
        match (&args[0], &args[1]) {
            (ColumnarValue::Array(arr1), ColumnarValue::Array(arr2)) => {
                let results = compute_pairwise_distance(arr1, arr2, &metric)?;
                Ok(ColumnarValue::Array(Arc::new(results)))
            }
            (ColumnarValue::Array(arr1), ColumnarValue::Scalar(scalar2)) => {
                let vec2 = scalar_to_f32_vec(scalar2)?;
                let results = compute_distance_to_constant(arr1, &vec2, &metric)?;
                Ok(ColumnarValue::Array(Arc::new(results)))
            }
            (ColumnarValue::Scalar(scalar1), ColumnarValue::Array(arr2)) => {
                let vec1 = scalar_to_f32_vec(scalar1)?;
                let results = compute_distance_to_constant(arr2, &vec1, &metric)?;
                Ok(ColumnarValue::Array(Arc::new(results)))
            }
            (ColumnarValue::Scalar(s1), ColumnarValue::Scalar(s2)) => {
                let v1 = scalar_to_f32_vec(s1)?;
                let v2 = scalar_to_f32_vec(s2)?;
                let dist = compute_single_distance(&v1, &v2, &metric)?;
                Ok(ColumnarValue::Scalar(
                    datafusion::common::ScalarValue::Float32(Some(dist)),
                ))
            }
        }
    }
}

// ============================================================================
// vector_norm(vec)
// ============================================================================

/// UDF: `vector_norm(vec)` — compute the L2 norm of a vector.
#[derive(Debug, PartialEq, Eq, Hash)]
pub struct VectorNormUdf {
    signature: Signature,
}

impl VectorNormUdf {
    pub fn new() -> Self {
        Self {
            signature: Signature::any(1, Volatility::Immutable),
        }
    }
}

impl ScalarUDFImpl for VectorNormUdf {
    fn as_any(&self) -> &dyn std::any::Any {
        self
    }

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

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

    fn return_type(&self, _args: &[DataType]) -> Result<DataType, DataFusionError> {
        Ok(DataType::Float32)
    }

    fn invoke_with_args(
        &self,
        args: ScalarFunctionArgs,
    ) -> Result<ColumnarValue, DataFusionError> {
        let args = &args.args;
        if args.is_empty() {
            return Err(DataFusionError::Execution(
                "vector_norm requires 1 argument".to_string(),
            ));
        }

        match &args[0] {
            ColumnarValue::Array(arr) => {
                let list_arr = arr
                    .as_any()
                    .downcast_ref::<FixedSizeListArray>()
                    .ok_or_else(|| {
                        DataFusionError::Execution(
                            "vector_norm expects FixedSizeList<Float32>".to_string(),
                        )
                    })?;

                let mut norms = Vec::with_capacity(list_arr.len());
                for i in 0..list_arr.len() {
                    if list_arr.is_null(i) {
                        norms.push(f32::NAN);
                    } else {
                        let vec = extract_f32_from_list(list_arr, i)?;
                        let norm: f32 = vec.iter().map(|x| x * x).sum::<f32>().sqrt();
                        norms.push(norm);
                    }
                }
                Ok(ColumnarValue::Array(Arc::new(Float32Array::from(norms))))
            }
            ColumnarValue::Scalar(s) => {
                let v = scalar_to_f32_vec(s)?;
                let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
                Ok(ColumnarValue::Scalar(
                    datafusion::common::ScalarValue::Float32(Some(norm)),
                ))
            }
        }
    }
}

// ============================================================================
// cosine_similarity(vec1, vec2)
// ============================================================================

/// UDF: `cosine_similarity(vec1, vec2)` — returns a similarity score in [-1, 1].
#[derive(Debug, PartialEq, Eq, Hash)]
pub struct CosineSimilarityUdf {
    signature: Signature,
}

impl CosineSimilarityUdf {
    pub fn new() -> Self {
        Self {
            signature: Signature::any(2, Volatility::Immutable),
        }
    }
}

impl ScalarUDFImpl for CosineSimilarityUdf {
    fn as_any(&self) -> &dyn std::any::Any {
        self
    }

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

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

    fn return_type(&self, _args: &[DataType]) -> Result<DataType, DataFusionError> {
        Ok(DataType::Float32)
    }

    fn invoke_with_args(
        &self,
        args: ScalarFunctionArgs,
    ) -> Result<ColumnarValue, DataFusionError> {
        let args = &args.args;
        if args.len() < 2 {
            return Err(DataFusionError::Execution(
                "cosine_similarity requires 2 arguments".to_string(),
            ));
        }

        match (&args[0], &args[1]) {
            (ColumnarValue::Array(arr1), ColumnarValue::Array(arr2)) => {
                let results = compute_pairwise_distance(arr1, arr2, "cosine_sim")?;
                Ok(ColumnarValue::Array(Arc::new(results)))
            }
            (ColumnarValue::Array(arr1), ColumnarValue::Scalar(scalar2)) => {
                let vec2 = scalar_to_f32_vec(scalar2)?;
                let results = compute_distance_to_constant(arr1, &vec2, "cosine_sim")?;
                Ok(ColumnarValue::Array(Arc::new(results)))
            }
            (ColumnarValue::Scalar(s1), ColumnarValue::Scalar(s2)) => {
                let v1 = scalar_to_f32_vec(s1)?;
                let v2 = scalar_to_f32_vec(s2)?;
                let sim = cosine_sim(&v1, &v2);
                Ok(ColumnarValue::Scalar(
                    datafusion::common::ScalarValue::Float32(Some(sim)),
                ))
            }
            _ => Err(DataFusionError::Execution(
                "cosine_similarity: unsupported argument combination".to_string(),
            )),
        }
    }
}

// ============================================================================
// Helper functions
// ============================================================================

/// Extract a `Vec<f32>` from a FixedSizeListArray at the given row index.
fn extract_f32_from_list(
    list_arr: &FixedSizeListArray,
    row: usize,
) -> Result<Vec<f32>, DataFusionError> {
    let value = list_arr.value(row);
    let float_arr = value
        .as_any()
        .downcast_ref::<Float32Array>()
        .ok_or_else(|| {
            DataFusionError::Execution("Expected Float32Array inside FixedSizeList".to_string())
        })?;
    Ok((0..float_arr.len()).map(|i| float_arr.value(i)).collect())
}

/// Convert a DataFusion ScalarValue to a `Vec<f32>`.
fn scalar_to_f32_vec(
    scalar: &datafusion::common::ScalarValue,
) -> Result<Vec<f32>, DataFusionError> {
    match scalar {
        datafusion::common::ScalarValue::List(arr) => {
            let float_arr = arr
                .values()
                .as_any()
                .downcast_ref::<Float32Array>()
                .ok_or_else(|| {
                    DataFusionError::Execution(
                        "Expected List<Float32> scalar for vector".to_string(),
                    )
                })?;
            Ok((0..float_arr.len()).map(|i| float_arr.value(i)).collect())
        }
        datafusion::common::ScalarValue::FixedSizeList(arr) => {
            let float_arr = arr
                .values()
                .as_any()
                .downcast_ref::<Float32Array>()
                .ok_or_else(|| {
                    DataFusionError::Execution(
                        "Expected FixedSizeList<Float32> scalar for vector".to_string(),
                    )
                })?;
            Ok((0..float_arr.len()).map(|i| float_arr.value(i)).collect())
        }
        _ => Err(DataFusionError::Execution(format!(
            "Cannot convert scalar {:?} to f32 vector",
            scalar.data_type()
        ))),
    }
}

/// Compute pairwise distances between two arrays of vectors.
fn compute_pairwise_distance(
    arr1: &dyn Array,
    arr2: &dyn Array,
    metric: &str,
) -> Result<Float32Array, DataFusionError> {
    let list1 = arr1
        .as_any()
        .downcast_ref::<FixedSizeListArray>()
        .ok_or_else(|| {
            DataFusionError::Execution("vector_distance arg1 must be FixedSizeList".to_string())
        })?;
    let list2 = arr2
        .as_any()
        .downcast_ref::<FixedSizeListArray>()
        .ok_or_else(|| {
            DataFusionError::Execution("vector_distance arg2 must be FixedSizeList".to_string())
        })?;

    if list1.len() != list2.len() {
        return Err(DataFusionError::Execution(
            "vector_distance: input arrays must have the same length".to_string(),
        ));
    }

    let mut results = Vec::with_capacity(list1.len());
    for i in 0..list1.len() {
        if list1.is_null(i) || list2.is_null(i) {
            results.push(f32::NAN);
        } else {
            let v1 = extract_f32_from_list(list1, i)?;
            let v2 = extract_f32_from_list(list2, i)?;
            let dist = compute_single_distance(&v1, &v2, metric)?;
            results.push(dist);
        }
    }

    Ok(Float32Array::from(results))
}

/// Compute distance from each vector in an array to a constant vector.
fn compute_distance_to_constant(
    arr: &dyn Array,
    constant: &[f32],
    metric: &str,
) -> Result<Float32Array, DataFusionError> {
    let list = arr
        .as_any()
        .downcast_ref::<FixedSizeListArray>()
        .ok_or_else(|| {
            DataFusionError::Execution("Expected FixedSizeList for vector column".to_string())
        })?;

    let mut results = Vec::with_capacity(list.len());
    for i in 0..list.len() {
        if list.is_null(i) {
            results.push(f32::NAN);
        } else {
            let v = extract_f32_from_list(list, i)?;
            let dist = compute_single_distance(&v, constant, metric)?;
            results.push(dist);
        }
    }

    Ok(Float32Array::from(results))
}

/// Compute distance between two f32 vectors using the specified metric.
fn compute_single_distance(v1: &[f32], v2: &[f32], metric: &str) -> Result<f32, DataFusionError> {
    if v1.len() != v2.len() {
        return Err(DataFusionError::Execution(format!(
            "Vector dimension mismatch: {} vs {}",
            v1.len(),
            v2.len()
        )));
    }

    match metric
        .to_lowercase()
        .trim_matches('\'')
        .trim_matches('"')
    {
        "l2" | "euclidean" => {
            let dist: f32 = v1
                .iter()
                .zip(v2.iter())
                .map(|(a, b)| (a - b) * (a - b))
                .sum::<f32>()
                .sqrt();
            Ok(dist)
        }
        "cosine" => {
            let sim = cosine_sim(v1, v2);
            Ok(1.0 - sim) // cosine distance = 1 - similarity
        }
        "cosine_sim" => Ok(cosine_sim(v1, v2)),
        "dot" | "inner_product" => {
            let dot: f32 = v1.iter().zip(v2.iter()).map(|(a, b)| a * b).sum();
            Ok(-dot) // negative dot product (lower = more similar)
        }
        "raw_dot" => {
            let dot: f32 = v1.iter().zip(v2.iter()).map(|(a, b)| a * b).sum();
            Ok(dot) // raw dot product without negation
        }
        other => Err(DataFusionError::Execution(format!(
            "Unknown distance metric: '{}'. Supported: l2, cosine, dot",
            other
        ))),
    }
}

// ============================================================================
// dot_product(vec1, vec2)
// ============================================================================

/// UDF: `dot_product(vec1, vec2)` — inner product of two vectors.
///
/// # Example SQL
/// ```sql
/// SELECT id, dot_product(embedding, ARRAY[0.1, 0.2, 0.3]) AS score
/// FROM memories
/// ORDER BY score DESC
/// LIMIT 10;
/// ```
#[derive(Debug, PartialEq, Eq, Hash)]
pub struct DotProductUdf {
    signature: Signature,
}

impl DotProductUdf {
    pub fn new() -> Self {
        Self {
            signature: Signature::any(2, Volatility::Immutable),
        }
    }
}

impl ScalarUDFImpl for DotProductUdf {
    fn as_any(&self) -> &dyn std::any::Any {
        self
    }

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

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

    fn return_type(&self, _args: &[DataType]) -> Result<DataType, DataFusionError> {
        Ok(DataType::Float32)
    }

    fn invoke_with_args(
        &self,
        args: ScalarFunctionArgs,
    ) -> Result<ColumnarValue, DataFusionError> {
        let args = &args.args;
        if args.len() < 2 {
            return Err(DataFusionError::Execution(
                "dot_product requires 2 arguments".to_string(),
            ));
        }

        match (&args[0], &args[1]) {
            (ColumnarValue::Array(arr1), ColumnarValue::Array(arr2)) => {
                let results = compute_pairwise_distance(arr1, arr2, "raw_dot")?;
                Ok(ColumnarValue::Array(Arc::new(results)))
            }
            (ColumnarValue::Array(arr1), ColumnarValue::Scalar(scalar2)) => {
                let vec2 = scalar_to_f32_vec(scalar2)?;
                let results = compute_distance_to_constant(arr1, &vec2, "raw_dot")?;
                Ok(ColumnarValue::Array(Arc::new(results)))
            }
            (ColumnarValue::Scalar(s1), ColumnarValue::Scalar(s2)) => {
                let v1 = scalar_to_f32_vec(s1)?;
                let v2 = scalar_to_f32_vec(s2)?;
                let dot: f32 = v1.iter().zip(v2.iter()).map(|(a, b)| a * b).sum();
                Ok(ColumnarValue::Scalar(
                    datafusion::common::ScalarValue::Float32(Some(dot)),
                ))
            }
            _ => Err(DataFusionError::Execution(
                "dot_product: unsupported argument combination".to_string(),
            )),
        }
    }
}

// ============================================================================
// normalize_vector(vec)
// ============================================================================

/// UDF: `normalize_vector(vec)` — returns a unit-length vector (L2 norm = 1).
///
/// Returns the same vector type (FixedSizeList<Float32>) with all elements
/// divided by the L2 norm.
///
/// # Example SQL
/// ```sql
/// SELECT id, normalize_vector(embedding) AS unit_embedding
/// FROM memories;
/// ```
#[derive(Debug, PartialEq, Eq, Hash)]
pub struct NormalizeVectorUdf {
    signature: Signature,
}

impl NormalizeVectorUdf {
    pub fn new() -> Self {
        Self {
            signature: Signature::any(1, Volatility::Immutable),
        }
    }
}

impl ScalarUDFImpl for NormalizeVectorUdf {
    fn as_any(&self) -> &dyn std::any::Any {
        self
    }

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

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

    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType, DataFusionError> {
        // Return the same type as input
        if arg_types.is_empty() {
            Ok(DataType::Null)
        } else {
            Ok(arg_types[0].clone())
        }
    }

    fn invoke_with_args(
        &self,
        args: ScalarFunctionArgs,
    ) -> Result<ColumnarValue, DataFusionError> {
        let args = &args.args;
        if args.is_empty() {
            return Err(DataFusionError::Execution(
                "normalize_vector requires 1 argument".to_string(),
            ));
        }

        match &args[0] {
            ColumnarValue::Array(arr) => {
                let list_arr = arr
                    .as_any()
                    .downcast_ref::<FixedSizeListArray>()
                    .ok_or_else(|| {
                        DataFusionError::Execution(
                            "normalize_vector expects FixedSizeList<Float32>".to_string(),
                        )
                    })?;

                let dim = list_arr.value_length() as usize;

                // Build normalized vectors for all rows
                let mut all_values: Vec<f32> = Vec::with_capacity(list_arr.len() * dim);
                for i in 0..list_arr.len() {
                    if list_arr.is_null(i) {
                        all_values.extend(std::iter::repeat(0.0f32).take(dim));
                    } else {
                        let vec = extract_f32_from_list(list_arr, i)?;
                        let norm: f32 = vec.iter().map(|x| x * x).sum::<f32>().sqrt();
                        if norm > 0.0 {
                            all_values.extend(vec.iter().map(|x| x / norm));
                        } else {
                            all_values.extend(std::iter::repeat(0.0f32).take(dim));
                        }
                    }
                }

                let values_array = Float32Array::from(all_values);
                let field = Arc::new(arrow_schema::Field::new("item", DataType::Float32, true));
                let result = FixedSizeListArray::try_new(field, dim as i32, Arc::new(values_array), None)
                    .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?;

                Ok(ColumnarValue::Array(Arc::new(result)))
            }
            ColumnarValue::Scalar(s) => {
                let v = scalar_to_f32_vec(s)?;
                let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
                let normalized: Vec<f32> = if norm > 0.0 {
                    v.iter().map(|x| x / norm).collect()
                } else {
                    vec![0.0; v.len()]
                };

                let values_array = Float32Array::from(normalized);
                let field = Arc::new(arrow_schema::Field::new("item", DataType::Float32, true));
                let dim = v.len() as i32;
                let result = FixedSizeListArray::try_new(field, dim, Arc::new(values_array), None)
                    .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?;

                Ok(ColumnarValue::Scalar(
                    datafusion::common::ScalarValue::FixedSizeList(Arc::new(result)),
                ))
            }
        }
    }
}

// ============================================================================
// Helper functions
// ============================================================================

/// Compute cosine similarity between two vectors.
fn cosine_sim(v1: &[f32], v2: &[f32]) -> f32 {
    let dot: f32 = v1.iter().zip(v2.iter()).map(|(a, b)| a * b).sum();
    let norm1: f32 = v1.iter().map(|x| x * x).sum::<f32>().sqrt();
    let norm2: f32 = v2.iter().map(|x| x * x).sum::<f32>().sqrt();

    if norm1 == 0.0 || norm2 == 0.0 {
        0.0
    } else {
        dot / (norm1 * norm2)
    }
}
