use std::collections::VecDeque;

use arrow::array::MutablePrimitiveArray;
use arrow::bitmap::MutableBitmap;
use arrow::datatypes::ArrowDataType;
use arrow::types::NativeType;
use num_traits::AsPrimitive;
use polars_error::PolarsResult;

use super::super::utils::MaybeNext;
use super::super::{utils, PagesIter};
use super::basic::{finish, PrimitiveDecoder, State as PrimitiveState};
use crate::arrow::read::deserialize::utils::{
    get_selected_rows, FilteredOptionalPageValidity, OptionalPageValidity,
};
use crate::parquet::deserialize::SliceFilteredIter;
use crate::parquet::encoding::{delta_bitpacked, Encoding};
use crate::parquet::page::{split_buffer, DataPage, DictPage};
use crate::parquet::schema::Repetition;
use crate::parquet::types::NativeType as ParquetNativeType;

/// The state of a [`DataPage`] of an integer parquet type (i32 or i64)
#[derive(Debug)]
enum State<'a, T>
where
    T: NativeType,
{
    Common(PrimitiveState<'a, T>),
    DeltaBinaryPackedRequired(delta_bitpacked::Decoder<'a>),
    DeltaBinaryPackedOptional(OptionalPageValidity<'a>, delta_bitpacked::Decoder<'a>),
    FilteredDeltaBinaryPackedRequired(SliceFilteredIter<delta_bitpacked::Decoder<'a>>),
    FilteredDeltaBinaryPackedOptional(
        FilteredOptionalPageValidity<'a>,
        delta_bitpacked::Decoder<'a>,
    ),
}

impl<'a, T> utils::PageState<'a> for State<'a, T>
where
    T: NativeType,
{
    fn len(&self) -> usize {
        match self {
            State::Common(state) => state.len(),
            State::DeltaBinaryPackedRequired(state) => state.size_hint().0,
            State::DeltaBinaryPackedOptional(state, _) => state.len(),
            State::FilteredDeltaBinaryPackedRequired(state) => state.size_hint().0,
            State::FilteredDeltaBinaryPackedOptional(state, _) => state.len(),
        }
    }
}

/// Decoder of integer parquet type
#[derive(Debug)]
struct IntDecoder<T, P, F>(PrimitiveDecoder<T, P, F>)
where
    T: NativeType,
    P: ParquetNativeType,
    i64: num_traits::AsPrimitive<P>,
    F: Fn(P) -> T;

impl<T, P, F> IntDecoder<T, P, F>
where
    T: NativeType,
    P: ParquetNativeType,
    i64: num_traits::AsPrimitive<P>,
    F: Fn(P) -> T,
{
    #[inline]
    fn new(op: F) -> Self {
        Self(PrimitiveDecoder::new(op))
    }
}

impl<'a, T, P, F> utils::Decoder<'a> for IntDecoder<T, P, F>
where
    T: NativeType,
    P: ParquetNativeType,
    i64: num_traits::AsPrimitive<P>,
    F: Copy + Fn(P) -> T,
{
    type State = State<'a, T>;
    type Dict = Vec<T>;
    type DecodedState = (Vec<T>, MutableBitmap);

    fn build_state(
        &self,
        page: &'a DataPage,
        dict: Option<&'a Self::Dict>,
    ) -> PolarsResult<Self::State> {
        let is_optional =
            page.descriptor.primitive_type.field_info.repetition == Repetition::Optional;
        let is_filtered = page.selected_rows().is_some();

        match (page.encoding(), dict, is_optional, is_filtered) {
            (Encoding::DeltaBinaryPacked, _, false, false) => {
                let (_, _, values) = split_buffer(page)?;
                Ok(delta_bitpacked::Decoder::try_new(values)
                    .map(State::DeltaBinaryPackedRequired)?)
            },
            (Encoding::DeltaBinaryPacked, _, true, false) => {
                let (_, _, values) = split_buffer(page)?;
                Ok(State::DeltaBinaryPackedOptional(
                    OptionalPageValidity::try_new(page)?,
                    delta_bitpacked::Decoder::try_new(values)?,
                ))
            },
            (Encoding::DeltaBinaryPacked, _, false, true) => {
                let (_, _, values) = split_buffer(page)?;
                let values = delta_bitpacked::Decoder::try_new(values)?;

                let rows = get_selected_rows(page);
                let values = SliceFilteredIter::new(values, rows);

                Ok(State::FilteredDeltaBinaryPackedRequired(values))
            },
            (Encoding::DeltaBinaryPacked, _, true, true) => {
                let (_, _, values) = split_buffer(page)?;
                let values = delta_bitpacked::Decoder::try_new(values)?;

                Ok(State::FilteredDeltaBinaryPackedOptional(
                    FilteredOptionalPageValidity::try_new(page)?,
                    values,
                ))
            },
            _ => self.0.build_state(page, dict).map(State::Common),
        }
    }

    fn with_capacity(&self, capacity: usize) -> Self::DecodedState {
        self.0.with_capacity(capacity)
    }

    fn extend_from_state(
        &self,
        state: &mut Self::State,
        decoded: &mut Self::DecodedState,
        remaining: usize,
    ) -> PolarsResult<()> {
        let (values, validity) = decoded;
        match state {
            State::Common(state) => self.0.extend_from_state(state, decoded, remaining)?,
            State::DeltaBinaryPackedRequired(state) => {
                values.extend(
                    state
                        .by_ref()
                        .map(|x| x.unwrap().as_())
                        .map(self.0.op)
                        .take(remaining),
                );
            },
            State::DeltaBinaryPackedOptional(page_validity, page_values) => {
                utils::extend_from_decoder(
                    validity,
                    page_validity,
                    Some(remaining),
                    values,
                    page_values
                        .by_ref()
                        .map(|x| x.unwrap().as_())
                        .map(self.0.op),
                )
            },
            State::FilteredDeltaBinaryPackedRequired(page) => {
                values.extend(
                    page.by_ref()
                        .map(|x| x.unwrap().as_())
                        .map(self.0.op)
                        .take(remaining),
                );
            },
            State::FilteredDeltaBinaryPackedOptional(page_validity, page_values) => {
                utils::extend_from_decoder(
                    validity,
                    page_validity,
                    Some(remaining),
                    values,
                    page_values
                        .by_ref()
                        .map(|x| x.unwrap().as_())
                        .map(self.0.op),
                );
            },
        }
        Ok(())
    }

    fn deserialize_dict(&self, page: &DictPage) -> Self::Dict {
        self.0.deserialize_dict(page)
    }
}

/// An [`Iterator`] adapter over [`PagesIter`] assumed to be encoded as primitive arrays
/// encoded as parquet integer types
#[derive(Debug)]
pub struct IntegerIter<T, I, P, F>
where
    I: PagesIter,
    T: NativeType,
    P: ParquetNativeType,
    F: Fn(P) -> T,
{
    iter: I,
    data_type: ArrowDataType,
    items: VecDeque<(Vec<T>, MutableBitmap)>,
    remaining: usize,
    chunk_size: Option<usize>,
    dict: Option<Vec<T>>,
    op: F,
    phantom: std::marker::PhantomData<P>,
}

impl<T, I, P, F> IntegerIter<T, I, P, F>
where
    I: PagesIter,
    T: NativeType,

    P: ParquetNativeType,
    F: Copy + Fn(P) -> T,
{
    pub fn new(
        iter: I,
        data_type: ArrowDataType,
        num_rows: usize,
        chunk_size: Option<usize>,
        op: F,
    ) -> Self {
        Self {
            iter,
            data_type,
            items: VecDeque::new(),
            dict: None,
            remaining: num_rows,
            chunk_size,
            op,
            phantom: Default::default(),
        }
    }
}

impl<T, I, P, F> Iterator for IntegerIter<T, I, P, F>
where
    I: PagesIter,
    T: NativeType,
    P: ParquetNativeType,
    i64: num_traits::AsPrimitive<P>,
    F: Copy + Fn(P) -> T,
{
    type Item = PolarsResult<MutablePrimitiveArray<T>>;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            let maybe_state = utils::next(
                &mut self.iter,
                &mut self.items,
                &mut self.dict,
                &mut self.remaining,
                self.chunk_size,
                &IntDecoder::new(self.op),
            );
            match maybe_state {
                MaybeNext::Some(Ok((values, validity))) => {
                    return Some(Ok(finish(&self.data_type, values, validity)))
                },
                MaybeNext::Some(Err(e)) => return Some(Err(e)),
                MaybeNext::None => return None,
                MaybeNext::More => continue,
            }
        }
    }
}
