use std::sync::Arc;

use arrow::{
    array::{Float64Array, RecordBatch, StringArray},
    datatypes::{DataType, Field, Schema},
};
use mpera_frontend::op::{Literal, OpRef};

use crate::{
    Result,
    op::{OpOut, StatefulBatchOp},
};

pub struct ConstantBatchOp {
    pub opref: OpRef,
    pub value: Literal,
    emitted: bool,
}

impl ConstantBatchOp {
    pub fn new(opref: OpRef, value: Literal) -> Self {
        Self {
            opref,
            value,
            emitted: false,
        }
    }
}

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

    fn consume(&mut self, _rtdb: &crate::runtimedb::RuntimeDB) -> OpOut {
        if self.emitted {
            return Ok(None);
        }

        self.emitted = true;
        Ok(Some(vec![literal_batch(&self.value)?]))
    }
}

fn literal_batch(value: &Literal) -> Result<RecordBatch> {
    match value {
        Literal::F64(value) => Ok(RecordBatch::try_new(
            Arc::new(Schema::new(vec![Field::new(
                "literal",
                DataType::Float64,
                true,
            )])),
            vec![Arc::new(Float64Array::from(vec![Some(*value)]))],
        )?),
        Literal::String(value) => Ok(RecordBatch::try_new(
            Arc::new(Schema::new(vec![Field::new(
                "literal",
                DataType::Utf8,
                true,
            )])),
            vec![Arc::new(StringArray::from(vec![Some(value.as_str())]))],
        )?),
    }
}
