use std::sync::Arc;

use arrow::{
    array::{ArrayRef, RecordBatch},
    datatypes::{Field, Schema},
};
use mpera_frontend::op::{BinaryOpKind, OpRef};

use crate::{
    Result,
    luminal_interop::binary_array,
    op::{OpOut, StatefulBatchOp},
    runtimedb::RuntimeDB,
    typecheck::binary_output_dtype,
};

pub struct BinaryOpBatchOp {
    pub opref: OpRef,
    pub lhs: OpRef,
    pub rhs: OpRef,
    pub kind: BinaryOpKind,
}

impl BinaryOpBatchOp {
    pub fn new(opref: OpRef, lhs: OpRef, rhs: OpRef, kind: BinaryOpKind) -> Self {
        Self {
            opref,
            lhs,
            rhs,
            kind,
        }
    }
}

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

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

        Ok(Some(vec![apply_binary(&self.kind, lhs, rhs)?]))
    }
}

fn apply_binary(kind: &BinaryOpKind, lhs: &RecordBatch, rhs: &RecordBatch) -> Result<RecordBatch> {
    let lhs_schema = lhs.schema();
    let rhs_schema = rhs.schema();
    let lhs_fields = lhs_schema.fields();
    let rhs_fields = rhs_schema.fields();
    let base_fields = if lhs_fields.len() >= rhs_fields.len() {
        lhs_fields.clone()
    } else {
        rhs_fields.clone()
    };

    let columns = (0..base_fields.len())
        .map(|ix| {
            let lhs = column_for_binary(lhs, ix)?;
            let rhs = column_for_binary(rhs, ix)?;
            binary_array(kind, lhs.as_ref(), rhs.as_ref())
        })
        .collect::<Result<Vec<_>>>()?;

    let fields = base_fields
        .iter()
        .enumerate()
        .map(|(ix, field)| {
            let lhs_dtype = column_for_binary(lhs, ix)?.data_type().clone();
            let rhs_dtype = column_for_binary(rhs, ix)?.data_type().clone();
            Ok(Field::new(
                field.name(),
                binary_output_dtype(kind, &lhs_dtype, &rhs_dtype)?,
                true,
            ))
        })
        .collect::<Result<Vec<_>>>()?;

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

fn column_for_binary(batch: &RecordBatch, ix: usize) -> Result<ArrayRef> {
    batch
        .columns()
        .get(ix)
        .cloned()
        .or_else(|| batch.columns().first().cloned())
        .ok_or_else(|| {
            autark_error::Error::Reader("binary op requires at least one input column".to_string())
        })
}
