use std::sync::Arc;

use arrow::{
    array::{ArrayRef, BooleanArray, Datum, RecordBatch, UInt32Array},
    compute::take,
    datatypes::{Field, Schema},
};
use arrow_select::zip::zip;
use mpera_frontend::op::OpRef;

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

pub struct WhereBatchOp {
    pub opref: OpRef,
    pub mask: OpRef,
    pub then_: OpRef,
    pub else_: OpRef,
}

impl WhereBatchOp {
    pub fn new(opref: OpRef, mask: OpRef, then_: OpRef, else_: OpRef) -> Self {
        Self {
            opref,
            mask,
            then_,
            else_,
        }
    }
}

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

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

        Ok(Some(vec![where_batch(mask, then_, else_)?]))
    }
}

fn where_batch(
    mask: &RecordBatch,
    then_: &RecordBatch,
    else_: &RecordBatch,
) -> crate::Result<RecordBatch> {
    let mask = mask.columns().first().ok_or_else(|| {
        autark_error::Error::Reader("where mask requires at least one column".to_string())
    })?;
    let mask = mask
        .as_any()
        .downcast_ref::<BooleanArray>()
        .ok_or_else(|| {
            autark_error::Error::Reader("where mask must be a boolean column".to_string())
        })?;

    let then_schema = then_.schema();
    let else_schema = else_.schema();
    let base_fields = if then_schema.fields().len() >= else_schema.fields().len() {
        then_schema.fields().clone()
    } else {
        else_schema.fields().clone()
    };

    let columns = (0..base_fields.len())
        .map(|ix| {
            let truthy = column_for_broadcast(then_, ix, mask.len())?;
            let falsy = column_for_broadcast(else_, ix, mask.len())?;
            if truthy.data_type() != falsy.data_type() {
                return Err(autark_error::Error::Reader(format!(
                    "where requires matching branch dtypes: then={} else={}",
                    truthy.data_type(),
                    falsy.data_type()
                )));
            }
            Ok(zip(
                mask,
                &truthy.as_ref() as &dyn Datum,
                &falsy.as_ref() as &dyn Datum,
            )?)
        })
        .collect::<crate::Result<Vec<_>>>()?;

    let fields = base_fields
        .iter()
        .enumerate()
        .map(|(ix, field)| {
            let truthy = column_for_broadcast(then_, ix, mask.len())?;
            Ok(Field::new(field.name(), truthy.data_type().clone(), true))
        })
        .collect::<crate::Result<Vec<_>>>()?;

    Ok(RecordBatch::try_new(
        Arc::new(Schema::new(fields)),
        columns,
    )?)
}

fn column_for_broadcast(batch: &RecordBatch, ix: usize, len: usize) -> crate::Result<ArrayRef> {
    let column = batch
        .columns()
        .get(ix)
        .cloned()
        .or_else(|| batch.columns().first().cloned())
        .ok_or_else(|| {
            autark_error::Error::Reader("where requires at least one branch column".to_string())
        })?;

    if batch.num_rows() == len {
        return Ok(column);
    }
    if batch.num_rows() == 1 {
        let indices = UInt32Array::from(vec![0; len]);
        return Ok(take(column.as_ref(), &indices, None)?);
    }

    Err(autark_error::Error::Reader(format!(
        "where requires branch row count to be {} or 1, got {}",
        len,
        batch.num_rows()
    )))
}
