use arrow::{
    array::{
        Array, Float32Array, Float64Array, Int16Array, Int32Array, Int64Array, UInt8Array,
        UInt16Array, UInt32Array, UInt64Array,
    },
    datatypes::DataType,
};
use mpera_frontend::op::{BinaryOpKind, ReduceKind};

use crate::Result;

pub fn binary_output_dtype(
    kind: &BinaryOpKind,
    lhs: &DataType,
    rhs: &DataType,
) -> Result<DataType> {
    match kind {
        BinaryOpKind::LesserThan
        | BinaryOpKind::GreaterThan
        | BinaryOpKind::LesserEquals
        | BinaryOpKind::GreaterEquals
        | BinaryOpKind::Equals
        | BinaryOpKind::NotEquals
        | BinaryOpKind::And
        | BinaryOpKind::Or => Ok(DataType::Boolean),
        BinaryOpKind::Add | BinaryOpKind::Sub | BinaryOpKind::Mul | BinaryOpKind::Div => {
            if matches!(lhs, DataType::Float64) || matches!(rhs, DataType::Float64) {
                Ok(DataType::Float64)
            } else if matches!(lhs, DataType::Float32) || matches!(rhs, DataType::Float32) {
                Ok(DataType::Float32)
            } else {
                Err(autark_error::Error::Reader(format!(
                    "binary {:?} only supports float32/float64 for now",
                    kind
                )))
            }
        }
    }
}

pub fn reduce_output_dtype(kind: &ReduceKind) -> Result<DataType> {
    match kind {
        ReduceKind::Count => Ok(DataType::UInt64),
        ReduceKind::Mean | ReduceKind::Stdev | ReduceKind::Percentile { .. } => {
            Ok(DataType::Float64)
        }
        ReduceKind::Sum
        | ReduceKind::Product
        | ReduceKind::Min
        | ReduceKind::Max
        | ReduceKind::Last => Err(autark_error::Error::Reader(
            "reduce output dtype depends on input dtype".to_string(),
        )),
    }
}

pub fn reduce_output_dtype_for_array(kind: &ReduceKind, array: &dyn Array) -> Result<DataType> {
    match kind {
        ReduceKind::Count => Ok(DataType::UInt64),
        ReduceKind::Mean | ReduceKind::Stdev | ReduceKind::Percentile { .. } => {
            Ok(DataType::Float64)
        }
        ReduceKind::Sum
        | ReduceKind::Product
        | ReduceKind::Min
        | ReduceKind::Max
        | ReduceKind::Last => Ok(array.data_type().clone()),
    }
}

pub fn numeric_values_as_f64(array: &dyn Array) -> Result<Vec<f64>> {
    macro_rules! values_primitive {
        ($arr:ty) => {{
            let array = array
                .as_any()
                .downcast_ref::<$arr>()
                .ok_or(autark_error::Error::UnsupportedArrowDataType)?;
            Ok(array.iter().flatten().map(|value| value as f64).collect())
        }};
    }

    match array.data_type() {
        DataType::Int16 => values_primitive!(Int16Array),
        DataType::Int32 => values_primitive!(Int32Array),
        DataType::Int64 => values_primitive!(Int64Array),
        DataType::UInt8 => values_primitive!(UInt8Array),
        DataType::UInt16 => values_primitive!(UInt16Array),
        DataType::UInt32 => values_primitive!(UInt32Array),
        DataType::UInt64 => values_primitive!(UInt64Array),
        DataType::Float32 => values_primitive!(Float32Array),
        DataType::Float64 => values_primitive!(Float64Array),
        _ => Err(autark_error::Error::Reader(format!(
            "numeric operation does not support dtype {} yet",
            array.data_type()
        ))),
    }
}
