// 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 common_telemetry::trace;
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::ArrayRef;
use datatypes::arrow::datatypes::{DataType, Float64Type};
use serde::{Deserialize, Serialize};
use uddsketch::{SketchHashKey, UDDSketch};

pub const UDDSKETCH_STATE_NAME: &str = "uddsketch_state";

pub const UDDSKETCH_MERGE_NAME: &str = "uddsketch_merge";

#[derive(Debug, Serialize, Deserialize)]
pub struct UddSketchState {
    uddsketch: UDDSketch,
    error_rate: f64,
}

impl UddSketchState {
    pub fn new(bucket_size: u64, error_rate: f64) -> Self {
        Self {
            uddsketch: UDDSketch::new(bucket_size, error_rate),
            error_rate,
        }
    }

    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 update(&mut self, value: f64) {
        self.uddsketch.add_value(value);
    }

    fn merge(&mut self, raw: &[u8]) -> DfResult<()> {
        if let Ok(uddsketch) = bincode::deserialize::<Self>(raw) {
            if uddsketch.uddsketch.count() != 0 {
                if self.uddsketch.max_allowed_buckets() != uddsketch.uddsketch.max_allowed_buckets()
                    || (self.error_rate - uddsketch.error_rate).abs() >= 1e-9
                {
                    return Err(DataFusionError::Plan(format!(
                        "Merging UDDSketch with different parameters: arguments={:?} vs actual input={:?}",
                        (self.uddsketch.max_allowed_buckets(), self.error_rate),
                        (
                            uddsketch.uddsketch.max_allowed_buckets(),
                            uddsketch.error_rate
                        )
                    )));
                }
                self.uddsketch.merge_sketch(&uddsketch.uddsketch);
            }
        } else {
            trace!("Warning: Failed to deserialize UDDSketch from {:?}", raw);
            return Err(DataFusionError::Plan(
                "Failed to deserialize UDDSketch from binary".to_string(),
            ));
        }

        Ok(())
    }
}

fn downcast_accumulator_args(args: AccumulatorArgs) -> DfResult<(u64, f64)> {
    let bucket_size = match args.exprs[0]
        .as_any()
        .downcast_ref::<Literal>()
        .map(|lit| lit.value())
    {
        Some(ScalarValue::Int64(Some(value))) => *value as u64,
        _ => {
            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)?;
                for v in f64_array.iter().flatten() {
                    self.update(v);
                }
            }
            // 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(
            bincode::serialize(&self).map_err(|e| {
                DataFusionError::Internal(format!("Failed to serialize UDDSketch: {}", e))
            })?,
        )))
    }

    fn size(&self) -> usize {
        // Base size of UDDSketch struct fields
        let mut total_size = std::mem::size_of::<f64>() * 3 + // alpha, gamma, values_sum
                            std::mem::size_of::<u32>() +      // compactions
                            std::mem::size_of::<u64>() * 2; // max_buckets, num_values

        // Size of buckets (SketchHashMap)
        // Each bucket entry contains:
        // - SketchHashKey (enum with i64/Zero/Invalid variants)
        // - SketchHashEntry (count: u64, next: SketchHashKey)
        let bucket_entry_size = std::mem::size_of::<SketchHashKey>() + // key
                               std::mem::size_of::<u64>() +            // count
                               std::mem::size_of::<SketchHashKey>(); // next

        total_size += self.uddsketch.current_buckets_count() * bucket_entry_size;

        total_size
    }

    fn state(&mut self) -> DfResult<Vec<ScalarValue>> {
        Ok(vec![ScalarValue::Binary(Some(
            bincode::serialize(&self).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 super::*;

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

        let result = state.evaluate().unwrap();
        if let ScalarValue::Binary(Some(bytes)) = result {
            let deserialized: UddSketchState = bincode::deserialize(&bytes).unwrap();
            assert_eq!(deserialized.uddsketch.count(), 3);
        } else {
            panic!("Expected binary scalar value");
        }
    }

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

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

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

            // Verify the merged state matches original by comparing deserialized values
            let original_sketch: UddSketchState = bincode::deserialize(bytes).unwrap();
            let original_sketch = original_sketch.uddsketch;
            let new_result = new_state.evaluate().unwrap();
            if let ScalarValue::Binary(Some(new_bytes)) = new_result {
                let new_sketch: UddSketchState = bincode::deserialize(&new_bytes).unwrap();
                let new_sketch = new_sketch.uddsketch;
                assert_eq!(original_sketch.count(), new_sketch.count());
                assert_eq!(original_sketch.sum(), new_sketch.sum());
                assert_eq!(original_sketch.mean(), new_sketch.mean());
                assert_eq!(original_sketch.max_error(), new_sketch.max_error());
                // Compare a few quantiles to ensure statistical equivalence
                for q in [0.1, 0.5, 0.9].iter() {
                    assert!(
                        (original_sketch.estimate_quantile(*q) - new_sketch.estimate_quantile(*q))
                            .abs()
                            < 1e-10,
                        "Quantile {} mismatch: original={}, new={}",
                        q,
                        original_sketch.estimate_quantile(*q),
                        new_sketch.estimate_quantile(*q)
                    );
                }
            } else {
                panic!("Expected binary scalar value");
            }
        } else {
            panic!("Expected binary scalar value");
        }
    }

    #[test]
    fn test_uddsketch_state_batch_update() {
        let mut state = UddSketchState::new(10, 0.01);
        let values = vec![1.0f64, 2.0, 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 deserialized: UddSketchState = bincode::deserialize(&bytes).unwrap();
            let deserialized = deserialized.uddsketch;
            assert_eq!(deserialized.count(), 3);
        } else {
            panic!("Expected binary scalar value");
        }
    }

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

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

        let mut merged_state = UddSketchState::new(10, 0.01);
        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 deserialized: UddSketchState = bincode::deserialize(&bytes).unwrap();
                let deserialized = deserialized.uddsketch;
                assert_eq!(deserialized.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);
        let initial_size = state.size();

        // Add some values to create buckets
        state.update(1.0);
        state.update(2.0);
        state.update(3.0);

        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
        );

        // Verify size increases with more buckets
        state.update(10.0); // This should create a new bucket
        assert!(
            state.size() > size_with_values,
            "Size should increase after adding new bucket: prev={}, new={}",
            size_with_values,
            state.size()
        );
    }
}
