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

pub struct Value {
    current: Option<RecordBatch>,
    all: Vec<RecordBatch>,
}

pub struct RuntimeDB {
    opstore: Vec<Value>,
    inp: Vec<Option<RecordBatch>>,
}

impl RuntimeDB {
    pub fn new(op_count: usize) -> RuntimeDB {
        RuntimeDB {
            opstore: std::iter::repeat_with(|| Value {
                current: None,
                all: Vec::new(),
            })
            .take(op_count)
            .collect(),
            inp: Vec::new(),
        }
    }

    pub fn input(&self, index: usize) -> Option<&RecordBatch> {
        self.inp.get(index)?.as_ref()
    }

    pub fn set_input(&mut self, index: usize, batch: RecordBatch) {
        self.ensure_input(index);
        self.inp.iter_mut().for_each(|slot| *slot = None);
        self.inp[index] = Some(batch);
    }

    pub fn read(&self, opref: OpRef) -> Option<&RecordBatch> {
        self.opstore.get(opref.0)?.current.as_ref()
    }

    pub fn read_batches(&self, opref: OpRef) -> Option<&[RecordBatch]> {
        let value = self.opstore.get(opref.0)?;
        if value.all.is_empty() {
            return None;
        }
        Some(value.all.as_slice())
    }

    pub fn write(&mut self, opref: OpRef, batches: Vec<RecordBatch>) {
        self.ensure(opref);
        let value = &mut self.opstore[opref.0];
        value.current = batches.last().cloned();
        value.all.extend(batches);
    }

    pub fn block(&mut self, opref: OpRef) {
        self.ensure(opref);
        self.opstore[opref.0].current = None;
    }

    fn ensure(&mut self, opref: OpRef) {
        if opref.0 >= self.opstore.len() {
            self.opstore.resize_with(opref.0 + 1, || Value {
                current: None,
                all: Vec::new(),
            });
        }
    }

    fn ensure_input(&mut self, index: usize) {
        if index >= self.inp.len() {
            self.inp.resize_with(index + 1, || None);
        }
    }
}
