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

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

pub struct SliceBatchOp {
    pub opref: OpRef,
    pub on: OpRef,
    pub start: Option<isize>,
    pub end: Option<isize>,
}

impl SliceBatchOp {
    pub fn new(opref: OpRef, on: OpRef, start: Option<isize>, end: Option<isize>) -> Self {
        Self {
            opref,
            on,
            start,
            end,
        }
    }

    fn apply(&self, batch: &RecordBatch) -> RecordBatch {
        let rows = batch.num_rows() as isize;
        let start = normalize_bound(self.start, rows, 0);
        let end = normalize_bound(self.end, rows, rows).max(start);
        batch.slice(start as usize, (end - start) as usize)
    }
}

impl StatefulBatchOp for SliceBatchOp {
    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![self.apply(batch)]))
    }
}

fn normalize_bound(bound: Option<isize>, rows: isize, default: isize) -> isize {
    let bound = bound.unwrap_or(default);
    if bound < 0 {
        (rows + bound).max(0)
    } else {
        bound.min(rows)
    }
}
