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

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

pub struct ConcatBatchOp {
    pub opref: OpRef,
    pub who: Vec<OpRef>,
}

impl ConcatBatchOp {
    pub fn new(opref: OpRef, who: Vec<OpRef>) -> Self {
        Self { opref, who }
    }
}

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

    fn consume(&mut self, rtdb: &RuntimeDB) -> OpOut {
        let mut batches = Vec::<RecordBatch>::with_capacity(self.who.len());
        for opref in &self.who {
            let Some(batch) = rtdb.read(*opref) else {
                return Ok(None);
            };
            batches.push(batch.clone());
        }

        let first = batches.first().ok_or_else(|| {
            autark_error::Error::Reader("concat requires at least one input".to_string())
        })?;
        let batch = if batches.len() == 1 {
            first.clone()
        } else {
            concat_batches(&first.schema(), &batches)?
        };
        Ok(Some(vec![batch]))
    }
}
