use std::sync::Arc;

use arrow::{
    array::{
        Array, ArrayRef, BooleanArray, BooleanBuilder, Float32Array, Float64Array,
        LargeStringArray, StringArray,
    },
    datatypes::DataType,
};
use luminal::prelude::*;
use mpera_frontend::op::BinaryOpKind;

use crate::{
    Result,
    luminal_interop::arrays::{MaskedF32, arrow_array_to_f32, broadcast_numeric_input},
    typecheck::binary_output_dtype,
};

pub fn binary_array(kind: &BinaryOpKind, lhs: &dyn Array, rhs: &dyn Array) -> Result<ArrayRef> {
    match kind {
        BinaryOpKind::Add | BinaryOpKind::Sub | BinaryOpKind::Mul | BinaryOpKind::Div => {
            binary_numeric_array(kind, lhs, rhs)
        }
        BinaryOpKind::LesserThan
        | BinaryOpKind::GreaterThan
        | BinaryOpKind::LesserEquals
        | BinaryOpKind::GreaterEquals => binary_numeric_comparison_array(kind, lhs, rhs),
        BinaryOpKind::Equals | BinaryOpKind::NotEquals => {
            if matches!(lhs.data_type(), DataType::Utf8 | DataType::LargeUtf8)
                && lhs.data_type() == rhs.data_type()
            {
                string_comparison_array(kind, lhs, rhs)
            } else if lhs.data_type() != rhs.data_type()
                && (matches!(lhs.data_type(), DataType::Utf8 | DataType::LargeUtf8)
                    || matches!(rhs.data_type(), DataType::Utf8 | DataType::LargeUtf8))
            {
                incompatible_type_comparison_array(kind, lhs, rhs)
            } else if matches!(lhs.data_type(), DataType::Boolean)
                && matches!(rhs.data_type(), DataType::Boolean)
            {
                bool_comparison_array(kind, lhs, rhs)
            } else {
                binary_numeric_comparison_array(kind, lhs, rhs)
            }
        }
        BinaryOpKind::And | BinaryOpKind::Or => bool_logic_array(kind, lhs, rhs),
    }
}

fn binary_numeric_array(kind: &BinaryOpKind, lhs: &dyn Array, rhs: &dyn Array) -> Result<ArrayRef> {
    let out_dtype = binary_output_dtype(kind, lhs.data_type(), rhs.data_type())?;
    let (lhs, rhs) =
        broadcast_numeric_inputs(kind, arrow_array_to_f32(lhs)?, arrow_array_to_f32(rhs)?)?;

    let validity = combine_validity(&lhs.validity, &rhs.validity);
    let (lhs_values, rhs_values) = compact_pair_values(&lhs.values, &rhs.values, &validity);
    let values = luminal_binary(kind, &lhs_values, &rhs_values)?;
    match out_dtype {
        DataType::Float32 => {
            Ok(Arc::new(Float32Array::from(with_validity(&validity, values))) as ArrayRef)
        }
        DataType::Float64 => Ok(Arc::new(Float64Array::from(
            with_validity(&validity, values)
                .into_iter()
                .map(|value| value.map(f64::from))
                .collect::<Vec<_>>(),
        )) as ArrayRef),
        _ => unreachable!(),
    }
}

fn binary_numeric_comparison_array(
    kind: &BinaryOpKind,
    lhs: &dyn Array,
    rhs: &dyn Array,
) -> Result<ArrayRef> {
    let (lhs, rhs) =
        broadcast_numeric_inputs(kind, arrow_array_to_f32(lhs)?, arrow_array_to_f32(rhs)?)?;
    let validity = combine_validity(&lhs.validity, &rhs.validity);

    let (lhs_values, rhs_values) = compact_pair_values(&lhs.values, &rhs.values, &validity);
    let values = luminal_binary_bool(kind, &lhs_values, &rhs_values)?;
    Ok(Arc::new(BooleanArray::from(with_validity(&validity, values))) as ArrayRef)
}

fn luminal_binary(kind: &BinaryOpKind, lhs: &[f32], rhs: &[f32]) -> Result<Vec<f32>> {
    let mut cx = Graph::new();
    let lhs_tensor = cx.tensor(lhs.len());
    let rhs_tensor = cx.tensor(rhs.len());
    let out = match kind {
        BinaryOpKind::Add => (lhs_tensor + rhs_tensor).output(),
        BinaryOpKind::Sub => (lhs_tensor - rhs_tensor).output(),
        BinaryOpKind::Mul => (lhs_tensor * rhs_tensor).output(),
        BinaryOpKind::Div => (lhs_tensor / rhs_tensor).output(),
        _ => {
            return Err(autark_error::Error::Reader(format!(
                "binary {:?} is not implemented yet",
                kind
            )));
        }
    };

    cx.build_search_space::<NativeRuntime>();
    let mut rt = cx.search(NativeRuntime::default(), 1);
    rt.set_data(lhs_tensor, lhs.to_vec());
    rt.set_data(rhs_tensor, rhs.to_vec());
    rt.execute(&cx.dyn_map);
    Ok(rt.get_f32(out.id).clone())
}

fn luminal_binary_bool(kind: &BinaryOpKind, lhs: &[f32], rhs: &[f32]) -> Result<Vec<bool>> {
    let mut cx = Graph::new();
    let lhs_tensor = cx.tensor(lhs.len());
    let rhs_tensor = cx.tensor(rhs.len());
    let out = match kind {
        BinaryOpKind::LesserThan => lhs_tensor.lt(rhs_tensor).cast(DType::F32).output(),
        BinaryOpKind::GreaterThan => lhs_tensor.gt(rhs_tensor).cast(DType::F32).output(),
        BinaryOpKind::LesserEquals => lhs_tensor.le(rhs_tensor).cast(DType::F32).output(),
        BinaryOpKind::GreaterEquals => lhs_tensor.ge(rhs_tensor).cast(DType::F32).output(),
        BinaryOpKind::Equals => lhs_tensor.eq(rhs_tensor).cast(DType::F32).output(),
        BinaryOpKind::NotEquals => lhs_tensor.ne(rhs_tensor).cast(DType::F32).output(),
        _ => {
            return Err(autark_error::Error::Reader(format!(
                "binary {:?} is not implemented yet",
                kind
            )));
        }
    };

    cx.build_search_space::<NativeRuntime>();
    let mut rt = cx.search(NativeRuntime::default(), 1);
    rt.set_data(lhs_tensor, lhs.to_vec());
    rt.set_data(rhs_tensor, rhs.to_vec());
    rt.execute(&cx.dyn_map);
    Ok(rt
        .get_f32(out.id)
        .iter()
        .map(|value| *value != 0.0)
        .collect())
}

fn bool_logic_array(kind: &BinaryOpKind, lhs: &dyn Array, rhs: &dyn Array) -> Result<ArrayRef> {
    let lhs = lhs
        .as_any()
        .downcast_ref::<BooleanArray>()
        .ok_or(autark_error::Error::UnsupportedArrowDataType)?;
    let rhs = rhs
        .as_any()
        .downcast_ref::<BooleanArray>()
        .ok_or(autark_error::Error::UnsupportedArrowDataType)?;
    append_bool_zip(kind, lhs, rhs, |kind, lhs, rhs| match kind {
        BinaryOpKind::And => Some(lhs && rhs),
        BinaryOpKind::Or => Some(lhs || rhs),
        _ => unreachable!(),
    })
}

fn bool_comparison_array(
    kind: &BinaryOpKind,
    lhs: &dyn Array,
    rhs: &dyn Array,
) -> Result<ArrayRef> {
    let lhs = lhs
        .as_any()
        .downcast_ref::<BooleanArray>()
        .ok_or(autark_error::Error::UnsupportedArrowDataType)?;
    let rhs = rhs
        .as_any()
        .downcast_ref::<BooleanArray>()
        .ok_or(autark_error::Error::UnsupportedArrowDataType)?;
    append_bool_zip(kind, lhs, rhs, |kind, lhs, rhs| match kind {
        BinaryOpKind::Equals => Some(lhs == rhs),
        BinaryOpKind::NotEquals => Some(lhs != rhs),
        _ => unreachable!(),
    })
}

fn incompatible_type_comparison_array(
    kind: &BinaryOpKind,
    lhs: &dyn Array,
    rhs: &dyn Array,
) -> Result<ArrayRef> {
    let len = broadcast_len(kind, lhs.len(), rhs.len())?;
    let value = match kind {
        BinaryOpKind::Equals => false,
        BinaryOpKind::NotEquals => true,
        _ => unreachable!(),
    };
    Ok(Arc::new(BooleanArray::from(vec![value; len])) as ArrayRef)
}

fn string_comparison_array(
    kind: &BinaryOpKind,
    lhs: &dyn Array,
    rhs: &dyn Array,
) -> Result<ArrayRef> {
    match lhs.data_type() {
        DataType::Utf8 => {
            let lhs = lhs.as_any().downcast_ref::<StringArray>().ok_or_else(|| {
                autark_error::Error::Reader(format!(
                    "binary string comparison expected Utf8 lhs, got {}",
                    lhs.data_type()
                ))
            })?;
            let rhs = rhs.as_any().downcast_ref::<StringArray>().ok_or_else(|| {
                autark_error::Error::Reader(format!(
                    "binary string comparison expected Utf8 rhs, got {}",
                    rhs.data_type()
                ))
            })?;
            append_string_comparison(kind, lhs, rhs)
        }
        DataType::LargeUtf8 => {
            let lhs = lhs
                .as_any()
                .downcast_ref::<LargeStringArray>()
                .ok_or_else(|| {
                    autark_error::Error::Reader(format!(
                        "binary string comparison expected LargeUtf8 lhs, got {}",
                        lhs.data_type()
                    ))
                })?;
            let rhs = rhs
                .as_any()
                .downcast_ref::<LargeStringArray>()
                .ok_or_else(|| {
                    autark_error::Error::Reader(format!(
                        "binary string comparison expected LargeUtf8 rhs, got {}",
                        rhs.data_type()
                    ))
                })?;
            append_string_comparison(kind, lhs, rhs)
        }
        _ => Err(autark_error::Error::Reader(format!(
            "binary string comparison does not support dtype {}",
            lhs.data_type()
        ))),
    }
}

fn append_string_comparison<T>(kind: &BinaryOpKind, lhs: &T, rhs: &T) -> Result<ArrayRef>
where
    T: arrow::array::Array + std::fmt::Debug,
    for<'a> &'a T: IntoIterator<Item = Option<&'a str>>,
{
    let lhs_values = lhs.into_iter().collect::<Vec<_>>();
    let rhs_values = rhs.into_iter().collect::<Vec<_>>();
    let len = broadcast_len(kind, lhs_values.len(), rhs_values.len())?;

    let mut out = BooleanBuilder::new();
    for ix in 0..len {
        let value = match (
            lhs_values[broadcast_index(lhs_values.len(), ix)],
            rhs_values[broadcast_index(rhs_values.len(), ix)],
        ) {
            (Some(lhs), Some(rhs)) => Some(match kind {
                BinaryOpKind::Equals => lhs == rhs,
                BinaryOpKind::NotEquals => lhs != rhs,
                _ => unreachable!(),
            }),
            _ => None,
        };
        match value {
            Some(value) => out.append_value(value),
            None => out.append_null(),
        }
    }
    Ok(Arc::new(out.finish()) as ArrayRef)
}

fn broadcast_numeric_inputs(
    kind: &BinaryOpKind,
    lhs: MaskedF32,
    rhs: MaskedF32,
) -> Result<(MaskedF32, MaskedF32)> {
    let len = broadcast_len(kind, lhs.values.len(), rhs.values.len())?;
    Ok((
        broadcast_numeric_input(lhs, len),
        broadcast_numeric_input(rhs, len),
    ))
}

fn broadcast_len(kind: &BinaryOpKind, lhs: usize, rhs: usize) -> Result<usize> {
    if lhs == rhs {
        return Ok(lhs);
    }
    if lhs == 1 {
        return Ok(rhs);
    }
    if rhs == 1 {
        return Ok(lhs);
    }
    Err(autark_error::Error::Reader(format!(
        "binary {:?} currently requires equal column lengths or scalar broadcast: lhs={} rhs={}",
        kind, lhs, rhs
    )))
}

fn broadcast_index(len: usize, ix: usize) -> usize {
    if len == 1 { 0 } else { ix }
}

fn boolean_value(array: &BooleanArray, ix: usize) -> Option<bool> {
    array.iter().nth(ix).flatten()
}

fn append_bool_zip(
    kind: &BinaryOpKind,
    lhs: &BooleanArray,
    rhs: &BooleanArray,
    f: impl Fn(&BinaryOpKind, bool, bool) -> Option<bool>,
) -> Result<ArrayRef> {
    let len = broadcast_len(kind, lhs.len(), rhs.len())?;
    let mut out = BooleanBuilder::new();
    for ix in 0..len {
        match (
            boolean_value(lhs, broadcast_index(lhs.len(), ix)),
            boolean_value(rhs, broadcast_index(rhs.len(), ix)),
        ) {
            (Some(lhs), Some(rhs)) => match f(kind, lhs, rhs) {
                Some(value) => out.append_value(value),
                None => out.append_null(),
            },
            _ => out.append_null(),
        }
    }
    Ok(Arc::new(out.finish()) as ArrayRef)
}

fn combine_validity(lhs: &[bool], rhs: &[bool]) -> Vec<bool> {
    lhs.iter().zip(rhs).map(|(lhs, rhs)| *lhs && *rhs).collect()
}

fn compact_pair_values(lhs: &[f32], rhs: &[f32], validity: &[bool]) -> (Vec<f32>, Vec<f32>) {
    let lhs = lhs
        .iter()
        .zip(validity)
        .filter_map(|(value, valid)| valid.then_some(*value))
        .collect();
    let rhs = rhs
        .iter()
        .zip(validity)
        .filter_map(|(value, valid)| valid.then_some(*value))
        .collect();
    (lhs, rhs)
}

fn with_validity<T>(validity: &[bool], values: Vec<T>) -> Vec<Option<T>> {
    let mut valid_values = values.into_iter();
    validity
        .iter()
        .map(|valid| valid.then(|| valid_values.next().unwrap()))
        .collect()
}
