use arrow::array::RecordBatch;
use mpera_frontend::op::OpRef;

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

pub struct RollingBatchOp {
    pub opref: OpRef,
    pub on: OpRef,
    pub n: usize,
}

impl RollingBatchOp {
    pub fn new(opref: OpRef, on: OpRef, n: usize) -> Self {
        Self { opref, on, n }
    }
}

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

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

fn rolling_batches(batch: &RecordBatch, n: usize) -> crate::Result<Vec<RecordBatch>> {
    if n == 0 {
        return Err(autark_error::Error::Reader(
            "rolling requires n > 0".to_string(),
        ));
    }
    if batch.num_rows() < n {
        return Ok(Vec::new());
    }

    Ok((0..=(batch.num_rows() - n))
        .map(|start| batch.slice(start, n))
        .collect())
}
