//! Metric aggregation implementations: avg, sum, min, max, value_count, stats.
//!
//! See [[aggregations]] and [[feature-aggregations-v010#Step 5]].

use luci_core::DocId;

use super::{AggCollector, AggCollectorFactory, AggResult, MetricResult};
use crate::segment::reader::SegmentReader;

/// Factory for creating metric collectors.
pub struct MetricAggFactory {
    pub field_name: String,
    pub metric_type: MetricType,
}

#[derive(Clone, Copy)]
pub enum MetricType {
    Avg,
    Sum,
    Min,
    Max,
    ValueCount,
    Stats,
    ExtendedStats,
}

impl AggCollectorFactory for MetricAggFactory {
    fn create_collector(&self, reader: &SegmentReader) -> Box<dyn AggCollector> {
        let field_id = reader.header().fields.iter()
            .find(|f| f.field_name == self.field_name)
            .map(|f| f.field_id);

        if let Some(fid) = field_id {
            if let Some(col) = reader.column(fid) {
                // Constant column — O(1) aggregation.
                if col.is_constant() {
                    let value = col.constant_value().unwrap();
                    let doc_count = col.doc_count();
                    let null_count = col.stats().map_or(0, |s| s.null_count);
                    return Box::new(ConstantMetricCollector {
                        value,
                        non_null_docs: doc_count - null_count,
                        collected: 0,
                    });
                }

                // Min/Max only — use column stats directly without loading values.
                if matches!(self.metric_type, MetricType::Min | MetricType::Max) {
                    if let Some(stats) = col.stats() {
                        return Box::new(StatsMetricCollector {
                            min: stats.min,
                            max: stats.max,
                            null_count: stats.null_count,
                            doc_count: col.doc_count(),
                            collected: 0,
                        });
                    }
                }
            }
        }

        // Pre-load numeric values once — avoid per-doc ColumnReader construction.
        let values = if let Some(fid) = field_id {
            if let Some(col) = reader.column(fid) {
                let doc_count = col.doc_count();
                (0..doc_count)
                    .map(|i| col.numeric_value(i).unwrap_or(f64::NAN))
                    .collect()
            } else {
                Vec::new()
            }
        } else {
            Vec::new()
        };

        Box::new(MetricCollector {
            values,
            sum: 0.0,
            sum_of_squares: 0.0,
            count: 0,
            min: f64::INFINITY,
            max: f64::NEG_INFINITY,
        })
    }

    fn merge_results(&self, results: Vec<AggResult>) -> AggResult {
        let mut total_sum = 0.0f64;
        let mut total_sum_of_squares = 0.0f64;
        let mut total_count = 0u64;
        let mut global_min = f64::INFINITY;
        let mut global_max = f64::NEG_INFINITY;

        for r in &results {
            if let AggResult::Metric(m) = r {
                let count = m.extra.get("count").copied().unwrap_or(0.0) as u64;
                let sum = m.extra.get("sum").copied().unwrap_or(0.0);
                let sum_sq = m.extra.get("sum_of_squares").copied().unwrap_or(0.0);
                let min = m.extra.get("min").copied().unwrap_or(f64::INFINITY);
                let max = m.extra.get("max").copied().unwrap_or(f64::NEG_INFINITY);
                total_sum += sum;
                total_sum_of_squares += sum_sq;
                total_count += count;
                if min < global_min { global_min = min; }
                if max > global_max { global_max = max; }
            }
        }

        if total_count == 0 {
            return AggResult::Metric(MetricResult::single(None));
        }

        let avg = total_sum / total_count as f64;

        match self.metric_type {
            MetricType::Avg => AggResult::Metric(MetricResult::single(Some(avg))),
            MetricType::Sum => AggResult::Metric(MetricResult::single(Some(total_sum))),
            MetricType::Min => AggResult::Metric(MetricResult::single(Some(global_min))),
            MetricType::Max => AggResult::Metric(MetricResult::single(Some(global_max))),
            MetricType::ValueCount => AggResult::Metric(MetricResult::single(Some(total_count as f64))),
            MetricType::Stats => AggResult::Metric(MetricResult::stats(
                total_count, global_min, global_max, avg, total_sum,
            )),
            MetricType::ExtendedStats => {
                let variance = (total_sum_of_squares / total_count as f64) - (avg * avg);
                let std_dev = variance.max(0.0).sqrt();
                let mut result = MetricResult::stats(
                    total_count, global_min, global_max, avg, total_sum,
                );
                result.extra.insert("sum_of_squares".into(), total_sum_of_squares);
                result.extra.insert("variance".into(), variance);
                result.extra.insert("std_deviation".into(), std_dev);
                result.extra.insert("std_deviation_bounds.upper".into(), avg + 2.0 * std_dev);
                result.extra.insert("std_deviation_bounds.lower".into(), avg - 2.0 * std_dev);
                AggResult::Metric(result)
            }
        }
    }
}

struct MetricCollector {
    /// Pre-loaded numeric values: one f64 per doc. NaN = null.
    /// Loaded once at construction, not per doc.
    values: Vec<f64>,
    sum: f64,
    sum_of_squares: f64,
    count: u64,
    min: f64,
    max: f64,
}

unsafe impl Send for MetricCollector {}

impl AggCollector for MetricCollector {
    fn collect(&mut self, doc_id: DocId) {
        let idx = doc_id.as_u32() as usize;
        if idx >= self.values.len() { return; }

        let v = self.values[idx];
        if v.is_nan() { return; } // null

        let value = Some(v);

        if let Some(v) = value {
            self.sum += v;
            self.sum_of_squares += v * v;
            self.count += 1;
            if v < self.min { self.min = v; }
            if v > self.max { self.max = v; }
        }
    }

    fn finish(self: Box<Self>) -> AggResult {
        if self.count == 0 {
            return AggResult::Metric(MetricResult::single(None));
        }
        let avg = self.sum / self.count as f64;
        // Always return full stats internally for merge to work
        let mut result = MetricResult::stats(
            self.count, self.min, self.max, avg, self.sum,
        );
        result.extra.insert("sum_of_squares".into(), self.sum_of_squares);
        AggResult::Metric(result)
    }
}

/// Collector for min/max aggregations using precomputed column stats.
/// Does not load any per-doc values — just counts collected docs and
/// returns the segment-level min/max from the zonemap.
struct StatsMetricCollector {
    min: f64,
    max: f64,
    null_count: u32,
    doc_count: u32,
    collected: u64,
}

unsafe impl Send for StatsMetricCollector {}

impl AggCollector for StatsMetricCollector {
    fn collect(&mut self, _doc_id: DocId) {
        self.collected += 1;
    }

    fn finish(self: Box<Self>) -> AggResult {
        let count = if self.doc_count == 0 { 0 } else { self.collected };
        if count == 0 {
            return AggResult::Metric(MetricResult::single(None));
        }
        // We don't have sum from stats, so use 0.0 — the merge phase
        // only uses sum for avg/sum, not min/max.
        AggResult::Metric(MetricResult::stats(
            count, self.min, self.max, 0.0, 0.0,
        ))
    }
}

/// Collector for constant-encoded columns. Counts collected docs and
/// computes metrics from the constant value × count. No per-doc value read.
struct ConstantMetricCollector {
    value: f64,
    /// Total non-null docs in the segment for this field.
    non_null_docs: u32,
    /// Number of docs actually collected (may be less if query filters).
    collected: u64,
}

unsafe impl Send for ConstantMetricCollector {}

impl AggCollector for ConstantMetricCollector {
    fn collect(&mut self, _doc_id: DocId) {
        // We don't need to read the value — it's constant.
        // Just count. If the query filters docs, we only count collected ones.
        self.collected += 1;
    }

    fn finish(self: Box<Self>) -> AggResult {
        // For match_all queries, collected == all matched docs.
        // The constant value applies to all non-null docs that were collected.
        // Approximation: assume all collected docs are non-null if the column
        // has no nulls. If it has nulls, we'd need per-doc null checks, but
        // constant columns with nulls are rare.
        let count = if self.non_null_docs == 0 {
            0
        } else {
            self.collected
        };

        if count == 0 {
            return AggResult::Metric(MetricResult::single(None));
        }

        let sum = self.value * count as f64;
        AggResult::Metric(MetricResult::stats(
            count, self.value, self.value, self.value, sum,
        ))
    }
}

// --- Geo bounds aggregation ---

pub struct GeoBoundsAggFactory {
    pub field_name: String,
}

impl AggCollectorFactory for GeoBoundsAggFactory {
    fn create_collector(&self, reader: &SegmentReader) -> Box<dyn AggCollector> {
        let field_id = reader.header().fields.iter()
            .find(|f| f.field_name == self.field_name)
            .map(|f| f.field_id);
        let store = field_id.and_then(|fid| reader.geo_points(fid));
        Box::new(GeoBoundsCollector {
            store,
            min_lat: f64::INFINITY,
            max_lat: f64::NEG_INFINITY,
            min_lon: f64::INFINITY,
            max_lon: f64::NEG_INFINITY,
            count: 0,
        })
    }

    fn merge_results(&self, results: Vec<AggResult>) -> AggResult {
        let mut min_lat = f64::INFINITY;
        let mut max_lat = f64::NEG_INFINITY;
        let mut min_lon = f64::INFINITY;
        let mut max_lon = f64::NEG_INFINITY;
        let mut count = 0u64;

        for r in &results {
            if let AggResult::Metric(m) = r {
                if let Some(&c) = m.extra.get("count") {
                    if c > 0.0 {
                        count += c as u64;
                        if let Some(&v) = m.extra.get("top_left.lat") { if v > max_lat { max_lat = v; } }
                        if let Some(&v) = m.extra.get("bottom_right.lat") { if v < min_lat { min_lat = v; } }
                        if let Some(&v) = m.extra.get("top_left.lon") { if v < min_lon { min_lon = v; } }
                        if let Some(&v) = m.extra.get("bottom_right.lon") { if v > max_lon { max_lon = v; } }
                    }
                }
            }
        }

        if count == 0 {
            return AggResult::Metric(MetricResult::single(None));
        }

        let mut result = MetricResult::single(None);
        result.extra.insert("count".into(), count as f64);
        result.extra.insert("top_left.lat".into(), max_lat);
        result.extra.insert("top_left.lon".into(), min_lon);
        result.extra.insert("bottom_right.lat".into(), min_lat);
        result.extra.insert("bottom_right.lon".into(), max_lon);
        AggResult::Metric(result)
    }
}

struct GeoBoundsCollector {
    store: Option<crate::spatial::geo::GeoPointStore>,
    min_lat: f64,
    max_lat: f64,
    min_lon: f64,
    max_lon: f64,
    count: u64,
}

unsafe impl Send for GeoBoundsCollector {}

impl AggCollector for GeoBoundsCollector {
    fn collect(&mut self, doc_id: DocId) {
        if let Some(store) = &self.store {
            if let Some(point) = store.get(doc_id.as_u32()) {
                if point.lat < self.min_lat { self.min_lat = point.lat; }
                if point.lat > self.max_lat { self.max_lat = point.lat; }
                if point.lon < self.min_lon { self.min_lon = point.lon; }
                if point.lon > self.max_lon { self.max_lon = point.lon; }
                self.count += 1;
            }
        }
    }

    fn finish(self: Box<Self>) -> AggResult {
        if self.count == 0 {
            return AggResult::Metric(MetricResult::single(None));
        }
        let mut result = MetricResult::single(None);
        result.extra.insert("count".into(), self.count as f64);
        result.extra.insert("top_left.lat".into(), self.max_lat);
        result.extra.insert("top_left.lon".into(), self.min_lon);
        result.extra.insert("bottom_right.lat".into(), self.min_lat);
        result.extra.insert("bottom_right.lon".into(), self.max_lon);
        AggResult::Metric(result)
    }
}

// --- Geo centroid aggregation ---

pub struct GeoCentroidAggFactory {
    pub field_name: String,
}

impl AggCollectorFactory for GeoCentroidAggFactory {
    fn create_collector(&self, reader: &SegmentReader) -> Box<dyn AggCollector> {
        let field_id = reader.header().fields.iter()
            .find(|f| f.field_name == self.field_name)
            .map(|f| f.field_id);
        let store = field_id.and_then(|fid| reader.geo_points(fid));
        Box::new(GeoCentroidCollector {
            store,
            sum_lat: 0.0,
            sum_lon: 0.0,
            count: 0,
        })
    }

    fn merge_results(&self, results: Vec<AggResult>) -> AggResult {
        let mut total_sum_lat = 0.0f64;
        let mut total_sum_lon = 0.0f64;
        let mut total_count = 0u64;

        for r in &results {
            if let AggResult::Metric(m) = r {
                let count = m.extra.get("count").copied().unwrap_or(0.0) as u64;
                let sum_lat = m.extra.get("sum_lat").copied().unwrap_or(0.0);
                let sum_lon = m.extra.get("sum_lon").copied().unwrap_or(0.0);
                total_count += count;
                total_sum_lat += sum_lat;
                total_sum_lon += sum_lon;
            }
        }

        if total_count == 0 {
            return AggResult::Metric(MetricResult::single(None));
        }

        let mut result = MetricResult::single(None);
        result.extra.insert("count".into(), total_count as f64);
        result.extra.insert("lat".into(), total_sum_lat / total_count as f64);
        result.extra.insert("lon".into(), total_sum_lon / total_count as f64);
        result.extra.insert("sum_lat".into(), total_sum_lat);
        result.extra.insert("sum_lon".into(), total_sum_lon);
        AggResult::Metric(result)
    }
}

struct GeoCentroidCollector {
    store: Option<crate::spatial::geo::GeoPointStore>,
    sum_lat: f64,
    sum_lon: f64,
    count: u64,
}

unsafe impl Send for GeoCentroidCollector {}

impl AggCollector for GeoCentroidCollector {
    fn collect(&mut self, doc_id: DocId) {
        if let Some(store) = &self.store {
            if let Some(point) = store.get(doc_id.as_u32()) {
                self.sum_lat += point.lat;
                self.sum_lon += point.lon;
                self.count += 1;
            }
        }
    }

    fn finish(self: Box<Self>) -> AggResult {
        if self.count == 0 {
            return AggResult::Metric(MetricResult::single(None));
        }
        let mut result = MetricResult::single(None);
        result.extra.insert("count".into(), self.count as f64);
        result.extra.insert("lat".into(), self.sum_lat / self.count as f64);
        result.extra.insert("lon".into(), self.sum_lon / self.count as f64);
        result.extra.insert("sum_lat".into(), self.sum_lat);
        result.extra.insert("sum_lon".into(), self.sum_lon);
        AggResult::Metric(result)
    }
}

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

    // Unit tests for merge logic
    #[test]
    fn merge_avg() {
        let factory = MetricAggFactory {
            field_name: "price".into(),
            metric_type: MetricType::Avg,
        };
        let results = vec![
            AggResult::Metric(MetricResult::stats(3, 1.0, 3.0, 2.0, 6.0)),
            AggResult::Metric(MetricResult::stats(2, 4.0, 5.0, 4.5, 9.0)),
        ];
        let merged = factory.merge_results(results);
        if let AggResult::Metric(m) = merged {
            assert_eq!(m.value, Some(3.0)); // (6+9) / 5 = 3.0
        } else { panic!(); }
    }

    #[test]
    fn merge_sum() {
        let factory = MetricAggFactory {
            field_name: "x".into(),
            metric_type: MetricType::Sum,
        };
        let results = vec![
            AggResult::Metric(MetricResult::stats(2, 0.0, 0.0, 0.0, 10.0)),
            AggResult::Metric(MetricResult::stats(3, 0.0, 0.0, 0.0, 20.0)),
        ];
        let merged = factory.merge_results(results);
        if let AggResult::Metric(m) = merged {
            assert_eq!(m.value, Some(30.0));
        } else { panic!(); }
    }

    #[test]
    fn merge_min_max() {
        let factory = MetricAggFactory {
            field_name: "x".into(),
            metric_type: MetricType::Min,
        };
        let results = vec![
            AggResult::Metric(MetricResult::stats(1, 5.0, 5.0, 5.0, 5.0)),
            AggResult::Metric(MetricResult::stats(1, 2.0, 2.0, 2.0, 2.0)),
        ];
        let merged = factory.merge_results(results);
        if let AggResult::Metric(m) = merged {
            assert_eq!(m.value, Some(2.0));
        } else { panic!(); }
    }

    #[test]
    fn merge_empty() {
        let factory = MetricAggFactory {
            field_name: "x".into(),
            metric_type: MetricType::Avg,
        };
        let merged = factory.merge_results(vec![
            AggResult::Metric(MetricResult::single(None)),
        ]);
        if let AggResult::Metric(m) = merged {
            assert!(m.value.is_none());
        } else { panic!(); }
    }
}
