use std::sync::Arc;

use arrow::{
    array::{Array, ArrayRef, RecordBatch},
    datatypes::{Field, Schema},
};
use mpera_frontend::op::{OpRef, ReduceKind};

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

pub struct ReduceBatchOp {
    pub opref: OpRef,
    pub on: OpRef,
    pub kind: ReduceKind,
}

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

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

    fn consume(&mut self, rtdb: &RuntimeDB) -> OpOut {
        let Some(batches) = rtdb.read_batches(self.on) else {
            return Ok(None);
        };
        Ok(Some(
            batches
                .iter()
                .map(|batch| reduce_batch(&self.kind, batch))
                .collect::<Result<Vec<_>>>()?,
        ))
    }
}

fn reduce_batch(kind: &ReduceKind, batch: &RecordBatch) -> Result<RecordBatch> {
    let fields = batch
        .schema()
        .fields()
        .iter()
        .zip(batch.columns())
        .map(|(field, column)| reduce_field(kind, field, column.as_ref()))
        .collect::<Result<Vec<_>>>()?;
    let columns = batch
        .columns()
        .iter()
        .map(|column| reduce_array(kind, column.as_ref()))
        .collect::<Result<Vec<_>>>()?;

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

fn reduce_field(kind: &ReduceKind, field: &Field, array: &dyn Array) -> Result<Field> {
    let data_type = reduce_output_dtype_for_array(kind, array)?;
    Ok(Field::new(field.name(), data_type, true))
}

fn reduce_array(kind: &ReduceKind, array: &dyn Array) -> Result<ArrayRef> {
    luminal_interop::reduce_array(kind, array)
}
