use std::cmp::Ordering;

use arrow::array::Date32Array;
use arrow::array::{
    Array, BinaryArray, BooleanArray, Float32Array, Float64Array, Int16Array, Int32Array,
    Int64Array, LargeBinaryArray, LargeStringArray, StringArray, UInt8Array, UInt16Array,
    UInt32Array, UInt64Array,
};

pub fn row_key(label: &str, batch: &dyn RowSource, row: usize) -> crate::Result<String> {
    let mut key = String::new();
    for column in batch.columns() {
        push_scalar_key(label, &mut key, column.as_ref(), row)?;
        key.push('\u{1f}');
    }
    Ok(key)
}

pub fn compare_scalar(
    label: &str,
    left: &dyn Array,
    left_row: usize,
    right: &dyn Array,
    right_row: usize,
) -> crate::Result<Ordering> {
    macro_rules! cmp_primitive {
        ($arr:ty) => {{
            let left = left.as_any().downcast_ref::<$arr>().ok_or_else(|| {
                autark_error::Error::Reader(format!(
                    "{label} key expected {}, got {}",
                    stringify!($arr),
                    left.data_type()
                ))
            })?;
            let right = right.as_any().downcast_ref::<$arr>().ok_or_else(|| {
                autark_error::Error::Reader(format!(
                    "{label} key expected {}, got {}",
                    stringify!($arr),
                    right.data_type()
                ))
            })?;
            Ok(left
                .iter()
                .nth(left_row)
                .flatten()
                .cmp(&right.iter().nth(right_row).flatten()))
        }};
    }

    macro_rules! cmp_float {
        ($arr:ty) => {{
            let left = left.as_any().downcast_ref::<$arr>().ok_or_else(|| {
                autark_error::Error::Reader(format!(
                    "{label} key expected {}, got {}",
                    stringify!($arr),
                    left.data_type()
                ))
            })?;
            let right = right.as_any().downcast_ref::<$arr>().ok_or_else(|| {
                autark_error::Error::Reader(format!(
                    "{label} key expected {}, got {}",
                    stringify!($arr),
                    right.data_type()
                ))
            })?;
            Ok(left
                .iter()
                .nth(left_row)
                .flatten()
                .partial_cmp(&right.iter().nth(right_row).flatten())
                .unwrap_or(Ordering::Equal))
        }};
    }

    match left.data_type() {
        arrow::datatypes::DataType::Boolean => cmp_primitive!(BooleanArray),
        arrow::datatypes::DataType::Utf8 => cmp_primitive!(StringArray),
        arrow::datatypes::DataType::LargeUtf8 => cmp_primitive!(LargeStringArray),
        arrow::datatypes::DataType::Binary => cmp_primitive!(BinaryArray),
        arrow::datatypes::DataType::LargeBinary => cmp_primitive!(LargeBinaryArray),
        arrow::datatypes::DataType::Int16 => cmp_primitive!(Int16Array),
        arrow::datatypes::DataType::Int32 => cmp_primitive!(Int32Array),
        arrow::datatypes::DataType::Int64 => cmp_primitive!(Int64Array),
        arrow::datatypes::DataType::Date32 => cmp_primitive!(Date32Array),
        arrow::datatypes::DataType::UInt8 => cmp_primitive!(UInt8Array),
        arrow::datatypes::DataType::UInt16 => cmp_primitive!(UInt16Array),
        arrow::datatypes::DataType::UInt32 => cmp_primitive!(UInt32Array),
        arrow::datatypes::DataType::UInt64 => cmp_primitive!(UInt64Array),
        arrow::datatypes::DataType::Float32 => cmp_float!(Float32Array),
        arrow::datatypes::DataType::Float64 => cmp_float!(Float64Array),
        _ => Err(autark_error::Error::Reader(format!(
            "{label} does not support key dtype {} yet",
            left.data_type()
        ))),
    }
}

pub trait RowSource {
    fn columns(&self) -> &[arrow::array::ArrayRef];
}

impl RowSource for arrow::array::RecordBatch {
    fn columns(&self) -> &[arrow::array::ArrayRef] {
        self.columns()
    }
}

fn push_scalar_key(
    label: &str,
    key: &mut String,
    array: &dyn Array,
    row: usize,
) -> crate::Result<()> {
    macro_rules! push_primitive {
        ($arr:ty) => {{
            let array = array.as_any().downcast_ref::<$arr>().ok_or_else(|| {
                autark_error::Error::Reader(format!(
                    "{label} key expected {}, got {}",
                    stringify!($arr),
                    array.data_type()
                ))
            })?;
            match array.iter().nth(row).flatten() {
                Some(value) => key.push_str(&value.to_string()),
                None => key.push_str("<null>"),
            }
            Ok(())
        }};
    }

    match array.data_type() {
        arrow::datatypes::DataType::Boolean => push_primitive!(BooleanArray),
        arrow::datatypes::DataType::Utf8 => push_string::<StringArray>("Utf8", key, array, row),
        arrow::datatypes::DataType::LargeUtf8 => {
            push_string::<LargeStringArray>("LargeUtf8", key, array, row)
        }
        arrow::datatypes::DataType::Binary => push_binary::<BinaryArray>("Binary", key, array, row),
        arrow::datatypes::DataType::LargeBinary => {
            push_binary::<LargeBinaryArray>("LargeBinary", key, array, row)
        }
        arrow::datatypes::DataType::Int16 => push_primitive!(Int16Array),
        arrow::datatypes::DataType::Int32 => push_primitive!(Int32Array),
        arrow::datatypes::DataType::Int64 => push_primitive!(Int64Array),
        arrow::datatypes::DataType::Date32 => push_primitive!(Date32Array),
        arrow::datatypes::DataType::UInt8 => push_primitive!(UInt8Array),
        arrow::datatypes::DataType::UInt16 => push_primitive!(UInt16Array),
        arrow::datatypes::DataType::UInt32 => push_primitive!(UInt32Array),
        arrow::datatypes::DataType::UInt64 => push_primitive!(UInt64Array),
        arrow::datatypes::DataType::Float32 => push_primitive!(Float32Array),
        arrow::datatypes::DataType::Float64 => push_primitive!(Float64Array),
        _ => Err(autark_error::Error::Reader(format!(
            "{label} does not support key dtype {} yet",
            array.data_type()
        ))),
    }
}

fn push_string<T>(name: &str, key: &mut String, array: &dyn Array, row: usize) -> crate::Result<()>
where
    T: Array + 'static,
    for<'a> &'a T: IntoIterator<Item = Option<&'a str>>,
{
    let array = array.as_any().downcast_ref::<T>().ok_or_else(|| {
        autark_error::Error::Reader(format!(
            "{name} key expected {name}, got {}",
            array.data_type()
        ))
    })?;
    let values = array.into_iter().collect::<Vec<_>>();
    match values[row] {
        Some(value) => key.push_str(value),
        None => key.push_str("<null>"),
    }
    Ok(())
}

fn push_binary<T>(name: &str, key: &mut String, array: &dyn Array, row: usize) -> crate::Result<()>
where
    T: Array + 'static,
    for<'a> &'a T: IntoIterator<Item = Option<&'a [u8]>>,
{
    let array = array.as_any().downcast_ref::<T>().ok_or_else(|| {
        autark_error::Error::Reader(format!(
            "{name} key expected {name}, got {}",
            array.data_type()
        ))
    })?;
    let values = array.into_iter().collect::<Vec<_>>();
    match values[row] {
        Some(value) => key.push_str(&format!("{value:?}")),
        None => key.push_str("<null>"),
    }
    Ok(())
}
