use arrow::{
    array::RecordBatch,
    compute::{SortOptions, take},
};
use arrow_ord::sort::{SortColumn, lexsort_to_indices};
use mpera_frontend::op::OpRef;

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

pub struct OrderByBatchOp {
    pub opref: OpRef,
    pub what: OpRef,
    pub by: OpRef,
    pub ascending: bool,
}

impl OrderByBatchOp {
    pub fn new(opref: OpRef, what: OpRef, by: OpRef, ascending: bool) -> Self {
        Self {
            opref,
            what,
            by,
            ascending,
        }
    }
}

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

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

        Ok(Some(vec![order_by_batch(what, by, self.ascending)?]))
    }
}

fn order_by_batch(
    what: &RecordBatch,
    by: &RecordBatch,
    ascending: bool,
) -> crate::Result<RecordBatch> {
    if what.num_rows() != by.num_rows() {
        return Err(autark_error::Error::Reader(format!(
            "orderby requires matching row counts: what={} by={}",
            what.num_rows(),
            by.num_rows()
        )));
    }
    if by.columns().is_empty() {
        return Err(autark_error::Error::Reader(
            "orderby requires at least one sort column".to_string(),
        ));
    }

    let indices = lexsort_to_indices(
        &by.columns()
            .iter()
            .cloned()
            .map(|values| SortColumn {
                values,
                options: Some(SortOptions {
                    descending: !ascending,
                    nulls_first: true,
                }),
            })
            .collect::<Vec<_>>(),
        None,
    )?;

    let columns = what
        .columns()
        .iter()
        .map(|column| Ok(take(column.as_ref(), &indices, None)?))
        .collect::<crate::Result<Vec<_>>>()?;

    Ok(RecordBatch::try_new(what.schema(), columns)?)
}
