mod array_impl;
mod from;
pub mod from_lit;
mod ops;
mod serdes;
mod series_like;
mod utils;
use std::{hash::Hash, ops::Sub, sync::Arc};

pub use array_impl::{ArrayWrapper, IntoSeries};
use arrow::array::ArrayRef;
use common_display::table_display::{StrValue, make_comfy_table};
use common_error::DaftResult;
use derive_more::Display;
use indexmap::{IndexMap, map::RawEntryApiV1};
pub use ops::cast_series_to_supertype;

pub(crate) use self::series_like::SeriesLike;
use crate::{
    array::{
        DataArray,
        ops::{DaftCompare, from_arrow::FromArrow, full::FullNull},
    },
    datatypes::{DaftDataType, DaftNumericType, DataType, Field, FieldRef, NumericNative},
    lit::Literal,
    utils::identity_hash_set::{IdentityBuildHasher, IndexHash},
    with_match_daft_types,
};

#[derive(Clone, Debug, Display)]
#[display("{}\n", self.to_comfy_table())]
pub struct Series {
    pub inner: Arc<dyn SeriesLike>,
}

impl PartialEq for Series {
    fn eq(&self, other: &Self) -> bool {
        match self.equal(other) {
            Ok(arr) => arr.into_iter().all(|x| x.unwrap_or(false)),
            Err(_) => false,
        }
    }
}

impl Hash for Series {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        let hash_result = self.hash(None);
        match hash_result {
            Ok(hash) => hash.into_iter().for_each(|i| i.hash(state)),
            Err(..) => panic!("Failed to hash series"),
        }
    }
}

impl Series {
    /// Build a hashset of the [`IndexHash`]s of each element in this [`Series`].
    ///
    /// The returned hashset can be used to probe for the existence of a given element in this [`Series`].
    /// Its length can also be used to determine the *exact* number of unique elements in this [`Series`].
    ///
    /// # Note
    /// 1. This function returns an `IndexMap<X, ()>` rather than a `HashSet<X>`. These two types are functionally equivalent.
    ///
    /// 2. `NULL`s are *not* inserted into the returned hashset. They won't be counted towards the final number of unique elements.
    pub fn build_probe_table_without_nulls(
        &self,
    ) -> DaftResult<IndexMap<IndexHash, (), IdentityBuildHasher>> {
        // Building a comparator function over a series of type `NULL` will result in a failure.
        // (I.e., `let comparator = build_is_equal(..)` will fail).
        //
        // Therefore, exit early with an empty hashmap.
        if matches!(self.data_type(), DataType::Null) {
            return Ok(IndexMap::default());
        }

        const DEFAULT_SIZE: usize = 20;
        let hashed_series = self.hash_with_nulls(None)?;

        let mut probe_table =
            IndexMap::<IndexHash, (), IdentityBuildHasher>::with_capacity_and_hasher(
                DEFAULT_SIZE,
                Default::default(),
            );

        for (idx, hash) in hashed_series.into_iter().enumerate() {
            let hash = match hash {
                Some(hash) => hash,
                None => continue,
            };
            let entry = probe_table.raw_entry_v1().from_hash(hash, |other| {
                (hash == other.hash) && self.get_lit(idx).eq(&self.get_lit(other.idx as _))
            });
            if entry.is_none() {
                probe_table.insert(
                    IndexHash {
                        idx: idx as u64,
                        hash,
                    },
                    (),
                );
            }
        }

        Ok(probe_table)
    }

    pub fn to_arrow(&self) -> DaftResult<ArrayRef> {
        self.inner.to_arrow()
    }

    pub fn from_arrow<F: Into<FieldRef>>(field: F, arrow_arr: ArrayRef) -> DaftResult<Self> {
        let field = field.into();
        with_match_daft_types!(field.dtype, |$T| {
            Ok(<<$T as DaftDataType>::ArrayType as FromArrow>::from_arrow(field, arrow_arr)?.into_series())
        })
    }

    /// Creates a Series that is all nulls
    pub fn full_null(name: &str, dtype: &DataType, length: usize) -> Self {
        with_match_daft_types!(dtype, |$T| {
            <<$T as DaftDataType>::ArrayType as FullNull>::full_null(name, dtype, length).into_series()
        })
    }

    /// Creates an empty [`Series`]
    pub fn empty(field_name: &str, dtype: &DataType) -> Self {
        with_match_daft_types!(dtype, |$T| {
            <<$T as DaftDataType>::ArrayType as FullNull>::empty(field_name, dtype).into_series()
        })
    }

    pub fn data_type(&self) -> &DataType {
        self.inner.data_type()
    }

    pub fn to_floating_data_type(&self) -> DaftResult<DataType> {
        self.inner.data_type().to_floating_representation()
    }

    pub fn name(&self) -> &str {
        self.inner.name()
    }

    #[must_use]
    pub fn rename<S: AsRef<str>>(&self, name: S) -> Self {
        self.inner.rename(name.as_ref())
    }

    pub fn field(&self) -> &Field {
        self.inner.field()
    }
    pub fn as_physical(&self) -> DaftResult<Self> {
        let physical_dtype = self.data_type().to_physical();
        if &physical_dtype == self.data_type() {
            Ok(self.clone())
        } else {
            self.inner.cast(&physical_dtype)
        }
    }

    pub fn to_comfy_table(&self) -> comfy_table::Table {
        let field = self.field();
        let field_disp = format!("{}\n---\n{}", field.name, field.dtype);

        make_comfy_table(
            [field_disp].as_slice(),
            Some([self as &dyn StrValue].as_slice()),
            Some(self.len()),
            Some(80),
        )
    }

    pub fn with_nulls(&self, nulls: Option<arrow::buffer::NullBuffer>) -> DaftResult<Self> {
        self.inner.with_nulls(nulls)
    }

    pub fn nulls(&self) -> Option<&arrow::buffer::NullBuffer> {
        self.inner.nulls()
    }

    pub fn null_count(&self) -> usize {
        self.nulls().map(|n| n.null_count()).unwrap_or(0)
    }

    pub fn is_valid(&self, idx: usize) -> bool {
        let Some(nulls) = self.nulls() else {
            return true;
        };
        nulls.is_valid(idx)
    }

    /// Attempts to downcast the Series to a primitive slice
    /// This will return an error if the Series is not of the physical type `T`
    /// # Example
    /// ```ignore
    /// let i32_arr: &[i32] = series.try_as_slice::<i32>()?;
    ///
    /// let f64_arr: &[f64] = series.try_as_slice::<f64>()?;
    /// ```
    pub fn try_as_slice<N: NumericNative>(
        &self,
    ) -> DaftResult<&[<N::DAFTTYPE as DaftNumericType>::Native]>
    where
        N::DAFTTYPE: DaftNumericType + DaftDataType,
    {
        let data: &DataArray<N::DAFTTYPE> = self.downcast()?;
        Ok(data.as_slice())
    }

    /// Helper function to check that two series of floats are within some `epsilon` of each other.
    pub fn fuzzy_eq(&self, other: &Self, epsilon: &Self) -> bool {
        if self.data_type() != other.data_type() {
            return false;
        }
        match self.data_type() {
            DataType::Float32 | DataType::Float64 => {
                let diff = self
                    .sub(other)
                    .expect("Failed to subtract one series from the other")
                    .abs()
                    .expect("Failed to get absolute difference between the two given series");
                match diff.lte(epsilon) {
                    Ok(arr) => arr.into_iter().all(|x| x.unwrap_or(false)),
                    Err(_) => false,
                }
            }
            _ => self == other,
        }
    }

    /// Get the value at `idx` as a [`Literal`]
    ///
    /// Panics if `idx` is out of bounds
    pub fn get_lit(&self, idx: usize) -> Literal {
        self.inner.get_lit(idx)
    }

    pub fn to_literals(&self) -> impl ExactSizeIterator<Item = Literal> + use<'_> {
        (0..self.len()).map(|i| self.get_lit(i))
    }
}

#[macro_export]
/// Convenient macro to create a Series from elements for testing purposes. Works for any types that implement `Into<Literal>`.
macro_rules! series {
    ($($element:expr),+) => {
        {
            // put into a vec first for compile-time type consistency checking
            let elements = vec![$($element),+];
            let elements_lit = elements.into_iter().map(Literal::from).collect::<Vec<_>>();
            Series::from_literals(elements_lit).unwrap()
        }
    };
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::datatypes::Int64Array;

    #[test]
    fn test_build_probe_table_without_nulls_basic() -> DaftResult<()> {
        let series = Int64Array::from_vec("", vec![1i64, 2, 3, 4, 5]).into_series();
        let probe_table = series.build_probe_table_without_nulls()?;
        assert_eq!(probe_table.len(), 5);
        Ok(())
    }

    #[test]
    fn test_build_probe_table_without_nulls_with_duplicates() -> DaftResult<()> {
        let series = Int64Array::from_vec("", vec![1i64, 2, 2, 3, 3, 3]).into_series();
        let probe_table = series.build_probe_table_without_nulls()?;
        assert_eq!(probe_table.len(), 3);
        Ok(())
    }

    #[test]
    fn test_build_probe_table_without_nulls_with_nulls() -> DaftResult<()> {
        let series = vec![Some(1i64), None, Some(2), None, Some(1)]
            .into_iter()
            .collect::<Int64Array>()
            .into_series();
        let probe_table = series.build_probe_table_without_nulls()?;
        assert_eq!(probe_table.len(), 2); // Only 1 and 2, nulls not counted
        Ok(())
    }

    #[test]
    fn test_build_probe_table_without_nulls_all_nulls() -> DaftResult<()> {
        let series = vec![None, None, None]
            .into_iter()
            .collect::<Int64Array>()
            .into_series();
        let probe_table = series.build_probe_table_without_nulls()?;
        assert_eq!(probe_table.len(), 0);
        Ok(())
    }

    #[test]
    fn test_build_probe_table_without_nulls_empty() -> DaftResult<()> {
        let series = Vec::<Option<i64>>::new()
            .into_iter()
            .collect::<Int64Array>()
            .into_series();
        let probe_table = series.build_probe_table_without_nulls()?;
        assert_eq!(probe_table.len(), 0);
        Ok(())
    }

    #[test]
    fn test_build_probe_table_without_nulls_null_dtype() -> DaftResult<()> {
        let series = Series::full_null("a", &DataType::Null, 5);
        let probe_table = series.build_probe_table_without_nulls()?;
        assert_eq!(probe_table.len(), 0);
        Ok(())
    }

    #[test]
    fn test_build_probe_table_without_nulls_all_same() -> DaftResult<()> {
        let series = Int64Array::from_vec("", vec![42i64; 100]).into_series();
        let probe_table = series.build_probe_table_without_nulls()?;
        assert_eq!(probe_table.len(), 1);
        Ok(())
    }
}
