use std::sync::Arc;

use arrow::{
    array::{
        Array, ArrayRef, BooleanArray, Float32Array, Float64Array, Int16Array, Int32Array,
        Int64Array, RecordBatch, StringArray, UInt8Array, UInt16Array, UInt32Array, UInt64Array,
    },
    datatypes::{Field, Schema},
};
use mpera_frontend::op::{OpRef, UnaryOpKind};

use crate::{
    Result, luminal_interop,
    op::{OpOut, StatefulBatchOp},
    runtimedb::RuntimeDB,
};

pub struct UnaryOpBatchOp {
    pub opref: OpRef,
    pub on: OpRef,
    pub kind: UnaryOpKind,
}

impl UnaryOpBatchOp {
    pub fn new(opref: OpRef, on: OpRef, kind: UnaryOpKind) -> Self {
        Self { opref, on, kind }
    }
}

impl StatefulBatchOp for UnaryOpBatchOp {
    fn opref(&self) -> OpRef {
        self.opref
    }

    fn consume(&mut self, rtdb: &RuntimeDB) -> OpOut {
        let Some(batch) = rtdb.read(self.on) else {
            return Ok(None);
        };
        Ok(Some(vec![apply_unary(&self.kind, batch)?]))
    }
}

fn apply_unary(kind: &UnaryOpKind, batch: &RecordBatch) -> Result<RecordBatch> {
    let fields = batch
        .schema()
        .fields()
        .iter()
        .zip(batch.columns())
        .map(|(field, column)| {
            Ok(Field::new(
                field.name(),
                unary_output_dtype(kind, column.as_ref())?,
                true,
            ))
        })
        .collect::<Result<Vec<_>>>()?;
    let columns = batch
        .columns()
        .iter()
        .map(|column| unary_array(kind, column.as_ref()))
        .collect::<Result<Vec<_>>>()?;

    Ok(RecordBatch::try_new(
        Arc::new(Schema::new(fields)),
        columns,
    )?)
}

fn unary_output_dtype(kind: &UnaryOpKind, array: &dyn Array) -> Result<arrow::datatypes::DataType> {
    match kind {
        UnaryOpKind::PctChange { .. } => Ok(arrow::datatypes::DataType::Float64),
        UnaryOpKind::Neg
        | UnaryOpKind::Cumsum
        | UnaryOpKind::Diff { .. }
        | UnaryOpKind::Lead { .. }
        | UnaryOpKind::Lag { .. } => Ok(array.data_type().clone()),
    }
}

fn unary_array(kind: &UnaryOpKind, array: &dyn Array) -> Result<ArrayRef> {
    match kind {
        UnaryOpKind::Neg | UnaryOpKind::Cumsum => luminal_interop::unary_array(kind, array),
        UnaryOpKind::Lead { n } => shift_array(array, -(*n as isize)),
        UnaryOpKind::Lag { n } => shift_array(array, *n as isize),
        UnaryOpKind::Diff { n } => diff_array(array, *n),
        UnaryOpKind::PctChange { n } => pct_change_array(array, *n),
    }
}

fn shift_array(array: &dyn Array, offset: isize) -> Result<ArrayRef> {
    macro_rules! shift_primitive {
        ($arr:ty, $out:ty) => {{
            let array = array
                .as_any()
                .downcast_ref::<$arr>()
                .ok_or(autark_error::Error::UnsupportedArrowDataType)?;
            let len = array.len() as isize;
            let values = (0..len)
                .map(|ix| {
                    let src = ix - offset;
                    if src < 0 || src >= len {
                        None
                    } else {
                        array.iter().nth(src as usize).flatten()
                    }
                })
                .collect::<Vec<_>>();
            Ok(Arc::new(<$out>::from(values)) as ArrayRef)
        }};
    }

    match array.data_type() {
        arrow::datatypes::DataType::Boolean => shift_primitive!(BooleanArray, BooleanArray),
        arrow::datatypes::DataType::Utf8 => shift_primitive!(StringArray, StringArray),
        arrow::datatypes::DataType::Int16 => shift_primitive!(Int16Array, Int16Array),
        arrow::datatypes::DataType::Int32 => shift_primitive!(Int32Array, Int32Array),
        arrow::datatypes::DataType::Int64 => shift_primitive!(Int64Array, Int64Array),
        arrow::datatypes::DataType::UInt8 => shift_primitive!(UInt8Array, UInt8Array),
        arrow::datatypes::DataType::UInt16 => shift_primitive!(UInt16Array, UInt16Array),
        arrow::datatypes::DataType::UInt32 => shift_primitive!(UInt32Array, UInt32Array),
        arrow::datatypes::DataType::UInt64 => shift_primitive!(UInt64Array, UInt64Array),
        arrow::datatypes::DataType::Float32 => shift_primitive!(Float32Array, Float32Array),
        arrow::datatypes::DataType::Float64 => shift_primitive!(Float64Array, Float64Array),
        _ => Err(autark_error::Error::UnsupportedArrowDataType),
    }
}

fn diff_array(array: &dyn Array, n: usize) -> Result<ArrayRef> {
    macro_rules! diff_primitive {
        ($arr:ty, $out:ty, $native:ty) => {{
            let array = array
                .as_any()
                .downcast_ref::<$arr>()
                .ok_or(autark_error::Error::UnsupportedArrowDataType)?;
            let len = array.len();
            let values = (0..len)
                .map(|ix| {
                    if ix < n {
                        return None;
                    }
                    match (
                        array.iter().nth(ix).flatten(),
                        array.iter().nth(ix - n).flatten(),
                    ) {
                        (Some(cur), Some(prev)) => Some((cur as $native) - (prev as $native)),
                        _ => None,
                    }
                })
                .collect::<Vec<_>>();
            Ok(Arc::new(<$out>::from(values)) as ArrayRef)
        }};
    }

    match array.data_type() {
        arrow::datatypes::DataType::Int16 => diff_primitive!(Int16Array, Int16Array, i16),
        arrow::datatypes::DataType::Int32 => diff_primitive!(Int32Array, Int32Array, i32),
        arrow::datatypes::DataType::Int64 => diff_primitive!(Int64Array, Int64Array, i64),
        arrow::datatypes::DataType::UInt8 => diff_primitive!(UInt8Array, UInt8Array, u8),
        arrow::datatypes::DataType::UInt16 => diff_primitive!(UInt16Array, UInt16Array, u16),
        arrow::datatypes::DataType::UInt32 => diff_primitive!(UInt32Array, UInt32Array, u32),
        arrow::datatypes::DataType::UInt64 => diff_primitive!(UInt64Array, UInt64Array, u64),
        arrow::datatypes::DataType::Float32 => diff_primitive!(Float32Array, Float32Array, f32),
        arrow::datatypes::DataType::Float64 => diff_primitive!(Float64Array, Float64Array, f64),
        _ => Err(autark_error::Error::Reader(format!(
            "diff does not support dtype {} yet",
            array.data_type()
        ))),
    }
}

fn pct_change_array(array: &dyn Array, n: usize) -> Result<ArrayRef> {
    macro_rules! pct_primitive {
        ($arr:ty) => {{
            let array = array
                .as_any()
                .downcast_ref::<$arr>()
                .ok_or(autark_error::Error::UnsupportedArrowDataType)?;
            let len = array.len();
            let values = (0..len)
                .map(|ix| {
                    if ix < n {
                        return None;
                    }
                    match (
                        array.iter().nth(ix).flatten(),
                        array.iter().nth(ix - n).flatten(),
                    ) {
                        (Some(cur), Some(prev)) if (prev as f64) != 0.0 => {
                            Some((cur as f64 - prev as f64) / prev as f64)
                        }
                        _ => None,
                    }
                })
                .collect::<Vec<_>>();
            Ok(Arc::new(Float64Array::from(values)) as ArrayRef)
        }};
    }

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