use arrow::{
    array::{
        Array, Float32Array, Float64Array, Int8Array, Int16Array, Int32Array, Int64Array,
        UInt8Array, UInt16Array, UInt32Array, UInt64Array,
    },
    datatypes::DataType,
};

use crate::Result;

pub struct MaskedF32 {
    pub values: Vec<f32>,
    pub validity: Vec<bool>,
}

impl MaskedF32 {
    pub fn new(values: Vec<f32>, validity: Vec<bool>) -> Self {
        assert_eq!(values.len(), validity.len());
        Self { values, validity }
    }

    pub fn compact_valid_values(&self) -> Vec<f32> {
        self.values
            .iter()
            .zip(&self.validity)
            .filter_map(|(value, valid)| valid.then_some(*value))
            .collect()
    }
}

pub fn arrow_array_to_f32(array: &dyn Array) -> Result<MaskedF32> {
    let validity = (0..array.len())
        .map(|ix| !array.is_null(ix))
        .collect::<Vec<_>>();

    macro_rules! primitive_to_f32 {
        ($arr:ty) => {
            array
                .as_any()
                .downcast_ref::<$arr>()
                .ok_or(autark_error::Error::UnsupportedArrowDataType)?
                .values()
                .iter()
                .map(|value| *value as f32)
                .collect()
        };
    }

    let values = match array.data_type() {
        DataType::Float32 => array
            .as_any()
            .downcast_ref::<Float32Array>()
            .ok_or(autark_error::Error::UnsupportedArrowDataType)?
            .values()
            .to_vec(),
        DataType::Float64 => primitive_to_f32!(Float64Array),
        DataType::Int8 => primitive_to_f32!(Int8Array),
        DataType::Int16 => primitive_to_f32!(Int16Array),
        DataType::Int32 => primitive_to_f32!(Int32Array),
        DataType::Int64 => primitive_to_f32!(Int64Array),
        DataType::UInt8 => primitive_to_f32!(UInt8Array),
        DataType::UInt16 => primitive_to_f32!(UInt16Array),
        DataType::UInt32 => primitive_to_f32!(UInt32Array),
        DataType::UInt64 => primitive_to_f32!(UInt64Array),
        _ => Err(autark_error::Error::Reader(format!(
            "luminal interop only supports numeric arrays for now, got {}",
            array.data_type()
        )))?,
    };

    Ok(MaskedF32::new(values, validity))
}

pub fn broadcast_numeric_input(input: MaskedF32, len: usize) -> MaskedF32 {
    if input.values.len() == len {
        input
    } else {
        MaskedF32::new(vec![input.values[0]; len], vec![input.validity[0]; len])
    }
}
