use std::{iter::repeat_n, sync::Arc};

use arrow::buffer::NullBuffer;
use common_error::DaftResult;

use super::{DaftCountAggable, GroupIndices};
#[cfg(feature = "python")]
use crate::prelude::PythonArray;
use crate::{
    array::{ListArray, StructArray},
    count_mode::CountMode,
    datatypes::*,
};

/// Helper to perform a grouped count on a validity map of type arrow::buffer::NullBuffer
fn grouped_count_arrow_bitmap(
    groups: &GroupIndices,
    mode: &CountMode,
    arrow_bitmap: Option<&NullBuffer>,
) -> Vec<u64> {
    match mode {
        CountMode::All => groups.iter().map(|g| g.len() as u64).collect(),
        CountMode::Valid => match arrow_bitmap {
            None => groups.iter().map(|g| g.len() as u64).collect(), // Equivalent to CountMode::All
            Some(nulls) => groups
                .iter()
                .map(|g| g.iter().map(|i| nulls.is_valid(*i as usize) as u64).sum())
                .collect(),
        },
        CountMode::Null => match arrow_bitmap {
            None => repeat_n(0, groups.len()).collect(), // None of the values are Null
            Some(nulls) => groups
                .iter()
                .map(|g| g.iter().map(|i| !nulls.is_valid(*i as usize) as u64).sum())
                .collect(),
        },
    }
}

/// Helper to perform a count on a validity map of type NullBuffer
fn count_arrow_bitmap(mode: &CountMode, arrow_bitmap: Option<&NullBuffer>, arr_len: usize) -> u64 {
    match mode {
        CountMode::All => arr_len as u64,
        CountMode::Valid => match arrow_bitmap {
            None => arr_len as u64,
            Some(nulls) => (nulls.len() - nulls.null_count()) as u64,
        },
        CountMode::Null => match arrow_bitmap {
            None => 0,
            Some(nulls) => nulls.null_count() as u64,
        },
    }
}

impl<T> DaftCountAggable for &DataArray<T>
where
    T: DaftPhysicalType,
{
    type Output = DaftResult<DataArray<UInt64Type>>;

    fn count(&self, mode: CountMode) -> Self::Output {
        let count = if self.data_type() == &DataType::Null {
            match &mode {
                CountMode::All => self.len() as u64,
                CountMode::Valid => 0u64,
                CountMode::Null => self.len() as u64,
            }
        } else {
            count_arrow_bitmap(&mode, self.nulls(), self.len())
        };
        let result_arrow_array = arrow::array::UInt64Array::new_scalar(count);

        DataArray::<UInt64Type>::from_arrow(
            Arc::new(Field::new(self.field.name.clone(), DataType::UInt64)),
            Arc::new(result_arrow_array.into_inner()),
        )
    }
    fn grouped_count(&self, groups: &GroupIndices, mode: CountMode) -> Self::Output {
        let counts_per_group: Vec<u64> = if self.data_type() == &DataType::Null {
            match &mode {
                CountMode::All => groups.iter().map(|g| g.len() as u64).collect(),
                CountMode::Valid => repeat_n(0, groups.len()).collect(),
                CountMode::Null => groups.iter().map(|g| g.len() as u64).collect(),
            }
        } else {
            grouped_count_arrow_bitmap(groups, &mode, self.nulls())
        };
        Ok(DataArray::<UInt64Type>::from_vec(
            self.field.name.as_ref(),
            counts_per_group,
        ))
    }
}

macro_rules! impl_daft_count_aggable {
    ($arr:ident) => {
        impl DaftCountAggable for &$arr {
            type Output = DaftResult<DataArray<UInt64Type>>;

            fn count(&self, mode: CountMode) -> Self::Output {
                let count = count_arrow_bitmap(&mode, self.nulls(), self.len());
                let result_arrow_array = arrow::array::UInt64Array::new_scalar(count);

                DataArray::<UInt64Type>::from_arrow(
                    Arc::new(Field::new(self.field().name.clone(), DataType::UInt64)),
                    Arc::new(result_arrow_array.into_inner()),
                )
            }

            fn grouped_count(&self, groups: &GroupIndices, mode: CountMode) -> Self::Output {
                let counts_per_group: Vec<_> =
                    grouped_count_arrow_bitmap(groups, &mode, self.nulls());
                Ok(DataArray::<UInt64Type>::from_vec(
                    self.field().name.as_ref(),
                    counts_per_group,
                ))
            }
        }
    };
}

impl_daft_count_aggable!(FixedSizeListArray);
impl_daft_count_aggable!(ListArray);
impl_daft_count_aggable!(StructArray);

#[cfg(feature = "python")]
impl_daft_count_aggable!(PythonArray);
