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

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

pub enum ColumnsSource {
    Input(usize),
    Op(OpRef),
}

pub struct ColumnsBatchOp {
    pub opref: OpRef,
    pub source: ColumnsSource,
    pub selector: ColumnSelector,
}

impl ColumnsBatchOp {
    pub fn new(opref: OpRef, source: ColumnsSource, selector: ColumnSelector) -> Self {
        Self {
            opref,
            source,
            selector,
        }
    }

    fn read<'a>(&self, rtdb: &'a RuntimeDB) -> Option<&'a RecordBatch> {
        match self.source {
            ColumnsSource::Input(index) => rtdb.input(index),
            ColumnsSource::Op(opref) => rtdb.read(opref),
        }
    }

    fn project(&self, batch: &RecordBatch) -> RecordBatch {
        match &self.selector {
            ColumnSelector::All => batch.clone(),
            ColumnSelector::Names(columns) => {
                let indices = columns
                    .iter()
                    .map(|name| {
                        batch
                            .schema()
                            .index_of(name)
                            .unwrap_or_else(|_| panic!("missing column `{name}`"))
                    })
                    .collect::<Vec<_>>();
                batch
                    .project(&indices)
                    .unwrap_or_else(|err| panic!("failed to project columns by name: {err}"))
            }
            ColumnSelector::Indices(indices) => {
                let indices = indices
                    .iter()
                    .map(|ix| {
                        if *ix < batch.num_columns() {
                            *ix
                        } else {
                            panic!(
                                "column index {ix} out of bounds for {} columns",
                                batch.num_columns()
                            );
                        }
                    })
                    .collect::<Vec<_>>();
                batch
                    .project(&indices)
                    .unwrap_or_else(|err| panic!("failed to project columns by index: {err}"))
            }
        }
    }
}

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

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

        Ok(Some(vec![self.project(batch)]))
    }
}
