use arrow::array::Array;
use common_error::{DaftError, DaftResult};

use crate::prelude::{AsArrow, BinaryArray, BooleanArray, Utf8Array};

impl Utf8Array {
    /// Returns an iterator of `&str` over the non-null values in this array.
    ///
    /// NOTE: this will error if there are any null values.
    /// If you need to handle nulls, use the `.iter()` method instead.
    pub fn values(&self) -> DaftResult<impl Iterator<Item = &str>> {
        let arr = self.as_arrow()?;
        if self.null_count() > 0 {
            return Err(DaftError::ComputeError(
                "Utf8Array::values with nulls".to_string(),
            ));
        }
        let iter = (0..arr.len()).map(|i| arr.value(i));
        Ok(iter)
    }
}

impl BinaryArray {
    /// Returns an iterator of `&[u8]` over the non-null values in this array.
    ///
    /// NOTE: this will error if there are any null values.
    /// If you need to handle nulls, use the `.iter()` method instead.
    pub fn values(&self) -> DaftResult<impl Iterator<Item = &[u8]>> {
        let arr = self.as_arrow()?;
        if self.null_count() > 0 {
            return Err(DaftError::ComputeError(
                "BinaryArray::values with nulls".to_string(),
            ));
        }
        let iter = (0..arr.len()).map(|i| arr.value(i));
        Ok(iter)
    }
}

impl BooleanArray {
    /// Returns an iterator of `bool` over the non-null values in this array.
    ///
    /// NOTE: this will error if there are any null values.
    /// If you need to handle nulls, use the `.iter()` method instead.
    pub fn values(&self) -> DaftResult<impl Iterator<Item = bool>> {
        if self.null_count() > 0 {
            return Err(DaftError::ComputeError(
                "BooleanArray::values with nulls".to_string(),
            ));
        }
        let bitmap = self.to_bitmap();
        let len = bitmap.len();
        Ok((0..len).map(move |i| bitmap.value(i)))
    }
}

#[cfg(test)]
mod tests {
    use crate::prelude::Utf8Array;

    #[test]
    fn test_values() {
        let array = Utf8Array::from_slice("test", &["hello", "world"]);
        let values: Vec<&str> = array.values().unwrap().collect();
        assert_eq!(values, vec!["hello", "world"]);
    }

    #[test]
    fn test_values_with_nulls() {
        let array =
            Utf8Array::from_iter("test", vec![Some("hello"), None, Some("world")].into_iter());
        let result = array.values();
        assert!(result.is_err());
        assert!(
            result
                .err()
                .unwrap()
                .to_string()
                .contains("Utf8Array::values with nulls")
        );
    }
}
