//! Distinct array aggregate UDAF for DataFusion.
//!
//! Implements `set_agg`, keeping duplicate elimination inside the
//! aggregate state instead of materializing a full `array_agg` result first.

use std::any::Any;
use std::mem::{size_of, size_of_val};
use std::sync::Arc;

use datafusion::arrow::array::{new_empty_array, Array, ArrayRef, AsArray, ListArray, StructArray};
use datafusion::arrow::buffer::OffsetBuffer;
use datafusion::arrow::datatypes::{DataType, Field, Fields};
use datafusion::common::{exec_err, DataFusionError, Result, ScalarValue};
use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
use datafusion::logical_expr::{
    Accumulator, AggregateUDF, AggregateUDFImpl, Signature, Volatility,
};
use vecmap::map::Entry;
use vecmap::VecMap;

pub fn set_agg_udaf() -> AggregateUDF {
    AggregateUDF::new_from_impl(SetAggUdaf::new())
}

#[derive(Debug, PartialEq, Eq, Hash)]
pub struct SetAggUdaf {
    signature: Signature,
}

impl Default for SetAggUdaf {
    fn default() -> Self {
        Self::new()
    }
}

impl SetAggUdaf {
    pub fn new() -> Self {
        Self {
            signature: Signature::any(1, Volatility::Immutable),
        }
    }
}

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

    fn name(&self) -> &str {
        "hamelin_set_agg"
    }

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

    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
        Ok(DataType::List(Arc::new(Field::new(
            "item",
            arg_types[0].clone(),
            true,
        ))))
    }

    fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<Arc<Field>>> {
        let element_type = args.input_fields[0].data_type().clone();
        let entry_type = DataType::Struct(entry_fields(element_type));

        Ok(vec![Arc::new(Field::new_list(
            "entries",
            Field::new("item", entry_type, false),
            true,
        ))])
    }

    fn accumulator(&self, acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
        let element_type = acc_args.exprs[0]
            .data_type(acc_args.schema)
            .map_err(|e| DataFusionError::External(Box::new(e)))?;

        Ok(Box::new(SetAggAccumulator::new(element_type)))
    }
}

fn entry_fields(element_type: DataType) -> Fields {
    Fields::from(vec![
        Field::new("value", element_type, true),
        Field::new("count", DataType::UInt64, false),
    ])
}

#[derive(Debug)]
struct SetAggAccumulator {
    values: VecMap<ScalarValue, usize>,
    element_type: DataType,
}

impl SetAggAccumulator {
    fn new(element_type: DataType) -> Self {
        Self {
            values: VecMap::new(),
            element_type,
        }
    }

    fn add_value(&mut self, mut value: ScalarValue, count: usize) {
        if count == 0 {
            return;
        }

        value.compact();

        match self.values.entry(value) {
            Entry::Occupied(mut entry) => {
                *entry.get_mut() += count;
            }
            Entry::Vacant(entry) => {
                entry.insert(count);
            }
        }
    }

    fn remove_value(&mut self, mut value: ScalarValue) -> Result<()> {
        value.compact();

        match self.values.entry(value) {
            Entry::Occupied(mut entry) => {
                let count = entry.get_mut();
                if *count == 1 {
                    entry.remove();
                } else {
                    *count -= 1;
                }
                Ok(())
            }
            Entry::Vacant(_) => exec_err!("set_agg retract received unseen value"),
        }
    }

    fn build_values_list(&self) -> Result<ScalarValue> {
        let field = Arc::new(Field::new("item", self.element_type.clone(), true));
        let values_array = if self.values.len() == 0 {
            new_empty_array(&self.element_type)
        } else {
            ScalarValue::iter_to_array(self.values.iter().map(|(value, _)| value.clone()))?
        };

        Ok(ScalarValue::List(Arc::new(ListArray::new(
            field,
            OffsetBuffer::from_lengths([values_array.len()]),
            values_array,
            None,
        ))))
    }

    fn build_entries_list(&self) -> Result<ScalarValue> {
        let fields = entry_fields(self.element_type.clone());
        let values_array = if self.values.len() == 0 {
            new_empty_array(&self.element_type)
        } else {
            ScalarValue::iter_to_array(self.values.iter().map(|(value, _)| value.clone()))?
        };

        let counts = self
            .values
            .iter()
            .map(|(_, count)| ScalarValue::UInt64(Some(*count as u64)))
            .collect::<Vec<_>>();

        let counts_array = if counts.is_empty() {
            new_empty_array(&DataType::UInt64)
        } else {
            ScalarValue::iter_to_array(counts.into_iter())?
        };

        let entries_array = Arc::new(StructArray::try_new(
            fields.clone(),
            vec![values_array, counts_array],
            None,
        )?);
        let field = Arc::new(Field::new("item", DataType::Struct(fields), false));

        Ok(ScalarValue::List(Arc::new(ListArray::new(
            field,
            OffsetBuffer::from_lengths([entries_array.len()]),
            entries_array,
            None,
        ))))
    }
}

impl Accumulator for SetAggAccumulator {
    fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
        if values.is_empty() {
            return Ok(());
        }

        let values_array = &values[0];
        for row in 0..values_array.len() {
            self.add_value(ScalarValue::try_from_array(values_array, row)?, 1);
        }

        Ok(())
    }

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

        let values_array = &values[0];
        for row in 0..values_array.len() {
            self.remove_value(ScalarValue::try_from_array(values_array, row)?)?;
        }

        Ok(())
    }

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

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

        let entries_list = states[0].as_list::<i32>();

        for row in 0..entries_list.len() {
            if entries_list.is_null(row) {
                continue;
            }

            let entries_array = entries_list.value(row);
            let Some(entries) = entries_array.as_any().downcast_ref::<StructArray>() else {
                return exec_err!(
                    "set_agg expected StructArray state entries, got {}",
                    entries_array.data_type()
                );
            };
            let Some(values_array) = entries.column_by_name("value") else {
                return exec_err!("set_agg state entries are missing value field");
            };
            let Some(counts_array) = entries.column_by_name("count") else {
                return exec_err!("set_agg state entries are missing count field");
            };

            for value_index in 0..entries.len() {
                if entries.is_null(value_index) {
                    continue;
                }

                let count = match ScalarValue::try_from_array(counts_array, value_index)? {
                    ScalarValue::UInt64(Some(count)) => count as usize,
                    ScalarValue::UInt64(None) => continue,
                    other => {
                        return exec_err!(
                            "set_agg expected UInt64 count, got {}",
                            other.data_type()
                        );
                    }
                };
                let value = ScalarValue::try_from_array(values_array, value_index)?;
                self.add_value(value, count);
            }
        }

        Ok(())
    }

    fn state(&mut self) -> Result<Vec<ScalarValue>> {
        Ok(vec![self.build_entries_list()?])
    }

    fn evaluate(&mut self) -> Result<ScalarValue> {
        self.build_values_list()
    }

    fn size(&self) -> usize {
        size_of_val(self)
            + self.values.capacity() * (size_of::<ScalarValue>() + size_of::<usize>())
            + self
                .values
                .iter()
                .map(|(value, _)| value.size() - size_of_val(value))
                .sum::<usize>()
    }
}
