//! Bucket aggregation implementations: terms, range, histogram, filter.
//!
//! See [[aggregations]] and [[feature-aggregations-v010#Step 6]].

use std::collections::HashMap;
use luci_core::{DocId, FieldId};

use super::*;
use crate::columnar::reader::ColumnReader;
use crate::segment::reader::SegmentReader;

/// Pre-load all numeric column values into a Vec<f64> for a field.
/// NaN represents null. Avoids per-doc ColumnReader construction.
fn preload_numeric_values(field_id: Option<FieldId>, reader: &SegmentReader) -> Vec<f64> {
    let Some(fid) = field_id else { return Vec::new() };
    let Some(col) = reader.column(fid) else { return Vec::new() };
    let doc_count = col.doc_count();
    (0..doc_count)
        .map(|i| col.numeric_value(i).unwrap_or(f64::NAN))
        .collect()
}

// --- Terms aggregation ---

pub struct TermsAggFactory {
    pub field_name: String,
    pub size: usize,
    pub sub_agg_factories: Vec<(String, Box<dyn AggCollectorFactory>)>,
}

impl AggCollectorFactory for TermsAggFactory {
    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);

        // Pre-load ordinals and dictionary once — avoid per-doc ColumnReader construction.
        // Follows the Lucene pattern: acquire DocValues once in getLeafCollector().
        let (ordinals, dict) = if let Some(fid) = field_id {
            if let Some(col) = reader.column(fid) {
                let doc_count = col.doc_count();
                let mut ords = Vec::with_capacity(doc_count as usize);
                for i in 0..doc_count {
                    ords.push(col.keyword_ordinal(i).unwrap_or(u32::MAX));
                }
                let dict: Vec<String> = (0..col.dict_size())
                    .map(|i| col.ordinal_to_string(i as u32).unwrap_or("?").to_string())
                    .collect();
                (ords, dict)
            } else {
                (Vec::new(), Vec::new())
            }
        } else {
            (Vec::new(), Vec::new())
        };

        let dict_size = dict.len();

        Box::new(TermsCollector {
            field_id,
            segment_data: reader as *const SegmentReader,
            ordinals,
            dict,
            ordinal_buckets: Vec::with_capacity(dict_size),
            sub_agg_factories: self.sub_agg_factories.as_slice() as *const [(String, Box<dyn AggCollectorFactory>)],
        })
    }

    fn merge_results(&self, results: Vec<AggResult>) -> AggResult {
        // Merge: combine bucket counts and sub-agg results across segments
        let mut merged: HashMap<String, (u64, HashMap<String, Vec<AggResult>>)> = HashMap::new();

        for r in results {
            if let AggResult::Bucket(br) = r {
                for bucket in br.buckets {
                    if let BucketKey::String(key) = bucket.key {
                        let entry = merged.entry(key).or_insert_with(|| (0, HashMap::new()));
                        entry.0 += bucket.doc_count;
                        for (name, result) in bucket.sub_aggs {
                            entry.1.entry(name).or_default().push(result);
                        }
                    }
                }
            }
        }

        // Sort by count descending, take top N
        let mut bucket_entries: Vec<_> = merged.into_iter().collect();
        bucket_entries.sort_by(|a, b| b.1.0.cmp(&a.1.0));
        bucket_entries.truncate(self.size);

        let result_buckets = bucket_entries
            .into_iter()
            .map(|(key, (doc_count, sub_partials))| {
                // Merge sub-agg results
                let mut sub_aggs = HashMap::new();
                for (i, (name, factory)) in self.sub_agg_factories.iter().enumerate() {
                    if let Some(partials) = sub_partials.get(name) {
                        sub_aggs.insert(name.clone(), factory.merge_results(partials.clone()));
                    }
                }
                Bucket {
                    key: BucketKey::String(key),
                    doc_count,
                    sub_aggs,
                }
            })
            .collect();

        AggResult::Bucket(BucketResult { buckets: result_buckets })
    }
}

/// Per-bucket state: count + sub-agg collectors.
struct BucketState {
    count: u64,
    sub_collectors: Vec<(String, Box<dyn AggCollector>)>,
}

/// Ordinal-indexed terms collector. Pre-loads ordinals into an owned Vec
/// at construction time, avoiding per-doc ColumnReader construction.
/// Follows the Lucene/Tantivy pattern: acquire column once per segment.
struct TermsCollector {
    field_id: Option<FieldId>,
    segment_data: *const SegmentReader,
    /// Pre-loaded ordinals: one u32 per doc. u32::MAX = null.
    ordinals: Vec<u32>,
    /// Pre-loaded dictionary for ordinal→string resolution at finish time.
    dict: Vec<String>,
    /// Vec indexed by ordinal. None = ordinal not yet seen.
    ordinal_buckets: Vec<Option<BucketState>>,
    sub_agg_factories: *const [(String, Box<dyn AggCollectorFactory>)],
}

unsafe impl Send for TermsCollector {}

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

        let ordinal = self.ordinals[idx];
        if ordinal == u32::MAX { return; } // null
        let ordinal = ordinal as usize;

        // Ensure Vec is large enough
        if ordinal >= self.ordinal_buckets.len() {
            self.ordinal_buckets.resize_with(ordinal + 1, || None);
        }

        let reader = unsafe { &*self.segment_data };
        let sub_factories = unsafe { &*self.sub_agg_factories };
        let state = self.ordinal_buckets[ordinal].get_or_insert_with(|| {
            let subs = sub_factories.iter()
                .map(|(name, factory)| {
                    (name.clone(), factory.create_collector(reader))
                })
                .collect();
            BucketState { count: 0, sub_collectors: subs }
        });
        state.count += 1;
        for (_, collector) in &mut state.sub_collectors {
            collector.collect(doc_id);
        }
    }

    fn finish(self: Box<Self>) -> AggResult {
        let mut buckets: Vec<Bucket> = self.ordinal_buckets
            .into_iter()
            .enumerate()
            .filter_map(|(ordinal, state)| {
                let state = state?;
                let key = self.dict.get(ordinal)
                    .cloned()
                    .unwrap_or_else(|| "?".to_string());
                let sub_aggs: HashMap<String, AggResult> = state.sub_collectors
                    .into_iter()
                    .map(|(name, collector)| (name, collector.finish()))
                    .collect();
                Some(Bucket {
                    key: BucketKey::String(key),
                    doc_count: state.count,
                    sub_aggs,
                })
            })
            .collect();

        buckets.sort_by(|a, b| b.doc_count.cmp(&a.doc_count));
        AggResult::Bucket(BucketResult { buckets })
    }
}

// --- Range aggregation ---

pub struct RangeAggFactory {
    pub field_name: String,
    pub ranges: Vec<RangeDef>,
}

impl AggCollectorFactory for RangeAggFactory {
    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);

        // Pre-load numeric values once — avoid per-doc ColumnReader construction.
        let values = preload_numeric_values(field_id, reader);

        Box::new(RangeCollector {
            values,
            ranges: self.ranges.clone(),
            counts: vec![0u64; self.ranges.len()],
        })
    }

    fn merge_results(&self, results: Vec<AggResult>) -> AggResult {
        let mut merged_counts = vec![0u64; self.ranges.len()];

        for r in results {
            if let AggResult::Bucket(br) = r {
                for (i, bucket) in br.buckets.iter().enumerate() {
                    if i < merged_counts.len() {
                        merged_counts[i] += bucket.doc_count;
                    }
                }
            }
        }

        let buckets = self.ranges.iter().zip(merged_counts.iter())
            .map(|(range, &count)| Bucket {
                key: BucketKey::Range { from: range.from, to: range.to },
                doc_count: count,
                sub_aggs: HashMap::new(),
            })
            .collect();

        AggResult::Bucket(BucketResult { buckets })
    }
}

struct RangeCollector {
    /// Pre-loaded numeric values: one f64 per doc. NaN = null.
    values: Vec<f64>,
    ranges: Vec<RangeDef>,
    counts: Vec<u64>,
}

unsafe impl Send for RangeCollector {}

impl AggCollector for RangeCollector {
    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; }

        for (i, range) in self.ranges.iter().enumerate() {
            let above_from = range.from.map_or(true, |f| v >= f);
            let below_to = range.to.map_or(true, |t| v < t);
            if above_from && below_to {
                self.counts[i] += 1;
            }
        }
    }

    fn finish(self: Box<Self>) -> AggResult {
        let buckets = self.ranges.iter().zip(self.counts.iter())
            .map(|(range, &count)| Bucket {
                key: BucketKey::Range { from: range.from, to: range.to },
                doc_count: count,
                sub_aggs: HashMap::new(),
            })
            .collect();
        AggResult::Bucket(BucketResult { buckets })
    }
}

// --- Histogram aggregation ---

pub struct HistogramAggFactory {
    pub field_name: String,
    pub interval: f64,
}

impl AggCollectorFactory for HistogramAggFactory {
    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);

        // Pre-load numeric values once — avoid per-doc ColumnReader construction.
        let values = preload_numeric_values(field_id, reader);

        Box::new(HistogramCollector {
            values,
            interval: self.interval,
            buckets: HashMap::new(),
        })
    }

    fn merge_results(&self, results: Vec<AggResult>) -> AggResult {
        let mut merged: HashMap<i64, u64> = HashMap::new();

        for r in results {
            if let AggResult::Bucket(br) = r {
                for bucket in br.buckets {
                    if let BucketKey::Number(key) = bucket.key {
                        *merged.entry(key as i64).or_insert(0) += bucket.doc_count;
                    }
                }
            }
        }

        let mut buckets: Vec<Bucket> = merged
            .into_iter()
            .map(|(key, count)| Bucket {
                key: BucketKey::Number(key as f64),
                doc_count: count,
                sub_aggs: HashMap::new(),
            })
            .collect();

        buckets.sort_by(|a, b| {
            let ka = if let BucketKey::Number(n) = a.key { n } else { 0.0 };
            let kb = if let BucketKey::Number(n) = b.key { n } else { 0.0 };
            ka.partial_cmp(&kb).unwrap_or(std::cmp::Ordering::Equal)
        });

        AggResult::Bucket(BucketResult { buckets })
    }
}

struct HistogramCollector {
    /// Pre-loaded numeric values: one f64 per doc. NaN = null.
    values: Vec<f64>,
    interval: f64,
    buckets: HashMap<i64, u64>,
}

unsafe impl Send for HistogramCollector {}

impl AggCollector for HistogramCollector {
    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; }

        let bucket_key = (v / self.interval).floor() as i64 * self.interval as i64;
        *self.buckets.entry(bucket_key).or_insert(0) += 1;
    }

    fn finish(self: Box<Self>) -> AggResult {
        let mut buckets: Vec<Bucket> = self.buckets
            .into_iter()
            .map(|(key, count)| Bucket {
                key: BucketKey::Number(key as f64),
                doc_count: count,
                sub_aggs: HashMap::new(),
            })
            .collect();

        buckets.sort_by(|a, b| {
            let ka = if let BucketKey::Number(n) = a.key { n } else { 0.0 };
            let kb = if let BucketKey::Number(n) = b.key { n } else { 0.0 };
            ka.partial_cmp(&kb).unwrap_or(std::cmp::Ordering::Equal)
        });

        AggResult::Bucket(BucketResult { buckets })
    }
}

// --- Date histogram aggregation ---

pub struct DateHistogramAggFactory {
    pub field_name: String,
    pub interval: DateInterval,
}

impl AggCollectorFactory for DateHistogramAggFactory {
    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 values = preload_numeric_values(field_id, reader);

        Box::new(DateHistogramCollector {
            values,
            interval: self.interval.clone(),
            buckets: HashMap::new(),
        })
    }

    fn merge_results(&self, results: Vec<AggResult>) -> AggResult {
        let mut merged: HashMap<i64, u64> = HashMap::new();
        for r in results {
            if let AggResult::Bucket(br) = r {
                for bucket in br.buckets {
                    if let BucketKey::Number(key) = bucket.key {
                        *merged.entry(key as i64).or_insert(0) += bucket.doc_count;
                    }
                }
            }
        }
        let mut buckets: Vec<Bucket> = merged.into_iter()
            .map(|(key, count)| Bucket {
                key: BucketKey::Number(key as f64),
                doc_count: count,
                sub_aggs: HashMap::new(),
            })
            .collect();
        buckets.sort_by(|a, b| {
            let ka = if let BucketKey::Number(n) = a.key { n } else { 0.0 };
            let kb = if let BucketKey::Number(n) = b.key { n } else { 0.0 };
            ka.partial_cmp(&kb).unwrap_or(std::cmp::Ordering::Equal)
        });
        AggResult::Bucket(BucketResult { buckets })
    }
}

struct DateHistogramCollector {
    values: Vec<f64>,
    interval: DateInterval,
    buckets: HashMap<i64, u64>,
}

unsafe impl Send for DateHistogramCollector {}

impl DateHistogramCollector {
    /// Compute the bucket key (epoch millis of the interval start) for a value.
    fn bucket_key(&self, epoch_millis: f64) -> i64 {
        match &self.interval {
            DateInterval::Fixed(ms) => {
                ((epoch_millis / ms).floor() as i64) * (*ms as i64)
            }
            DateInterval::Calendar(cal) => {
                calendar_floor(epoch_millis as i64, cal)
            }
        }
    }
}

/// Floor an epoch millis timestamp to the start of a calendar interval.
/// Simple UTC-only implementation (no timezone support).
fn calendar_floor(epoch_ms: i64, interval: &CalendarInterval) -> i64 {
    const MS_PER_SEC: i64 = 1_000;
    const MS_PER_MIN: i64 = 60 * MS_PER_SEC;
    const MS_PER_HOUR: i64 = 60 * MS_PER_MIN;
    const MS_PER_DAY: i64 = 24 * MS_PER_HOUR;

    match interval {
        CalendarInterval::Minute => (epoch_ms / MS_PER_MIN) * MS_PER_MIN,
        CalendarInterval::Hour => (epoch_ms / MS_PER_HOUR) * MS_PER_HOUR,
        CalendarInterval::Day => (epoch_ms / MS_PER_DAY) * MS_PER_DAY,
        CalendarInterval::Week => {
            // ISO week: Monday is day 0. Epoch (1970-01-01) was Thursday (day 3).
            let days = epoch_ms / MS_PER_DAY;
            let week_start = days - ((days + 3) % 7); // align to Monday
            week_start * MS_PER_DAY
        }
        CalendarInterval::Month | CalendarInterval::Quarter | CalendarInterval::Year => {
            // Convert epoch_ms to (year, month, day) then floor
            let days_since_epoch = epoch_ms / MS_PER_DAY;
            let (y, m, _d) = days_to_ymd(days_since_epoch);
            let floored_month = match interval {
                CalendarInterval::Month => m,
                CalendarInterval::Quarter => ((m - 1) / 3) * 3 + 1,
                CalendarInterval::Year => 1,
                _ => unreachable!(),
            };
            let floored_year = if matches!(interval, CalendarInterval::Year) { y } else { y };
            ymd_to_epoch_ms(floored_year, floored_month, 1)
        }
    }
}

/// Convert days since epoch to (year, month, day). Simplified algorithm.
fn days_to_ymd(days: i64) -> (i32, u32, u32) {
    // Algorithm from Howard Hinnant's civil_from_days
    let z = days + 719468;
    let era = if z >= 0 { z } else { z - 146096 } / 146097;
    let doe = (z - era * 146097) as u32;
    let yoe = (doe - doe / 1460 + doe / 36524 - doe / 146096) / 365;
    let y = yoe as i64 + era * 400;
    let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
    let mp = (5 * doy + 2) / 153;
    let d = doy - (153 * mp + 2) / 5 + 1;
    let m = if mp < 10 { mp + 3 } else { mp - 9 };
    let y = if m <= 2 { y + 1 } else { y };
    (y as i32, m, d)
}

/// Convert (year, month, day) to epoch milliseconds.
fn ymd_to_epoch_ms(y: i32, m: u32, d: u32) -> i64 {
    let y = if m <= 2 { y as i64 - 1 } else { y as i64 };
    let era = if y >= 0 { y } else { y - 399 } / 400;
    let yoe = (y - era * 400) as u32;
    let doy = (153 * (if m > 2 { m - 3 } else { m + 9 }) + 2) / 5 + d - 1;
    let doe = yoe * 365 + yoe / 4 - yoe / 100 + doy;
    let days = era * 146097 + doe as i64 - 719468;
    days * 86_400_000
}

impl AggCollector for DateHistogramCollector {
    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; }
        let key = self.bucket_key(v);
        *self.buckets.entry(key).or_insert(0) += 1;
    }

    fn finish(self: Box<Self>) -> AggResult {
        let mut buckets: Vec<Bucket> = self.buckets.into_iter()
            .map(|(key, count)| Bucket {
                key: BucketKey::Number(key as f64),
                doc_count: count,
                sub_aggs: HashMap::new(),
            })
            .collect();
        buckets.sort_by(|a, b| {
            let ka = if let BucketKey::Number(n) = a.key { n } else { 0.0 };
            let kb = if let BucketKey::Number(n) = b.key { n } else { 0.0 };
            ka.partial_cmp(&kb).unwrap_or(std::cmp::Ordering::Equal)
        });
        AggResult::Bucket(BucketResult { buckets })
    }
}

// --- Nested aggregation ---

pub struct NestedAggFactory {
    pub path: String,
    pub sub_agg_factories: Vec<(String, Box<dyn AggCollectorFactory>)>,
}

impl AggCollectorFactory for NestedAggFactory {
    fn create_collector(&self, reader: &SegmentReader) -> Box<dyn AggCollector> {
        let parent_bitset = reader.parent_bitset().map(|b| b.to_vec());
        let sub_collectors: Vec<(String, Box<dyn AggCollector>)> = self.sub_agg_factories
            .iter()
            .map(|(name, factory)| (name.clone(), factory.create_collector(reader)))
            .collect();
        Box::new(NestedAggCollector {
            parent_bitset,
            sub_collectors,
        })
    }

    fn merge_results(&self, results: Vec<AggResult>) -> AggResult {
        // Nested agg is a single-bucket agg: merge doc_counts and sub-aggs
        let mut total_count = 0u64;
        let mut sub_results: HashMap<String, Vec<AggResult>> = HashMap::new();

        for r in results {
            if let AggResult::Bucket(br) = r {
                for b in br.buckets {
                    total_count += b.doc_count;
                    for (name, sub_r) in b.sub_aggs {
                        sub_results.entry(name).or_default().push(sub_r);
                    }
                }
            }
        }

        let mut merged_sub_aggs = HashMap::new();
        for ((name, factory), (_, results)) in self.sub_agg_factories.iter()
            .zip(sub_results.iter())
        {
            merged_sub_aggs.insert(name.clone(), factory.merge_results(results.clone()));
        }

        AggResult::Bucket(BucketResult {
            buckets: vec![Bucket {
                key: BucketKey::String("nested".into()),
                doc_count: total_count,
                sub_aggs: merged_sub_aggs,
            }],
        })
    }
}

struct NestedAggCollector {
    parent_bitset: Option<Vec<bool>>,
    sub_collectors: Vec<(String, Box<dyn AggCollector>)>,
}

unsafe impl Send for NestedAggCollector {}

impl NestedAggCollector {
    /// Given a parent doc_id, find all nested child doc_ids.
    fn children_of(&self, parent_doc: u32) -> Vec<u32> {
        let Some(bitset) = &self.parent_bitset else { return vec![]; };
        let mut children = Vec::new();
        let start = parent_doc as usize + 1;
        for i in start..bitset.len() {
            if bitset[i] {
                break; // next parent
            }
            children.push(i as u32);
        }
        children
    }
}

impl AggCollector for NestedAggCollector {
    fn collect(&mut self, doc_id: DocId) {
        // Expand parent doc to its nested children and collect each
        let children = self.children_of(doc_id.as_u32());
        for child_id in children {
            for (_, collector) in &mut self.sub_collectors {
                collector.collect(DocId::new(child_id));
            }
        }
    }

    fn finish(self: Box<Self>) -> AggResult {
        let mut sub_aggs = HashMap::new();
        let mut total_children = 0u64;
        for (name, collector) in self.sub_collectors {
            let result = collector.finish();
            // Count total nested docs collected
            if let AggResult::Bucket(ref br) = result {
                for b in &br.buckets {
                    total_children += b.doc_count;
                }
            }
            sub_aggs.insert(name, result);
        }
        AggResult::Bucket(BucketResult {
            buckets: vec![Bucket {
                key: BucketKey::String("nested".into()),
                doc_count: total_children,
                sub_aggs,
            }],
        })
    }
}

// --- Reverse nested aggregation ---

pub struct ReverseNestedAggFactory {
    pub sub_agg_factories: Vec<(String, Box<dyn AggCollectorFactory>)>,
}

impl AggCollectorFactory for ReverseNestedAggFactory {
    fn create_collector(&self, reader: &SegmentReader) -> Box<dyn AggCollector> {
        let parent_bitset = reader.parent_bitset().map(|b| b.to_vec());
        let sub_collectors: Vec<(String, Box<dyn AggCollector>)> = self.sub_agg_factories
            .iter()
            .map(|(name, factory)| (name.clone(), factory.create_collector(reader)))
            .collect();
        Box::new(ReverseNestedAggCollector {
            parent_bitset,
            sub_collectors,
            seen_parents: std::collections::HashSet::new(),
        })
    }

    fn merge_results(&self, results: Vec<AggResult>) -> AggResult {
        let mut total_count = 0u64;
        let mut sub_results: HashMap<String, Vec<AggResult>> = HashMap::new();

        for r in results {
            if let AggResult::Bucket(br) = r {
                for b in br.buckets {
                    total_count += b.doc_count;
                    for (name, sub_r) in b.sub_aggs {
                        sub_results.entry(name).or_default().push(sub_r);
                    }
                }
            }
        }

        let mut merged_sub_aggs = HashMap::new();
        for ((name, factory), (_, results)) in self.sub_agg_factories.iter()
            .zip(sub_results.iter())
        {
            merged_sub_aggs.insert(name.clone(), factory.merge_results(results.clone()));
        }

        AggResult::Bucket(BucketResult {
            buckets: vec![Bucket {
                key: BucketKey::String("reverse_nested".into()),
                doc_count: total_count,
                sub_aggs: merged_sub_aggs,
            }],
        })
    }
}

struct ReverseNestedAggCollector {
    parent_bitset: Option<Vec<bool>>,
    sub_collectors: Vec<(String, Box<dyn AggCollector>)>,
    /// Track which parents we've already collected to avoid duplicates
    /// (multiple nested children may map to the same parent).
    seen_parents: std::collections::HashSet<u32>,
}

unsafe impl Send for ReverseNestedAggCollector {}

impl ReverseNestedAggCollector {
    /// Find the parent doc for a nested child doc.
    fn find_parent(&self, nested_doc: u32) -> Option<u32> {
        let bitset = self.parent_bitset.as_ref()?;
        let mut i = nested_doc as usize;
        while i > 0 {
            if i < bitset.len() && bitset[i] {
                return Some(i as u32);
            }
            i -= 1;
        }
        if !bitset.is_empty() && bitset[0] { Some(0) } else { None }
    }
}

impl AggCollector for ReverseNestedAggCollector {
    fn collect(&mut self, doc_id: DocId) {
        // Map nested child doc back to its parent
        if let Some(parent_id) = self.find_parent(doc_id.as_u32()) {
            if self.seen_parents.insert(parent_id) {
                for (_, collector) in &mut self.sub_collectors {
                    collector.collect(DocId::new(parent_id));
                }
            }
        }
    }

    fn finish(self: Box<Self>) -> AggResult {
        let parent_count = self.seen_parents.len() as u64;
        let mut sub_aggs = HashMap::new();
        for (name, collector) in self.sub_collectors {
            sub_aggs.insert(name, collector.finish());
        }
        AggResult::Bucket(BucketResult {
            buckets: vec![Bucket {
                key: BucketKey::String("reverse_nested".into()),
                doc_count: parent_count,
                sub_aggs,
            }],
        })
    }
}

// --- Geohash grid aggregation ---

pub struct GeohashGridAggFactory {
    pub field_name: String,
    pub precision: usize,
    pub size: usize,
}

impl AggCollectorFactory for GeohashGridAggFactory {
    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(GeohashGridCollector {
            store,
            precision: self.precision,
            buckets: HashMap::new(),
        })
    }

    fn merge_results(&self, results: Vec<AggResult>) -> AggResult {
        let mut merged: HashMap<String, u64> = HashMap::new();
        for r in results {
            if let AggResult::Bucket(br) = r {
                for b in br.buckets {
                    if let BucketKey::String(key) = b.key {
                        *merged.entry(key).or_insert(0) += b.doc_count;
                    }
                }
            }
        }
        let mut buckets: Vec<Bucket> = merged.into_iter()
            .map(|(key, count)| Bucket {
                key: BucketKey::String(key),
                doc_count: count,
                sub_aggs: HashMap::new(),
            })
            .collect();
        // Sort by doc_count descending, truncate to size
        buckets.sort_by(|a, b| b.doc_count.cmp(&a.doc_count));
        buckets.truncate(self.size);
        AggResult::Bucket(BucketResult { buckets })
    }
}

struct GeohashGridCollector {
    store: Option<crate::spatial::geo::GeoPointStore>,
    precision: usize,
    buckets: HashMap<String, u64>,
}

unsafe impl Send for GeohashGridCollector {}

impl AggCollector for GeohashGridCollector {
    fn collect(&mut self, doc_id: DocId) {
        if let Some(store) = &self.store {
            if let Some(point) = store.get(doc_id.as_u32()) {
                if let Ok(hash) = geohash::encode(
                    geohash::Coord { x: point.lon, y: point.lat },
                    self.precision,
                ) {
                    *self.buckets.entry(hash).or_insert(0) += 1;
                }
            }
        }
    }

    fn finish(self: Box<Self>) -> AggResult {
        let mut buckets: Vec<Bucket> = self.buckets.into_iter()
            .map(|(key, count)| Bucket {
                key: BucketKey::String(key),
                doc_count: count,
                sub_aggs: HashMap::new(),
            })
            .collect();
        buckets.sort_by(|a, b| b.doc_count.cmp(&a.doc_count));
        AggResult::Bucket(BucketResult { buckets })
    }
}

// --- Top hits aggregation ---

pub struct TopHitsAggFactory {
    pub size: usize,
}

impl AggCollectorFactory for TopHitsAggFactory {
    fn create_collector(&self, reader: &SegmentReader) -> Box<dyn AggCollector> {
        Box::new(TopHitsCollector {
            segment: reader as *const SegmentReader,
            doc_ids: Vec::new(),
            size: self.size,
        })
    }

    fn merge_results(&self, results: Vec<AggResult>) -> AggResult {
        let mut all_hits: Vec<serde_json::Value> = Vec::new();
        for r in results {
            if let AggResult::Hits(h) = r {
                all_hits.extend(h.hits);
            }
        }
        all_hits.truncate(self.size);
        AggResult::Hits(HitsResult { hits: all_hits })
    }
}

struct TopHitsCollector {
    segment: *const SegmentReader,
    doc_ids: Vec<u32>,
    size: usize,
}

unsafe impl Send for TopHitsCollector {}

impl AggCollector for TopHitsCollector {
    fn collect(&mut self, doc_id: DocId) {
        if self.doc_ids.len() < self.size {
            self.doc_ids.push(doc_id.as_u32());
        }
    }

    fn finish(self: Box<Self>) -> AggResult {
        let reader = unsafe { &*self.segment };
        let doc_store = reader.doc_store();
        let mut hits = Vec::new();
        for &doc_id in &self.doc_ids {
            if let Some(source_bytes) = doc_store.get(doc_id) {
                if let Ok(source) = serde_json::from_slice::<serde_json::Value>(&source_bytes) {
                    hits.push(serde_json::json!({
                        "_doc_id": doc_id,
                        "_source": source,
                    }));
                }
            }
        }
        AggResult::Hits(HitsResult { hits })
    }
}

// --- Filter aggregation ---

pub struct FilterAggFactory {
    pub query: Box<dyn crate::query::Query>,
}

impl AggCollectorFactory for FilterAggFactory {
    fn create_collector(&self, _reader: &SegmentReader) -> Box<dyn AggCollector> {
        // Filter agg: counts docs matching the filter query.
        // For simplicity in M3, we count all collected docs (the filter is
        // applied at query level). A proper implementation would re-evaluate
        // the filter per doc.
        Box::new(FilterCollector { count: 0 })
    }

    fn merge_results(&self, results: Vec<AggResult>) -> AggResult {
        let mut total = 0u64;
        for r in results {
            if let AggResult::Bucket(br) = r {
                for b in &br.buckets {
                    total += b.doc_count;
                }
            }
        }
        AggResult::Bucket(BucketResult {
            buckets: vec![Bucket {
                key: BucketKey::String("filter".into()),
                doc_count: total,
                sub_aggs: HashMap::new(),
            }],
        })
    }
}

struct FilterCollector {
    count: u64,
}

// --- Filters (multi-filter) aggregation ---

pub struct FiltersAggFactory {
    /// Named filters: (bucket_name, query)
    pub filters: Vec<(String, Box<dyn crate::query::Query>)>,
}

impl AggCollectorFactory for FiltersAggFactory {
    fn create_collector(&self, reader: &SegmentReader) -> Box<dyn AggCollector> {
        use luci_core::{ScoreMode, Scorer as _};
        use crate::search::searcher::Searcher;

        // Build a bitset per filter by running each filter query against the segment
        let searcher = Searcher::new(vec![]); // empty searcher for weight creation
        let mut filter_bitsets: Vec<(String, Vec<bool>)> = Vec::new();
        let doc_count = reader.doc_count() as usize;

        for (name, query) in &self.filters {
            let mut bitset = vec![false; doc_count];
            if let Ok(weight) = query.create_weight(&searcher, ScoreMode::CompleteNoScores) {
                if let Ok(Some(supplier)) = weight.scorer_supplier(reader) {
                    if let Ok(mut scorer) = supplier.scorer() {
                        while scorer.doc_id() != luci_core::NO_MORE_DOCS {
                            let id = scorer.doc_id().as_u32() as usize;
                            if id < doc_count {
                                bitset[id] = true;
                            }
                            scorer.next();
                        }
                    }
                }
            }
            filter_bitsets.push((name.clone(), bitset));
        }

        Box::new(FiltersCollector {
            filter_bitsets,
            counts: vec![0u64; self.filters.len()],
        })
    }

    fn merge_results(&self, results: Vec<AggResult>) -> AggResult {
        let num_filters = self.filters.len();
        let mut totals = vec![0u64; num_filters];
        let names: Vec<String> = self.filters.iter().map(|(n, _)| n.clone()).collect();

        for r in results {
            if let AggResult::Bucket(br) = r {
                for (i, b) in br.buckets.iter().enumerate() {
                    if i < num_filters {
                        totals[i] += b.doc_count;
                    }
                }
            }
        }

        AggResult::Bucket(BucketResult {
            buckets: names.iter().zip(totals.iter())
                .map(|(name, &count)| Bucket {
                    key: BucketKey::String(name.clone()),
                    doc_count: count,
                    sub_aggs: HashMap::new(),
                })
                .collect(),
        })
    }
}

struct FiltersCollector {
    filter_bitsets: Vec<(String, Vec<bool>)>,
    counts: Vec<u64>,
}

unsafe impl Send for FiltersCollector {}

impl AggCollector for FiltersCollector {
    fn collect(&mut self, doc_id: DocId) {
        let id = doc_id.as_u32() as usize;
        for (i, (_, bitset)) in self.filter_bitsets.iter().enumerate() {
            if id < bitset.len() && bitset[id] {
                self.counts[i] += 1;
            }
        }
    }

    fn finish(self: Box<Self>) -> AggResult {
        AggResult::Bucket(BucketResult {
            buckets: self.filter_bitsets.iter().zip(self.counts.iter())
                .map(|((name, _), &count)| Bucket {
                    key: BucketKey::String(name.clone()),
                    doc_count: count,
                    sub_aggs: HashMap::new(),
                })
                .collect(),
        })
    }
}

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

    fn finish(self: Box<Self>) -> AggResult {
        AggResult::Bucket(BucketResult {
            buckets: vec![Bucket {
                key: BucketKey::String("filter".into()),
                doc_count: self.count,
                sub_aggs: HashMap::new(),
            }],
        })
    }
}
