use arrow::{
    array::{BooleanArray, RecordBatch},
    compute::filter_record_batch,
};
use mpera_frontend::op::OpRef;

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

pub struct FilterBatchOp {
    pub opref: OpRef,
    pub on: OpRef,
    pub mask: OpRef,
}

impl FilterBatchOp {
    pub fn new(opref: OpRef, on: OpRef, mask: OpRef) -> Self {
        Self { opref, on, mask }
    }
}

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

    fn consume(&mut self, rtdb: &RuntimeDB) -> OpOut {
        let Some(batch) = rtdb.read(self.on) else {
            return Ok(None);
        };
        let Some(mask) = rtdb.read(self.mask) else {
            return Ok(None);
        };

        let mask = mask.columns().first().ok_or_else(|| {
            autark_error::Error::Reader("filter mask requires at least one column".to_string())
        })?;
        let mask = mask
            .as_any()
            .downcast_ref::<BooleanArray>()
            .ok_or_else(|| {
                autark_error::Error::Reader("filter mask must be a boolean column".to_string())
            })?;

        Ok(Some(vec![filter_batch(batch, mask)?]))
    }
}

fn filter_batch(batch: &RecordBatch, mask: &BooleanArray) -> crate::Result<RecordBatch> {
    Ok(filter_record_batch(batch, mask)?)
}
