// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Implementation of the `uddsketch_state` UDAF that generate the state of
//! UDDSketch for a given set of values.
//!
//! The generated state can be used to compute approximate quantiles using
//! `uddsketch_calc` UDF.

use std::sync::Arc;

use common_query::prelude::*;
use datafusion::common::cast::{as_binary_array, as_primitive_array};
use datafusion::common::not_impl_err;
use datafusion::error::{DataFusionError, Result as DfResult};
use datafusion::logical_expr::function::AccumulatorArgs;
use datafusion::logical_expr::{Accumulator as DfAccumulator, AggregateUDF, Volatility};
use datafusion::physical_plan::expressions::Literal;
use datafusion::prelude::create_udaf;
use datatypes::arrow::array::{Array, ArrayRef};
use datatypes::arrow::datatypes::{DataType, Float64Type};
use uddsketch::{BatchWorkspace, UddSketch};

use crate::uddsketch_compat;

pub const UDDSKETCH_STATE_NAME: &str = "uddsketch_state";

pub const UDDSKETCH_MERGE_NAME: &str = "uddsketch_merge";

const MAX_BUCKETS: u32 = 1_000_000;

#[derive(Debug)]
pub struct UddSketchState {
    uddsketch: UddSketch,
    workspace: BatchWorkspace,
    values: Vec<f64>,
}

impl UddSketchState {
    pub fn new(bucket_size: u32, error_rate: f64) -> DfResult<Self> {
        if bucket_size > MAX_BUCKETS {
            return Err(DataFusionError::Plan(format!(
                "UDDSketch bucket size exceeds the maximum of {}",
                MAX_BUCKETS
            )));
        }
        let uddsketch = UddSketch::new(bucket_size, error_rate)
            .map_err(|e| DataFusionError::Plan(e.to_string()))?;
        Ok(Self {
            uddsketch,
            workspace: BatchWorkspace::default(),
            values: Vec::new(),
        })
    }

    pub fn state_udf_impl() -> AggregateUDF {
        create_udaf(
            UDDSKETCH_STATE_NAME,
            vec![DataType::Int64, DataType::Float64, DataType::Float64],
            Arc::new(DataType::Binary),
            Volatility::Immutable,
            Arc::new(|args| {
                let (bucket_size, error_rate) = downcast_accumulator_args(args)?;
                Ok(Box::new(UddSketchState::new(bucket_size, error_rate)?))
            }),
            Arc::new(vec![DataType::Binary]),
        )
    }

    /// Create a UDF for the `uddsketch_merge` function.
    ///
    /// `uddsketch_merge` accepts bucket size, error rate, and a binary column of states generated by `uddsketch_state`
    /// and merges them into a single state.
    ///
    /// The bucket size and error rate must be the same as the original state.
    pub fn merge_udf_impl() -> AggregateUDF {
        create_udaf(
            UDDSKETCH_MERGE_NAME,
            vec![DataType::Int64, DataType::Float64, DataType::Binary],
            Arc::new(DataType::Binary),
            Volatility::Immutable,
            Arc::new(|args| {
                let (bucket_size, error_rate) = downcast_accumulator_args(args)?;
                Ok(Box::new(UddSketchState::new(bucket_size, error_rate)?))
            }),
            Arc::new(vec![DataType::Binary]),
        )
    }

    fn merge(&mut self, raw: &[u8]) -> DfResult<()> {
        let uddsketch = uddsketch_compat::decode(raw).map_err(|e| {
            common_telemetry::trace!("Failed to deserialize UDDSketch: {}", e);
            DataFusionError::Plan("Failed to deserialize UDDSketch from binary".to_string())
        })?;
        if uddsketch.count() == 0 {
            return Ok(());
        }
        if self.uddsketch.max_buckets() != uddsketch.max_buckets()
            || self.uddsketch.initial_error().to_bits() != uddsketch.initial_error().to_bits()
        {
            return Err(DataFusionError::Plan(format!(
                "Merging UDDSketch with different parameters: arguments={:?} vs actual input={:?}",
                (self.uddsketch.max_buckets(), self.uddsketch.initial_error()),
                (uddsketch.max_buckets(), uddsketch.initial_error())
            )));
        }
        self.uddsketch
            .merge(&uddsketch)
            .map_err(|e| DataFusionError::Plan(e.to_string()))
    }
}

fn downcast_accumulator_args(args: AccumulatorArgs) -> DfResult<(u32, f64)> {
    let bucket_size = match args.exprs[0]
        .as_any()
        .downcast_ref::<Literal>()
        .map(|lit| lit.value())
    {
        Some(ScalarValue::Int64(Some(value))) => u32::try_from(*value).map_err(|_| {
            DataFusionError::Plan(format!("Invalid UDDSketch bucket size: {}", value))
        })?,
        _ => {
            return not_impl_err!(
                "{} not supported for bucket size: {}",
                UDDSKETCH_STATE_NAME,
                &args.exprs[0]
            );
        }
    };

    let error_rate = match args.exprs[1]
        .as_any()
        .downcast_ref::<Literal>()
        .map(|lit| lit.value())
    {
        Some(ScalarValue::Float64(Some(value))) => *value,
        _ => {
            return not_impl_err!(
                "{} not supported for error rate: {}",
                UDDSKETCH_STATE_NAME,
                &args.exprs[1]
            );
        }
    };

    Ok((bucket_size, error_rate))
}

impl DfAccumulator for UddSketchState {
    fn update_batch(&mut self, values: &[ArrayRef]) -> DfResult<()> {
        let array = &values[2]; // the third column is data value
        match array.data_type() {
            DataType::Float64 => {
                let f64_array = as_primitive_array::<Float64Type>(array)?;
                let values: &[f64] = if f64_array.null_count() == 0 {
                    f64_array.values().as_ref()
                } else {
                    self.values.clear();
                    self.values.extend(f64_array.iter().flatten());
                    self.values.as_slice()
                };
                self.uddsketch
                    .add_batch_with_workspace(values, &mut self.workspace)
                    .map_err(|e| DataFusionError::Execution(e.to_string()))?;
            }
            // meaning instantiate as `uddsketch_merge`
            DataType::Binary => self.merge_batch(std::slice::from_ref(array))?,
            _ => {
                return not_impl_err!(
                    "UDDSketch functions do not support data type: {}",
                    array.data_type()
                );
            }
        }

        Ok(())
    }

    fn evaluate(&mut self) -> DfResult<ScalarValue> {
        Ok(ScalarValue::Binary(Some(self.uddsketch.encode().map_err(
            |e| DataFusionError::Internal(format!("Failed to serialize UDDSketch: {}", e)),
        )?)))
    }

    fn size(&self) -> usize {
        std::mem::size_of::<Self>() - std::mem::size_of::<UddSketch>()
            + self.uddsketch.allocated_size()
            + self.workspace.allocated_size()
            + self.values.capacity() * std::mem::size_of::<f64>()
    }

    fn state(&mut self) -> DfResult<Vec<ScalarValue>> {
        Ok(vec![ScalarValue::Binary(Some(
            self.uddsketch.encode().map_err(|e| {
                DataFusionError::Internal(format!("Failed to serialize UDDSketch: {}", e))
            })?,
        ))])
    }

    fn merge_batch(&mut self, states: &[ArrayRef]) -> DfResult<()> {
        let array = &states[0];
        let binary_array = as_binary_array(array)?;
        for v in binary_array.iter().flatten() {
            self.merge(v)?;
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use datafusion::arrow::array::{BinaryArray, Float64Array};
    use uddsketch::UddSketchRef;

    use super::*;

    #[test]
    fn test_uddsketch_state_basic() {
        let mut state = UddSketchState::new(10, 0.01).unwrap();
        state.uddsketch.add(1.0).unwrap();
        state.uddsketch.add(2.0).unwrap();
        state.uddsketch.add(3.0).unwrap();

        let result = state.evaluate().unwrap();
        if let ScalarValue::Binary(Some(bytes)) = result {
            let encoded = UddSketchRef::parse(&bytes).unwrap();
            assert_eq!(encoded.count(), 3);
        } else {
            panic!("Expected binary scalar value");
        }
    }

    #[test]
    fn test_uddsketch_state_roundtrip() {
        let mut state = UddSketchState::new(10, 0.01).unwrap();
        state.uddsketch.add(1.0).unwrap();
        state.uddsketch.add(2.0).unwrap();

        // Serialize
        let serialized = state.evaluate().unwrap();

        // Create new state and merge the serialized data
        let mut new_state = UddSketchState::new(10, 0.01).unwrap();
        if let ScalarValue::Binary(Some(bytes)) = &serialized {
            new_state.merge(bytes).unwrap();

            let original_sketch = UddSketchRef::parse(bytes).unwrap();
            let new_result = new_state.evaluate().unwrap();
            if let ScalarValue::Binary(Some(new_bytes)) = new_result {
                let new_sketch = UddSketchRef::parse(&new_bytes).unwrap();
                assert_eq!(original_sketch.count(), new_sketch.count());
                assert_eq!(original_sketch.sum(), new_sketch.sum());
                assert_eq!(
                    original_sketch.max_error().unwrap(),
                    new_sketch.max_error().unwrap()
                );
                // Compare a few quantiles to ensure statistical equivalence
                for q in [0.1, 0.5, 0.9].iter() {
                    let original = original_sketch.quantile(*q).unwrap().unwrap();
                    let merged = new_sketch.quantile(*q).unwrap().unwrap();
                    assert!(
                        (original - merged).abs() < 1e-10,
                        "Quantile {} mismatch: original={}, new={}",
                        q,
                        original,
                        merged
                    );
                }
            } else {
                panic!("Expected binary scalar value");
            }
        } else {
            panic!("Expected binary scalar value");
        }
    }

    #[test]
    fn test_uddsketch_state_merges_legacy_state() {
        let mut state = UddSketchState::new(128, 0.01).unwrap();

        state.merge(uddsketch_compat::LEGACY_STATE).unwrap();

        let ScalarValue::Binary(Some(encoded)) = state.evaluate().unwrap() else {
            panic!("Expected binary scalar value");
        };
        let sketch = UddSketchRef::parse(&encoded).unwrap();
        assert_eq!(sketch.count(), 4);
        assert_eq!(sketch.sum(), 1.0);
        assert_eq!(sketch.quantile(0.5).unwrap(), Some(0.9900000000000001));
    }

    #[test]
    fn test_uddsketch_state_merges_compacted_legacy_state() {
        let mut legacy_state = uddsketch_compat::COMPACTED_LEGACY_SKETCH.to_vec();
        legacy_state.extend_from_slice(&0.01_f64.to_le_bytes());
        let mut state = UddSketchState::new(7, 0.01).unwrap();

        state.merge(&legacy_state).unwrap();

        let ScalarValue::Binary(Some(encoded)) = state.evaluate().unwrap() else {
            panic!("Expected binary scalar value");
        };
        let sketch = UddSketchRef::parse(&encoded).unwrap();
        assert_eq!(sketch.count(), 201);
        assert_eq!(sketch.times_compacted(), 12);
    }

    #[test]
    fn test_uddsketch_state_batch_update() {
        let mut state = UddSketchState::new(10, 0.01).unwrap();
        let values = vec![Some(1.0f64), None, Some(2.0), Some(3.0)];
        let array = Arc::new(Float64Array::from(values)) as ArrayRef;

        state
            .update_batch(&[array.clone(), array.clone(), array])
            .unwrap();

        let result = state.evaluate().unwrap();
        if let ScalarValue::Binary(Some(bytes)) = result {
            let encoded = UddSketchRef::parse(&bytes).unwrap();
            assert_eq!(encoded.count(), 3);
        } else {
            panic!("Expected binary scalar value");
        }
    }

    #[test]
    fn test_uddsketch_state_non_null_batch_avoids_values_buffer() {
        let mut state = UddSketchState::new(10, 0.01).unwrap();
        let array = Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0])) as ArrayRef;

        state
            .update_batch(&[array.clone(), array.clone(), array])
            .unwrap();

        assert_eq!(state.uddsketch.count(), 3);
        assert_eq!(state.values.capacity(), 0);
    }

    #[test]
    fn test_uddsketch_state_non_null_sliced_batch() {
        let mut state = UddSketchState::new(10, 0.01).unwrap();
        let array = Float64Array::from(vec![1.0, 2.0, 3.0, 4.0]);
        let array = array.slice(1, 2);
        let array = Arc::new(array) as ArrayRef;

        state
            .update_batch(&[array.clone(), array.clone(), array])
            .unwrap();

        assert_eq!(state.uddsketch.count(), 2);
        assert_eq!(state.uddsketch.sum(), 5.0);
    }

    #[test]
    fn test_uddsketch_state_merge_batch() {
        let mut state1 = UddSketchState::new(10, 0.01).unwrap();
        state1.uddsketch.add(1.0).unwrap();
        let state1_binary = state1.evaluate().unwrap();

        let mut state2 = UddSketchState::new(10, 0.01).unwrap();
        state2.uddsketch.add(2.0).unwrap();
        let state2_binary = state2.evaluate().unwrap();

        let mut merged_state = UddSketchState::new(10, 0.01).unwrap();
        if let (ScalarValue::Binary(Some(bytes1)), ScalarValue::Binary(Some(bytes2))) =
            (&state1_binary, &state2_binary)
        {
            let binary_array = Arc::new(BinaryArray::from(vec![
                bytes1.as_slice(),
                bytes2.as_slice(),
            ])) as ArrayRef;
            merged_state.merge_batch(&[binary_array]).unwrap();

            let result = merged_state.evaluate().unwrap();
            if let ScalarValue::Binary(Some(bytes)) = result {
                let encoded = UddSketchRef::parse(&bytes).unwrap();
                assert_eq!(encoded.count(), 2);
            } else {
                panic!("Expected binary scalar value");
            }
        } else {
            panic!("Expected binary scalar values");
        }
    }

    #[test]
    fn test_uddsketch_state_size() {
        let mut state = UddSketchState::new(10, 0.01).unwrap();
        let initial_size = state.size();

        // Add some values to create buckets
        let array = Arc::new(Float64Array::from_iter_values((0..64).map(f64::from))) as ArrayRef;
        state
            .update_batch(&[array.clone(), array.clone(), array])
            .unwrap();

        let size_with_values = state.size();
        assert!(
            size_with_values > initial_size,
            "Size should increase after adding values: initial={}, with_values={}",
            initial_size,
            size_with_values
        );
    }

    #[test]
    fn test_uddsketch_state_rejects_invalid_config() {
        assert!(UddSketchState::new(6, 0.01).is_err());
        assert!(UddSketchState::new(10, 1.0).is_err());

        let mut maximum = UddSketchState::new(1_000_000, 0.01).unwrap();
        let ScalarValue::Binary(Some(encoded)) = maximum.evaluate().unwrap() else {
            panic!("Expected binary scalar value");
        };
        UddSketchRef::parse(&encoded).unwrap();
        assert!(UddSketchState::new(1_000_001, 0.01).is_err());
    }

    #[test]
    fn test_uddsketch_state_rejects_nan_batch() {
        let mut state = UddSketchState::new(10, 0.01).unwrap();
        let array = Arc::new(Float64Array::from(vec![1.0, f64::NAN])) as ArrayRef;

        let error = state
            .update_batch(&[array.clone(), array.clone(), array])
            .unwrap_err();
        assert!(error.to_string().contains("NaN values are not supported"));
    }
}
