use crate::{lit::Literal, prelude::*};

impl NullArray {
    pub fn get_lit(&self, idx: usize) -> Literal {
        assert!(
            idx < self.len(),
            "Out of bounds: {} vs len: {}",
            idx,
            self.len()
        );

        Literal::Null
    }
}

impl StructArray {
    pub fn get_lit(&self, idx: usize) -> Literal {
        assert!(
            idx < self.len(),
            "Out of bounds: {} vs len: {}",
            idx,
            self.len()
        );

        if self.is_valid(idx) {
            Literal::Struct(
                self.children
                    .iter()
                    .map(|child| (child.field().clone(), child.get_lit(idx)))
                    .collect(),
            )
        } else {
            Literal::Null
        }
    }
}

#[cfg(feature = "python")]
impl PythonArray {
    pub fn get_lit(&self, idx: usize) -> Literal {
        use pyo3::prelude::*;

        assert!(
            idx < self.len(),
            "Out of bounds: {} vs len: {}",
            idx,
            self.len()
        );

        let v = self.get(idx);
        if Python::with_gil(|py| v.is_none(py)) {
            Literal::Null
        } else {
            Literal::Python(v.into())
        }
    }
}

fn map_or_null<T, U, F>(o: Option<T>, f: F) -> Literal
where
    F: FnOnce(U) -> Literal,
    U: From<T>,
{
    match o {
        Some(v) => f(v.into()),
        None => Literal::Null,
    }
}

macro_rules! impl_array_get_lit {
    ($type:ty, $variant:ident) => {
        impl $type {
            pub fn get_lit(&self, idx: usize) -> Literal {
                assert!(
                    idx < self.len(),
                    "Out of bounds: {} vs len: {}",
                    idx,
                    self.len()
                );
                map_or_null(self.get(idx), Literal::$variant)
            }
        }
    };

    ($type:ty, $dtype:pat => $mapper:expr) => {
        impl $type {
            pub fn get_lit(&self, idx: usize) -> Literal {
                assert!(
                    idx < self.len(),
                    "Out of bounds: {} vs len: {}",
                    idx,
                    self.len()
                );
                let $dtype = self.data_type() else {
                    unreachable!(
                        "Unexpected data type for {}: {}",
                        stringify!($type),
                        self.data_type()
                    )
                };

                map_or_null(self.get(idx), $mapper)
            }
        }
    };
}

macro_rules! unimplemented_get_lit {
    ($type:ty) => {
        impl $type {
            pub fn get_lit(&self, _: usize) -> Literal {
                unimplemented!("Unsupported data type: {}", self.data_type())
            }
        }
    };
}

impl_array_get_lit!(BooleanArray, Boolean);
impl_array_get_lit!(BinaryArray, Binary);
impl_array_get_lit!(FixedSizeBinaryArray, Binary);
impl_array_get_lit!(Int8Array, Int8);
impl_array_get_lit!(Int16Array, Int16);
impl_array_get_lit!(Int32Array, Int32);
impl_array_get_lit!(Int64Array, Int64);
impl_array_get_lit!(UInt8Array, UInt8);
impl_array_get_lit!(UInt16Array, UInt16);
impl_array_get_lit!(UInt32Array, UInt32);
impl_array_get_lit!(UInt64Array, UInt64);
impl_array_get_lit!(Float32Array, Float32);
impl_array_get_lit!(Float64Array, Float64);
impl_array_get_lit!(Utf8Array, Utf8);
impl_array_get_lit!(IntervalArray, Interval);
impl_array_get_lit!(DateArray, Date);
impl_array_get_lit!(ListArray, List);
impl_array_get_lit!(FixedSizeListArray, List);

impl_array_get_lit!(Decimal128Array, DataType::Decimal128(precision, scale) => |v| Literal::Decimal(v, *precision as _, *scale as _));
impl_array_get_lit!(TimestampArray, DataType::Timestamp(tu, tz) => |v| Literal::Timestamp(v, *tu, tz.clone()));
impl_array_get_lit!(TimeArray, DataType::Time(tu) => |v| Literal::Time(v, *tu));
impl_array_get_lit!(DurationArray, DataType::Duration(tu) => |v| Literal::Duration(v, *tu));

unimplemented_get_lit!(ExtensionArray);
unimplemented_get_lit!(ImageArray);
unimplemented_get_lit!(FixedShapeImageArray);
unimplemented_get_lit!(TensorArray);
unimplemented_get_lit!(EmbeddingArray);
unimplemented_get_lit!(FixedShapeTensorArray);
unimplemented_get_lit!(SparseTensorArray);
unimplemented_get_lit!(FixedShapeSparseTensorArray);
unimplemented_get_lit!(MapArray);
