use std::collections::HashMap;
use std::sync::Arc;

use arrow::array::{
    Array, ArrayRef, Float32Array, Float64Array, Int32Array, Int64Array, TimestampNanosecondArray,
    UInt32Array, UInt64Array, UInt8Array,
};
use bt_kernel::{BtError, Result};

pub type RegisteredFunction = Arc<dyn Fn(&[ArrayRef]) -> Result<ArrayRef> + Send + Sync>;

#[derive(Clone)]
struct FunctionDef {
    min_args: usize,
    max_args: usize,
    implementation: RegisteredFunction,
}

#[derive(Clone, Default)]
pub struct FunctionRegistry {
    functions: HashMap<String, FunctionDef>,
}

impl FunctionRegistry {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn with_builtins() -> Self {
        let mut registry = Self::new();
        registry.register(
            "abs",
            1,
            1,
            Arc::new(|args| apply_unary_f64(args, f64::abs)),
        );
        registry.register(
            "sqrt",
            1,
            1,
            Arc::new(|args| apply_unary_f64(args, f64::sqrt)),
        );
        registry.register("log", 1, 1, Arc::new(|args| apply_unary_f64(args, f64::ln)));
        registry.register(
            "exp",
            1,
            1,
            Arc::new(|args| apply_unary_f64(args, f64::exp)),
        );
        registry.register(
            "min",
            2,
            2,
            Arc::new(|args| apply_binary_f64(args, f64::min)),
        );
        registry.register(
            "max",
            2,
            2,
            Arc::new(|args| apply_binary_f64(args, f64::max)),
        );
        registry
    }

    pub fn register(
        &mut self,
        name: impl Into<String>,
        min_args: usize,
        max_args: usize,
        implementation: RegisteredFunction,
    ) {
        self.functions.insert(
            name.into(),
            FunctionDef {
                min_args,
                max_args,
                implementation,
            },
        );
    }

    pub fn call(&self, name: &str, args: &[ArrayRef]) -> Result<ArrayRef> {
        let function = self
            .functions
            .get(name)
            .ok_or_else(|| BtError::Unsupported(format!("unknown expression function '{name}'")))?;

        if args.len() < function.min_args || args.len() > function.max_args {
            return Err(BtError::Expression(format!(
                "function '{name}' expected {}..={} args, got {}",
                function.min_args,
                function.max_args,
                args.len()
            )));
        }

        (function.implementation)(args)
    }
}

fn apply_unary_f64(args: &[ArrayRef], op: impl Fn(f64) -> f64) -> Result<ArrayRef> {
    let value = args
        .first()
        .ok_or_else(|| BtError::Expression("unary function missing argument".to_string()))?;
    let value = coerce_f64(value)?;

    let mut out = Vec::with_capacity(value.len());
    for idx in 0..value.len() {
        if value.is_null(idx) {
            out.push(None);
        } else {
            out.push(Some(op(value.value(idx))));
        }
    }

    Ok(Arc::new(Float64Array::from(out)))
}

fn apply_binary_f64(args: &[ArrayRef], op: impl Fn(f64, f64) -> f64) -> Result<ArrayRef> {
    let lhs = args
        .first()
        .ok_or_else(|| BtError::Expression("binary function missing lhs".to_string()))?;
    let rhs = args
        .get(1)
        .ok_or_else(|| BtError::Expression("binary function missing rhs".to_string()))?;
    let lhs = coerce_f64(lhs)?;
    let rhs = coerce_f64(rhs)?;

    if lhs.len() != rhs.len() {
        return Err(BtError::Expression(format!(
            "binary function argument length mismatch: {} vs {}",
            lhs.len(),
            rhs.len()
        )));
    }

    let mut out = Vec::with_capacity(lhs.len());
    for idx in 0..lhs.len() {
        if lhs.is_null(idx) || rhs.is_null(idx) {
            out.push(None);
        } else {
            out.push(Some(op(lhs.value(idx), rhs.value(idx))));
        }
    }

    Ok(Arc::new(Float64Array::from(out)))
}

fn coerce_f64(array: &ArrayRef) -> Result<Float64Array> {
    if let Some(array) = array.as_any().downcast_ref::<Float64Array>() {
        return Ok(array.clone());
    }

    let mut out = Vec::with_capacity(array.len());
    if let Some(array) = array.as_any().downcast_ref::<Float32Array>() {
        for idx in 0..array.len() {
            out.push(if array.is_null(idx) {
                None
            } else {
                Some(array.value(idx) as f64)
            });
        }
    } else if let Some(array) = array.as_any().downcast_ref::<Int64Array>() {
        for idx in 0..array.len() {
            out.push(if array.is_null(idx) {
                None
            } else {
                Some(array.value(idx) as f64)
            });
        }
    } else if let Some(array) = array.as_any().downcast_ref::<Int32Array>() {
        for idx in 0..array.len() {
            out.push(if array.is_null(idx) {
                None
            } else {
                Some(array.value(idx) as f64)
            });
        }
    } else if let Some(array) = array.as_any().downcast_ref::<UInt64Array>() {
        for idx in 0..array.len() {
            out.push(if array.is_null(idx) {
                None
            } else {
                Some(array.value(idx) as f64)
            });
        }
    } else if let Some(array) = array.as_any().downcast_ref::<UInt32Array>() {
        for idx in 0..array.len() {
            out.push(if array.is_null(idx) {
                None
            } else {
                Some(array.value(idx) as f64)
            });
        }
    } else if let Some(array) = array.as_any().downcast_ref::<UInt8Array>() {
        for idx in 0..array.len() {
            out.push(if array.is_null(idx) {
                None
            } else {
                Some(array.value(idx) as f64)
            });
        }
    } else if let Some(array) = array.as_any().downcast_ref::<TimestampNanosecondArray>() {
        for idx in 0..array.len() {
            out.push(if array.is_null(idx) {
                None
            } else {
                Some(array.value(idx) as f64)
            });
        }
    } else {
        return Err(BtError::Expression(format!(
            "expected numeric array for function evaluation, got {:?}",
            array.data_type()
        )));
    }

    Ok(Float64Array::from(out))
}
