//! Aggregate functions.

use std::collections::HashSet;

use crate::error::{ExecutionError, ExecutionResult};
use crate::result::CypherValue;

/// Accumulator for aggregate functions.
#[derive(Debug, Clone)]
pub enum AggregateAccumulator {
    Count { count: i64, distinct_values: Option<HashSet<String>> },
    Sum { sum: f64, is_int: bool, has_value: bool },
    Avg { sum: f64, count: i64 },
    Min { min: Option<CypherValue> },
    Max { max: Option<CypherValue> },
    Collect { values: Vec<CypherValue>, distinct: bool, seen: HashSet<String> },
    StdDev { values: Vec<f64>, population: bool },
    PercentileCont { values: Vec<f64>, percentile: f64 },
    PercentileDisc { values: Vec<f64>, percentile: f64 },
}

impl AggregateAccumulator {
    /// Create a new COUNT accumulator.
    pub fn count(distinct: bool) -> Self {
        AggregateAccumulator::Count {
            count: 0,
            distinct_values: if distinct { Some(HashSet::new()) } else { None },
        }
    }

    /// Create a new SUM accumulator.
    pub fn sum() -> Self {
        AggregateAccumulator::Sum {
            sum: 0.0,
            is_int: true,
            has_value: false,
        }
    }

    /// Create a new AVG accumulator.
    pub fn avg() -> Self {
        AggregateAccumulator::Avg { sum: 0.0, count: 0 }
    }

    /// Create a new MIN accumulator.
    pub fn min() -> Self {
        AggregateAccumulator::Min { min: None }
    }

    /// Create a new MAX accumulator.
    pub fn max() -> Self {
        AggregateAccumulator::Max { max: None }
    }

    /// Create a new COLLECT accumulator.
    pub fn collect(distinct: bool) -> Self {
        AggregateAccumulator::Collect {
            values: Vec::new(),
            distinct,
            seen: HashSet::new(),
        }
    }

    /// Create a new STDEV accumulator.
    pub fn stdev(population: bool) -> Self {
        AggregateAccumulator::StdDev {
            values: Vec::new(),
            population,
        }
    }

    /// Create a new percentileCont accumulator.
    pub fn percentile_cont(percentile: f64) -> Self {
        AggregateAccumulator::PercentileCont {
            values: Vec::new(),
            percentile,
        }
    }

    /// Create a new percentileDisc accumulator.
    pub fn percentile_disc(percentile: f64) -> Self {
        AggregateAccumulator::PercentileDisc {
            values: Vec::new(),
            percentile,
        }
    }

    /// Accumulate a value.
    pub fn accumulate(&mut self, value: CypherValue) -> ExecutionResult<()> {
        match self {
            AggregateAccumulator::Count { count, distinct_values } => {
                if !value.is_null() {
                    if let Some(seen) = distinct_values {
                        let key = format!("{:?}", value);
                        if seen.insert(key) {
                            *count += 1;
                        }
                    } else {
                        *count += 1;
                    }
                }
            }
            AggregateAccumulator::Sum { sum, is_int, has_value } => {
                if !value.is_null() {
                    match value {
                        CypherValue::Integer(i) => {
                            *sum += i as f64;
                            *has_value = true;
                        }
                        CypherValue::Float(f) => {
                            *sum += f;
                            *is_int = false;
                            *has_value = true;
                        }
                        _ => {
                            return Err(ExecutionError::Type(format!(
                                "sum() requires numeric values, got {}",
                                value.type_name()
                            )))
                        }
                    }
                }
            }
            AggregateAccumulator::Avg { sum, count } => {
                if !value.is_null() {
                    match value {
                        CypherValue::Integer(i) => {
                            *sum += i as f64;
                            *count += 1;
                        }
                        CypherValue::Float(f) => {
                            *sum += f;
                            *count += 1;
                        }
                        _ => {
                            return Err(ExecutionError::Type(format!(
                                "avg() requires numeric values, got {}",
                                value.type_name()
                            )))
                        }
                    }
                }
            }
            AggregateAccumulator::Min { min } => {
                if !value.is_null() {
                    match min {
                        None => *min = Some(value),
                        Some(current) => {
                            // Use order_compare for proper type ordering (Lists < Strings < Numbers)
                            if value.order_compare(current) == std::cmp::Ordering::Less {
                                *min = Some(value);
                            }
                        }
                    }
                }
            }
            AggregateAccumulator::Max { max } => {
                if !value.is_null() {
                    match max {
                        None => *max = Some(value),
                        Some(current) => {
                            // Use order_compare for proper type ordering (Lists < Strings < Numbers)
                            if value.order_compare(current) == std::cmp::Ordering::Greater {
                                *max = Some(value);
                            }
                        }
                    }
                }
            }
            AggregateAccumulator::Collect { values, distinct, seen } => {
                if !value.is_null() {
                    if *distinct {
                        let key = format!("{:?}", value);
                        if seen.insert(key) {
                            values.push(value);
                        }
                    } else {
                        values.push(value);
                    }
                }
            }
            AggregateAccumulator::StdDev { values, .. } => {
                if !value.is_null() {
                    match value {
                        CypherValue::Integer(i) => values.push(i as f64),
                        CypherValue::Float(f) => values.push(f),
                        _ => {
                            return Err(ExecutionError::Type(format!(
                                "stdev() requires numeric values, got {}",
                                value.type_name()
                            )))
                        }
                    }
                }
            }
            AggregateAccumulator::PercentileCont { values, .. }
            | AggregateAccumulator::PercentileDisc { values, .. } => {
                if !value.is_null() {
                    match value {
                        CypherValue::Integer(i) => values.push(i as f64),
                        CypherValue::Float(f) => values.push(f),
                        _ => {
                            return Err(ExecutionError::Type(format!(
                                "percentile functions require numeric values, got {}",
                                value.type_name()
                            )))
                        }
                    }
                }
            }
        }
        Ok(())
    }

    /// Finalize and return the aggregate result.
    pub fn finalize(self) -> CypherValue {
        match self {
            AggregateAccumulator::Count { count, .. } => CypherValue::Integer(count),
            AggregateAccumulator::Sum { sum, is_int, has_value } => {
                if !has_value {
                    CypherValue::Integer(0)
                } else if is_int && sum.fract() == 0.0 {
                    CypherValue::Integer(sum as i64)
                } else {
                    CypherValue::Float(sum)
                }
            }
            AggregateAccumulator::Avg { sum, count } => {
                if count == 0 {
                    CypherValue::Null
                } else {
                    CypherValue::Float(sum / count as f64)
                }
            }
            AggregateAccumulator::Min { min } => min.unwrap_or(CypherValue::Null),
            AggregateAccumulator::Max { max } => max.unwrap_or(CypherValue::Null),
            AggregateAccumulator::Collect { values, .. } => CypherValue::List(values),
            AggregateAccumulator::StdDev { values, population } => {
                if values.is_empty() {
                    return CypherValue::Null;
                }
                let n = values.len() as f64;
                let mean = values.iter().sum::<f64>() / n;
                let variance: f64 = values.iter().map(|x| (x - mean).powi(2)).sum::<f64>();
                let divisor = if population { n } else { n - 1.0 };
                if divisor <= 0.0 {
                    return CypherValue::Null;
                }
                CypherValue::Float((variance / divisor).sqrt())
            }
            AggregateAccumulator::PercentileCont { mut values, percentile } => {
                if values.is_empty() {
                    return CypherValue::Null;
                }
                values.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
                let n = values.len();
                let index = percentile * (n - 1) as f64;
                let lower = index.floor() as usize;
                let upper = index.ceil() as usize;
                let frac = index - lower as f64;
                if lower == upper || upper >= n {
                    CypherValue::Float(values[lower])
                } else {
                    CypherValue::Float(values[lower] * (1.0 - frac) + values[upper] * frac)
                }
            }
            AggregateAccumulator::PercentileDisc { mut values, percentile } => {
                if values.is_empty() {
                    return CypherValue::Null;
                }
                values.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
                let n = values.len();
                let index = (percentile * (n - 1) as f64).round() as usize;
                CypherValue::Float(values[index.min(n - 1)])
            }
        }
    }
}

/// Create an aggregate accumulator from a function name.
pub fn create_accumulator(name: &str, distinct: bool, args: &[CypherValue]) -> ExecutionResult<AggregateAccumulator> {
    match name.to_lowercase().as_str() {
        "count" => Ok(AggregateAccumulator::count(distinct)),
        "sum" => Ok(AggregateAccumulator::sum()),
        "avg" => Ok(AggregateAccumulator::avg()),
        "min" => Ok(AggregateAccumulator::min()),
        "max" => Ok(AggregateAccumulator::max()),
        "collect" => Ok(AggregateAccumulator::collect(distinct)),
        "stdev" => Ok(AggregateAccumulator::stdev(false)),
        "stdevp" => Ok(AggregateAccumulator::stdev(true)),
        "percentilecont" => {
            let percentile = args.first()
                .and_then(|v| v.as_float())
                .ok_or_else(|| ExecutionError::InvalidArgument(
                    "percentileCont requires a percentile argument".to_string()
                ))?;
            Ok(AggregateAccumulator::percentile_cont(percentile))
        }
        "percentiledisc" => {
            let percentile = args.first()
                .and_then(|v| v.as_float())
                .ok_or_else(|| ExecutionError::InvalidArgument(
                    "percentileDisc requires a percentile argument".to_string()
                ))?;
            Ok(AggregateAccumulator::percentile_disc(percentile))
        }
        _ => Err(ExecutionError::UnknownFunction(name.to_string())),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_count() {
        let mut acc = AggregateAccumulator::count(false);
        acc.accumulate(CypherValue::Integer(1)).unwrap();
        acc.accumulate(CypherValue::Integer(2)).unwrap();
        acc.accumulate(CypherValue::Null).unwrap();
        assert_eq!(acc.finalize(), CypherValue::Integer(2));
    }

    #[test]
    fn test_count_distinct() {
        let mut acc = AggregateAccumulator::count(true);
        acc.accumulate(CypherValue::Integer(1)).unwrap();
        acc.accumulate(CypherValue::Integer(1)).unwrap();
        acc.accumulate(CypherValue::Integer(2)).unwrap();
        assert_eq!(acc.finalize(), CypherValue::Integer(2));
    }

    #[test]
    fn test_sum() {
        let mut acc = AggregateAccumulator::sum();
        acc.accumulate(CypherValue::Integer(1)).unwrap();
        acc.accumulate(CypherValue::Integer(2)).unwrap();
        acc.accumulate(CypherValue::Integer(3)).unwrap();
        assert_eq!(acc.finalize(), CypherValue::Integer(6));
    }

    #[test]
    fn test_avg() {
        let mut acc = AggregateAccumulator::avg();
        acc.accumulate(CypherValue::Integer(1)).unwrap();
        acc.accumulate(CypherValue::Integer(2)).unwrap();
        acc.accumulate(CypherValue::Integer(3)).unwrap();
        assert_eq!(acc.finalize(), CypherValue::Float(2.0));
    }

    #[test]
    fn test_min_max() {
        let mut min_acc = AggregateAccumulator::min();
        let mut max_acc = AggregateAccumulator::max();

        for v in [3, 1, 4, 1, 5] {
            min_acc.accumulate(CypherValue::Integer(v)).unwrap();
            max_acc.accumulate(CypherValue::Integer(v)).unwrap();
        }

        assert_eq!(min_acc.finalize(), CypherValue::Integer(1));
        assert_eq!(max_acc.finalize(), CypherValue::Integer(5));
    }

    #[test]
    fn test_collect() {
        let mut acc = AggregateAccumulator::collect(false);
        acc.accumulate(CypherValue::Integer(1)).unwrap();
        acc.accumulate(CypherValue::Integer(2)).unwrap();
        acc.accumulate(CypherValue::Integer(3)).unwrap();
        assert_eq!(
            acc.finalize(),
            CypherValue::List(vec![
                CypherValue::Integer(1),
                CypherValue::Integer(2),
                CypherValue::Integer(3),
            ])
        );
    }
}
