//! Boolean aggregate UDAFs with sliding-window retraction support.
//!
//! Hamelin maps `any(boolean)` and `all(boolean)` through these UDAFs instead
//! of DataFusion's built-in boolean aggregates because Hamelin `WINDOW ... WITHIN`
//! can produce bounded sliding frames. DataFusion evaluates those frames by
//! retracting rows as they leave the frame, so the aggregate accumulator must
//! implement `retract_batch`. The built-ins are fine for plain aggregation, but
//! do not currently support that sliding-window contract.
//!
//! What comes from DataFusion is the UDAF/Accumulator interface and lifecycle:
//! `AggregateUDFImpl`, `Accumulator`, `update_batch`, `merge_batch`,
//! `retract_batch`, `state`, and `evaluate`. The implementation here is
//! Hamelin-owned rather than copied from upstream: we keep `true_count` and
//! `false_count` state so both `any` and `all` can add rows, retract rows, merge
//! partial states, and preserve the same null behavior as the standard boolean
//! aggregates.

use std::any::Any;
use std::mem::size_of_val;

use datafusion::arrow::array::ArrayRef;
use datafusion::arrow::datatypes::{DataType, Field};
use datafusion::common::{exec_err, DataFusionError, Result, ScalarValue};
use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
use datafusion::logical_expr::{
    Accumulator, AggregateUDF, AggregateUDFImpl, Signature, Volatility,
};

pub fn bool_or_udaf() -> AggregateUDF {
    AggregateUDF::new_from_impl(BoolAggUdaf::new(BoolAggKind::Any))
}

pub fn bool_and_udaf() -> AggregateUDF {
    AggregateUDF::new_from_impl(BoolAggUdaf::new(BoolAggKind::All))
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
enum BoolAggKind {
    Any,
    All,
}

#[derive(Debug, PartialEq, Eq, Hash)]
struct BoolAggUdaf {
    kind: BoolAggKind,
    signature: Signature,
}

impl BoolAggUdaf {
    fn new(kind: BoolAggKind) -> Self {
        Self {
            kind,
            signature: Signature::exact(vec![DataType::Boolean], Volatility::Immutable),
        }
    }
}

impl AggregateUDFImpl for BoolAggUdaf {
    fn as_any(&self) -> &dyn Any {
        self
    }

    fn name(&self) -> &str {
        match self.kind {
            BoolAggKind::Any => "hamelin_bool_or",
            BoolAggKind::All => "hamelin_bool_and",
        }
    }

    fn signature(&self) -> &Signature {
        &self.signature
    }

    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
        Ok(DataType::Boolean)
    }

    fn state_fields(&self, _args: StateFieldsArgs) -> Result<Vec<std::sync::Arc<Field>>> {
        Ok(vec![
            std::sync::Arc::new(Field::new("true_count", DataType::UInt64, false)),
            std::sync::Arc::new(Field::new("false_count", DataType::UInt64, false)),
        ])
    }

    fn accumulator(&self, acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
        let input_type = acc_args.exprs[0]
            .data_type(acc_args.schema)
            .map_err(|e| DataFusionError::External(Box::new(e)))?;
        if input_type != DataType::Boolean {
            return exec_err!("boolean aggregate expected Boolean input, got {input_type}");
        }

        Ok(Box::new(BoolAggAccumulator::new(self.kind)))
    }
}

#[derive(Debug)]
struct BoolAggAccumulator {
    kind: BoolAggKind,
    true_count: u64,
    false_count: u64,
}

impl BoolAggAccumulator {
    fn new(kind: BoolAggKind) -> Self {
        Self {
            kind,
            true_count: 0,
            false_count: 0,
        }
    }

    fn add_value(&mut self, value: Option<bool>) {
        match value {
            Some(true) => self.true_count += 1,
            Some(false) => self.false_count += 1,
            None => {}
        }
    }

    fn remove_value(&mut self, value: Option<bool>) -> Result<()> {
        match value {
            Some(true) if self.true_count > 0 => self.true_count -= 1,
            Some(false) if self.false_count > 0 => self.false_count -= 1,
            Some(true) => return exec_err!("boolean aggregate retract received unseen true value"),
            Some(false) => {
                return exec_err!("boolean aggregate retract received unseen false value")
            }
            None => {}
        }
        Ok(())
    }

    fn update_counts(&mut self, values: &[ArrayRef], retract: bool) -> Result<()> {
        if values.is_empty() {
            return Ok(());
        }

        let values_array = &values[0];
        for row in 0..values_array.len() {
            let value = match ScalarValue::try_from_array(values_array, row)? {
                ScalarValue::Boolean(value) => value,
                other => {
                    return exec_err!("boolean aggregate expected Boolean input, got {other}");
                }
            };

            if retract {
                self.remove_value(value)?;
            } else {
                self.add_value(value);
            }
        }

        Ok(())
    }
}

impl Accumulator for BoolAggAccumulator {
    fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
        self.update_counts(values, false)
    }

    fn retract_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
        self.update_counts(values, true)
    }

    fn supports_retract_batch(&self) -> bool {
        true
    }

    fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
        if states.len() != 2 {
            return exec_err!(
                "boolean aggregate merge expects 2 state arrays, got {}",
                states.len()
            );
        }

        for row in 0..states[0].len() {
            match ScalarValue::try_from_array(&states[0], row)? {
                ScalarValue::UInt64(Some(count)) => self.true_count += count,
                ScalarValue::UInt64(None) => {}
                other => {
                    return exec_err!("boolean aggregate expected true_count UInt64, got {other}")
                }
            }
            match ScalarValue::try_from_array(&states[1], row)? {
                ScalarValue::UInt64(Some(count)) => self.false_count += count,
                ScalarValue::UInt64(None) => {}
                other => {
                    return exec_err!("boolean aggregate expected false_count UInt64, got {other}");
                }
            }
        }

        Ok(())
    }

    fn state(&mut self) -> Result<Vec<ScalarValue>> {
        Ok(vec![
            ScalarValue::UInt64(Some(self.true_count)),
            ScalarValue::UInt64(Some(self.false_count)),
        ])
    }

    fn evaluate(&mut self) -> Result<ScalarValue> {
        let value = match self.kind {
            BoolAggKind::Any if self.true_count > 0 => Some(true),
            BoolAggKind::Any if self.false_count > 0 => Some(false),
            BoolAggKind::All if self.false_count > 0 => Some(false),
            BoolAggKind::All if self.true_count > 0 => Some(true),
            _ => None,
        };
        Ok(ScalarValue::Boolean(value))
    }

    fn size(&self) -> usize {
        size_of_val(self)
    }
}
