//! Searcher: multi-segment term search with BM25 scoring and source retrieval.
//!
//! Provides the `search(field, term, top_k)` API that searches across all
//! committed segments, scores with BM25, merges results, and retrieves
//! `_source` for the top hits.
//!
//! See [[architecture-query-execution]] and [[architecture-overview#Step 10]].

use std::sync::atomic::{AtomicU32, Ordering};

use rayon::prelude::*;

use luci_analysis::AnalyzerRegistry;
use luci_core::{DocId, FieldId, Scorer as _, SegmentId, NO_MORE_DOCS};

use crate::search::bm25::{Bm25Scorer, Bm25Weight};
use crate::search::collector::TopDocsCollector;
use crate::segment::reader::SegmentReader;

/// Shared minimum competitive score across parallel search threads.
///
/// Each thread reads the global min to initialize its local threshold,
/// and updates it when its local min increases. The global min only
/// increases monotonically. Relaxed ordering is sufficient — stale
/// reads only cause slightly delayed pruning, not correctness issues.
struct SharedMinScore(AtomicU32);

impl SharedMinScore {
    fn new() -> Self {
        Self(AtomicU32::new(0.0f32.to_bits()))
    }

    fn get(&self) -> f32 {
        f32::from_bits(self.0.load(Ordering::Relaxed))
    }

    fn update(&self, new_min: f32) {
        let mut current = self.0.load(Ordering::Relaxed);
        loop {
            if new_min <= f32::from_bits(current) {
                break;
            }
            match self.0.compare_exchange_weak(
                current,
                new_min.to_bits(),
                Ordering::Relaxed,
                Ordering::Relaxed,
            ) {
                Ok(_) => break,
                Err(actual) => current = actual,
            }
        }
    }
}

/// A search hit with score, source document, and optional field values.
#[derive(Clone, Debug)]
pub struct SearchHit {
    pub doc_id: DocId,
    pub segment_id: SegmentId,
    pub score: f32,
    pub source: Option<serde_json::Value>,
    /// Field values retrieved from the columnar store via the `fields`
    /// parameter. Each value is an array (ES compat). None when `fields`
    /// is not requested. See [[feature-search-fields-retrieval]].
    pub fields: Option<serde_json::Map<String, serde_json::Value>>,
    /// Sort values used for ordering this hit. Present when `sort` is
    /// specified in the request. See [[feature-sort-by-field]].
    pub sort_values: Option<Vec<crate::search::SortValue>>,
    /// Score explanation tree. Present when `explain: true`.
    /// See [[feature-search-explain]].
    pub explanation: Option<crate::search::Explanation>,
    /// Inner hits from nested queries. See [[feature-search-inner-hits]].
    pub inner_hits: Option<std::collections::HashMap<String, serde_json::Value>>,
}

/// Search results from a query.
#[derive(Clone, Debug)]
pub struct SearchResults {
    pub hits: Vec<SearchHit>,
    pub total_hits: crate::search::TotalHits,
    pub aggregations: std::collections::HashMap<String, crate::agg::AggResult>,
}

/// Searches across multiple segment readers.
pub struct Searcher {
    segments: Vec<SegmentReader>,
    analyzers: AnalyzerRegistry,
    /// Cached average field lengths, keyed by field name.
    /// Computed once on construction, reused across all weight creations.
    avg_field_lengths: std::collections::HashMap<String, f32>,
}

impl Searcher {
    pub fn new(segments: Vec<SegmentReader>) -> Self {
        let avg_field_lengths = Self::compute_avg_field_lengths(&segments);
        Self {
            segments,
            analyzers: AnalyzerRegistry::new(),
            avg_field_lengths,
        }
    }

    /// Create a searcher with a custom analyzer registry.
    pub fn with_analyzers(segments: Vec<SegmentReader>, analyzers: AnalyzerRegistry) -> Self {
        let avg_field_lengths = Self::compute_avg_field_lengths(&segments);
        Self { segments, analyzers, avg_field_lengths }
    }

    /// Precompute average field lengths for all indexed fields.
    fn compute_avg_field_lengths(segments: &[SegmentReader]) -> std::collections::HashMap<String, f32> {
        // Collect all field names from all segments
        let mut field_names: std::collections::HashSet<String> = std::collections::HashSet::new();
        for segment in segments {
            for field in &segment.header().fields {
                field_names.insert(field.field_name.clone());
            }
        }

        let mut result = std::collections::HashMap::new();
        for field_name in &field_names {
            let mut total_length = 0.0f64;
            let mut total_docs = 0u32;
            for segment in segments {
                let field_id = segment.header().fields.iter()
                    .find(|f| &f.field_name == field_name)
                    .map(|f| f.field_id);
                if let Some(fid) = field_id {
                    if let Some(norms) = segment.norms(fid) {
                        for i in 0..norms.doc_count() {
                            total_length += norms.norm(DocId::new(i)) as f64;
                        }
                        total_docs += norms.doc_count();
                    }
                }
            }
            if total_docs > 0 {
                result.insert(field_name.clone(), (total_length / total_docs as f64) as f32);
            }
        }
        result
    }

    /// Access the analyzer registry.
    pub fn analyzers(&self) -> &AnalyzerRegistry {
        &self.analyzers
    }

    /// Total number of documents across all segments.
    pub fn total_docs(&self) -> u32 {
        self.segments.iter().map(|s| s.doc_count()).sum()
    }

    /// Number of segments.
    pub fn num_segments(&self) -> usize {
        self.segments.len()
    }

    /// Access the segment readers.
    pub fn segments(&self) -> &[SegmentReader] {
        &self.segments
    }

    /// Consume the searcher and return the segment readers.
    pub fn into_segments(self) -> Vec<SegmentReader> {
        self.segments
    }

    /// Total doc frequency for a term across all segments.
    pub fn doc_freq(&self, field_name: &str, term: &str) -> u32 {
        self.segments.iter().map(|s| {
            let field_id = s.header().fields.iter()
                .find(|f| f.field_name == field_name)
                .map(|f| f.field_id);
            match field_id {
                Some(fid) => s.doc_freq(fid, term),
                None => 0,
            }
        }).sum()
    }

    /// Average field length across all segments (for BM25).
    ///
    /// Returns the precomputed value from construction time. O(1).
    pub fn avg_field_length(&self, field_name: &str) -> f32 {
        self.avg_field_lengths.get(field_name).copied().unwrap_or(0.0)
    }

    /// Search using a Query object, returning top-K results with scoring.
    ///
    /// This is the primary search API. Accepts any `Query` implementation
    /// and executes the `Query → Weight → Scorer` pipeline per segment.
    /// Segments are searched in parallel via rayon when there are 2+ segments.
    pub fn search_query(
        &self,
        query: &dyn crate::query::Query,
        top_k: usize,
        from: usize,
        source_filter: &crate::search::SourceFilter,
        requested_fields: &Option<Vec<String>>,
        explain: bool,
        inner_hit_specs: &[crate::query::nested::InnerHitSpec],
    ) -> SearchResults {
        use luci_core::ScoreMode;

        let weight = match query.create_weight(self, ScoreMode::Complete) {
            Ok(w) => w,
            Err(_) => return SearchResults { hits: Vec::new(), total_hits: crate::search::TotalHits::exact(0), aggregations: std::collections::HashMap::new() },
        };

        // Single-segment fast path: skip rayon overhead
        if self.segments.len() <= 1 {
            let mut collector = TopDocsCollector::new(top_k);
            let mut total_hits: u64 = 0;
            if let Some(segment) = self.segments.first() {
                total_hits = Self::search_segment(segment, weight.as_ref(), &mut collector, None);
            }
            let ew = if explain { Some(weight.as_ref()) } else { None };
            let mut results = self.build_results(collector, total_hits, from, source_filter, requested_fields, ew);
            if !inner_hit_specs.is_empty() {
                self.resolve_inner_hits(&mut results, inner_hit_specs);
            }
            return results;
        }

        // Multi-segment parallel search
        let shared_min = SharedMinScore::new();

        let partial_results: Vec<(TopDocsCollector, u64)> = self.segments
            .par_iter()
            .map(|segment| {
                let mut collector = TopDocsCollector::new(top_k);
                let hits = Self::search_segment(segment, weight.as_ref(), &mut collector, Some(&shared_min));
                // Propagate local min back to global
                shared_min.update(collector.min_score());
                (collector, hits)
            })
            .collect();

        // Merge thread-local results
        let mut merged = TopDocsCollector::new(top_k);
        let mut total_hits: u64 = 0;
        for (collector, hits) in partial_results {
            merged.merge(collector);
            total_hits += hits;
        }

        let ew = if explain { Some(weight.as_ref()) } else { None };
        let mut results = self.build_results(merged, total_hits, from, source_filter, requested_fields, ew);
        if !inner_hit_specs.is_empty() {
            self.resolve_inner_hits(&mut results, inner_hit_specs);
        }
        results
    }

    /// Search with field collapse (deduplication by field value).
    ///
    /// Keeps only the highest-scoring doc per unique collapse field value.
    /// See [[feature-search-collapse]].
    pub fn search_query_collapsed(
        &self,
        query: &dyn crate::query::Query,
        top_k: usize,
        from: usize,
        source_filter: &crate::search::SourceFilter,
        requested_fields: &Option<Vec<String>>,
        collapse_field: &str,
        explain: bool,
    ) -> SearchResults {
        use crate::search::collector::{CollapsingCollector, CollapseKey};
        use luci_core::ScoreMode;

        let weight = match query.create_weight(self, ScoreMode::Complete) {
            Ok(w) => w,
            Err(_) => return SearchResults { hits: Vec::new(), total_hits: crate::search::TotalHits::exact(0), aggregations: std::collections::HashMap::new() },
        };

        let search_one = |segment: &SegmentReader| -> (CollapsingCollector, u64) {
            let mut collector = CollapsingCollector::new(top_k);
            let seg_id = segment.segment_id();

            // Resolve collapse field
            let field_id = segment.header().fields.iter()
                .find(|f| f.field_name == collapse_field)
                .map(|f| f.field_id);
            let col_reader = field_id.and_then(|fid| segment.column(fid));

            let supplier = match weight.scorer_supplier(segment) {
                Ok(Some(s)) => s,
                _ => return (collector, 0),
            };
            let mut scorer = match supplier.scorer() {
                Ok(s) => s,
                Err(_) => return (collector, 0),
            };

            let mut total_hits: u64 = 0;
            loop {
                let doc = scorer.doc_id();
                if doc == NO_MORE_DOCS { break; }
                let score = scorer.score();

                let key = match &col_reader {
                    Some(col) => {
                        use crate::columnar::writer::ColumnType;
                        let d = doc.as_u32();
                        match col.col_type() {
                            ColumnType::Keyword => col.keyword_value(d)
                                .map(|s| CollapseKey::Keyword(s.to_string()))
                                .unwrap_or(CollapseKey::Null),
                            ColumnType::I64 | ColumnType::BitpackedI64 | ColumnType::ConstantI64 => col.i64_value(d)
                                .map(CollapseKey::Numeric)
                                .unwrap_or(CollapseKey::Null),
                            ColumnType::F64 | ColumnType::ConstantF64 => col.f64_value(d)
                                .map(|f| CollapseKey::Numeric(f.to_bits() as i64))
                                .unwrap_or(CollapseKey::Null),
                            _ => CollapseKey::Null,
                        }
                    }
                    None => CollapseKey::Null,
                };

                collector.collect(doc, seg_id, score, key);
                total_hits += 1;
                scorer.next();
            }

            (collector, total_hits)
        };

        let segment_results: Vec<(CollapsingCollector, u64)> = if self.segments.len() <= 1 {
            self.segments.iter().map(search_one).collect()
        } else {
            self.segments.par_iter().map(search_one).collect()
        };

        let mut merged = CollapsingCollector::new(top_k);
        let mut total_hits: u64 = 0;
        for (collector, hits) in segment_results {
            merged.merge(collector);
            total_hits += hits;
        }

        let collapsed_docs = merged.into_sorted_results();
        let ew = if explain { Some(weight.as_ref()) } else { None };
        let hits = collapsed_docs
            .into_iter()
            .skip(from)
            .map(|cd| {
                let source = match source_filter {
                    crate::search::SourceFilter::Disabled => None,
                    crate::search::SourceFilter::Enabled => self.get_source(cd.segment_id, cd.doc_id),
                    filter => {
                        self.get_source(cd.segment_id, cd.doc_id)
                            .and_then(|v| crate::search::filter_source(&v, filter))
                    }
                };
                let fields = requested_fields.as_ref().map(|names| {
                    let mut f = self.retrieve_fields(cd.segment_id, cd.doc_id, names);
                    // Add collapse field value
                    f.insert(collapse_field.to_string(), serde_json::json!([cd.collapse_key.to_json()]));
                    f
                }).or_else(|| {
                    let mut f = serde_json::Map::new();
                    f.insert(collapse_field.to_string(), serde_json::json!([cd.collapse_key.to_json()]));
                    Some(f)
                });
                let explanation = ew.and_then(|w| {
                    let segment = self.segments.iter().find(|s| s.segment_id() == cd.segment_id)?;
                    w.explain(segment, cd.doc_id).ok()
                });
                SearchHit {
                    doc_id: cd.doc_id,
                    segment_id: cd.segment_id,
                    score: cd.score,
                    source,
                    fields,
                    sort_values: None,
                    explanation,
                    inner_hits: None,
                }
            })
            .collect();

        SearchResults { hits, total_hits: crate::search::TotalHits::exact(total_hits), aggregations: std::collections::HashMap::new() }
    }

    /// Search with field-based sorting.
    ///
    /// Uses `TopFieldCollector` instead of `TopDocsCollector`. Reads sort
    /// values from `ColumnReader` per segment. When sort doesn't include
    /// `_score`, scoring is skipped entirely.
    /// See [[feature-sort-by-field]].
    pub fn search_query_sorted(
        &self,
        query: &dyn crate::query::Query,
        top_k: usize,
        from: usize,
        source_filter: &crate::search::SourceFilter,
        requested_fields: &Option<Vec<String>>,
        sort_fields: &[crate::search::SortField],
        search_after: &Option<Vec<crate::search::SortValue>>,
    ) -> SearchResults {
        use crate::search::{SortField, SortFieldType, SortValue};
        use crate::search::collector::TopFieldCollector;
        use luci_core::ScoreMode;

        let needs_scores = sort_fields.iter().any(|sf| matches!(sf.field, SortFieldType::Score));
        let score_mode = if needs_scores { ScoreMode::Complete } else { ScoreMode::CompleteNoScores };

        let weight = match query.create_weight(self, score_mode) {
            Ok(w) => w,
            Err(_) => return SearchResults { hits: Vec::new(), total_hits: crate::search::TotalHits::exact(0), aggregations: std::collections::HashMap::new() },
        };

        let search_one = |segment: &SegmentReader| -> (TopFieldCollector, u64) {
            let mut collector = TopFieldCollector::new(top_k, sort_fields.to_vec());
            if let Some(after) = search_after {
                collector.set_search_after(after.clone());
            }
            let seg_id = segment.segment_id();

            let supplier = match weight.scorer_supplier(segment) {
                Ok(Some(s)) => s,
                _ => return (collector, 0),
            };
            let mut scorer = match supplier.scorer() {
                Ok(s) => s,
                Err(_) => return (collector, 0),
            };

            // Open ColumnReaders for sort fields
            let col_readers: Vec<Option<crate::columnar::reader::ColumnReader>> = sort_fields.iter().map(|sf| {
                match &sf.field {
                    SortFieldType::Field(name) => {
                        let field_id = segment.header().fields.iter()
                            .find(|f| f.field_name == *name)
                            .map(|f| f.field_id)?;
                        segment.column(field_id)
                    }
                    _ => None,
                }
            }).collect();

            // Helper: read a sort value for field index `i` from columns
            let read_sort_value = |i: usize, sf: &SortField, doc: DocId, score: f32| -> SortValue {
                match &sf.field {
                    SortFieldType::Score => SortValue::Score(score),
                    SortFieldType::Doc => SortValue::Doc(seg_id.as_u64() * 1_000_000_000 + doc.as_u32() as u64),
                    SortFieldType::Field(_) => {
                        match &col_readers[i] {
                            Some(col) => {
                                use crate::columnar::writer::ColumnType;
                                let d = doc.as_u32();
                                match col.col_type() {
                                    ColumnType::Keyword => col.keyword_value(d)
                                        .map(|s| SortValue::Str(s.to_string()))
                                        .unwrap_or(SortValue::Null),
                                    ColumnType::F64 | ColumnType::ConstantF64 => col.f64_value(d)
                                        .map(SortValue::F64)
                                        .unwrap_or(SortValue::Null),
                                    ColumnType::I64 | ColumnType::BitpackedI64 | ColumnType::ConstantI64 => col.i64_value(d)
                                        .map(SortValue::I64)
                                        .unwrap_or(SortValue::Null),
                                    ColumnType::Bool => col.bool_value(d)
                                        .map(SortValue::Bool)
                                        .unwrap_or(SortValue::Null),
                                    _ => SortValue::Null,
                                }
                            }
                            None => SortValue::Null,
                        }
                    }
                }
            };

            let mut total_hits: u64 = 0;
            loop {
                let doc = scorer.doc_id();
                if doc == NO_MORE_DOCS { break; }

                let score = if needs_scores { scorer.score() } else { 0.0 };

                // Early rejection: check competitiveness before allocating.
                // For keyword fields, compare &str directly to avoid String alloc.
                let primary_competitive = match &sort_fields[0].field {
                    SortFieldType::Field(_) => {
                        match &col_readers[0] {
                            Some(col) if col.col_type() == crate::columnar::writer::ColumnType::Keyword => {
                                let kv = col.keyword_value(doc.as_u32());
                                collector.is_competitive_keyword(kv)
                            }
                            _ => {
                                let primary = read_sort_value(0, &sort_fields[0], doc, score);
                                collector.is_competitive_primary(&primary)
                            }
                        }
                    }
                    _ => {
                        let primary = read_sort_value(0, &sort_fields[0], doc, score);
                        collector.is_competitive_primary(&primary)
                    }
                };
                if !primary_competitive {
                    total_hits += 1;
                    scorer.next();
                    continue;
                }

                // Competitive: build full sort values (re-read primary)
                let mut sort_values = Vec::with_capacity(sort_fields.len());
                sort_values.push(read_sort_value(0, &sort_fields[0], doc, score));
                for i in 1..sort_fields.len() {
                    sort_values.push(read_sort_value(i, &sort_fields[i], doc, score));
                }

                collector.collect(doc, seg_id, score, sort_values);
                total_hits += 1;
                scorer.next();
            }

            (collector, total_hits)
        };

        let segment_results: Vec<(TopFieldCollector, u64)> = if self.segments.len() <= 1 {
            self.segments.iter().map(search_one).collect()
        } else {
            self.segments.par_iter().map(search_one).collect()
        };

        let mut merged = TopFieldCollector::new(top_k, sort_fields.to_vec());
        let mut total_hits: u64 = 0;
        for (collector, hits) in segment_results {
            merged.merge(collector);
            total_hits += hits;
        }

        // Build results from sorted collector
        let field_docs = merged.into_sorted_results();
        let hits = field_docs
            .into_iter()
            .skip(from)
            .map(|fd| {
                let source = match source_filter {
                    crate::search::SourceFilter::Disabled => None,
                    crate::search::SourceFilter::Enabled => self.get_source(fd.segment_id, fd.doc_id),
                    filter => {
                        self.get_source(fd.segment_id, fd.doc_id)
                            .and_then(|v| crate::search::filter_source(&v, filter))
                    }
                };
                let fields = requested_fields.as_ref().map(|names| {
                    self.retrieve_fields(fd.segment_id, fd.doc_id, names)
                });
                SearchHit {
                    doc_id: fd.doc_id,
                    segment_id: fd.segment_id,
                    score: fd.score,
                    source,
                    fields,
                    sort_values: Some(fd.sort_values),
                    explanation: None,
                    inner_hits: None,
                }
            })
            .collect();

        SearchResults { hits, total_hits: crate::search::TotalHits::exact(total_hits), aggregations: std::collections::HashMap::new() }
    }

    /// Search a single segment, collecting into a thread-local collector.
    fn search_segment(
        segment: &SegmentReader,
        weight: &dyn crate::query::Weight,
        collector: &mut TopDocsCollector,
        shared_min: Option<&SharedMinScore>,
    ) -> u64 {
        let seg_id = segment.segment_id();
        let mut total_hits: u64 = 0;

        // Try bulk scoring first (bypasses doc-at-a-time overhead)
        if let Ok(Some(hits)) = weight.bulk_score(segment, collector, seg_id) {
            if let Some(sm) = shared_min {
                sm.update(collector.min_score());
            }
            return hits;
        }

        let supplier = match weight.scorer_supplier(segment) {
            Ok(Some(s)) => s,
            _ => return 0,
        };

        let mut scorer = match supplier.scorer() {
            Ok(s) => s,
            Err(_) => return 0,
        };

        // Seed scorer with global min competitive score from other threads
        if let Some(sm) = shared_min {
            let global_min = sm.get();
            if global_min > 0.0 {
                scorer.set_min_competitive_score(global_min);
            }
        }

        // Hot loop: score + collect each matching document.
        // Only call set_min_competitive_score when the threshold actually
        // changes, avoiding 99.99% wasted indirect calls.
        // See [[optimization-collector-scaling]].
        let mut last_min: f32 = 0.0;
        loop {
            let doc = scorer.doc_id();
            if doc == NO_MORE_DOCS {
                break;
            }
            let score = scorer.score();
            collector.collect(doc, seg_id, score);
            total_hits += 1;
            let min = collector.min_score();
            if min > last_min {
                scorer.set_min_competitive_score(min);
                last_min = min;
            }
            scorer.next();
        }

        total_hits
    }

    /// Build SearchResults from a collector and total hit count.
    ///
    /// Skips the first `from` results and only retrieves `_source` for
    /// the remaining docs. This makes `from`-based pagination nearly
    /// free — source retrieval (~2.5μs/doc) is the dominant per-result
    /// cost, and we only pay it for the final page.
    /// See [[feature-search-from-offset]].
    fn build_results(
        &self,
        collector: TopDocsCollector,
        total_hits: u64,
        from: usize,
        source_filter: &crate::search::SourceFilter,
        requested_fields: &Option<Vec<String>>,
        explain_weight: Option<&dyn crate::query::Weight>,
    ) -> SearchResults {
        use crate::search::SourceFilter;
        let score_docs = collector.into_sorted_results();
        let hits = score_docs
            .into_iter()
            .skip(from)
            .map(|sd| {
                let source = match source_filter {
                    SourceFilter::Disabled => None,
                    SourceFilter::Enabled => self.get_source(sd.segment_id, sd.doc_id),
                    filter => {
                        self.get_source(sd.segment_id, sd.doc_id)
                            .and_then(|v| crate::search::filter_source(&v, filter))
                    }
                };
                let fields = requested_fields.as_ref().map(|names| {
                    self.retrieve_fields(sd.segment_id, sd.doc_id, names)
                });
                let explanation = explain_weight.and_then(|w| {
                    let segment = self.segments.iter().find(|s| s.segment_id() == sd.segment_id)?;
                    w.explain(segment, sd.doc_id).ok()
                });
                SearchHit {
                    doc_id: sd.doc_id,
                    segment_id: sd.segment_id,
                    score: sd.score,
                    source,
                    fields,
                    sort_values: None,
                    explanation,
                    inner_hits: None,
                }
            })
            .collect();

        SearchResults { hits, total_hits: crate::search::TotalHits::exact(total_hits), aggregations: std::collections::HashMap::new() }
    }

    /// Search with query and aggregations.
    ///
    /// Segments are searched in parallel via rayon when there are 2+ segments.
    pub fn search_with_aggs(
        &self,
        query: &dyn crate::query::Query,
        agg_factories: &[(String, Box<dyn crate::agg::AggCollectorFactory>)],
        top_k: usize,
        from: usize,
        source_filter: &crate::search::SourceFilter,
        requested_fields: &Option<Vec<String>>,
    ) -> SearchResults {
        use luci_core::ScoreMode;

        let weight = match query.create_weight(self, ScoreMode::Complete) {
            Ok(w) => w,
            Err(_) => return SearchResults { hits: Vec::new(), total_hits: crate::search::TotalHits::exact(0), aggregations: std::collections::HashMap::new() },
        };

        // Per-segment search + agg collection (runs in parallel for multi-segment)
        let search_one = |segment: &SegmentReader| -> (TopDocsCollector, u64, Vec<crate::agg::AggResult>) {
            let mut collector = TopDocsCollector::new(top_k);
            let mut total_hits: u64 = 0;

            let mut agg_collectors: Vec<Box<dyn crate::agg::AggCollector>> = agg_factories
                .iter()
                .map(|(_, factory)| factory.create_collector(segment))
                .collect();

            let seg_id = segment.segment_id();

            // Match-all + agg-only: bypass scorer, iterate raw doc IDs
            if top_k == 0 && weight.is_match_all() {
                let doc_count = segment.doc_count();
                total_hits = doc_count as u64;
                for i in 0..doc_count {
                    let doc_id = DocId::new(i);
                    for agg_collector in &mut agg_collectors {
                        agg_collector.collect(doc_id);
                    }
                }
                let agg_results = agg_collectors.into_iter().map(|c| c.finish()).collect();
                return (collector, total_hits, agg_results);
            }

            let supplier = match weight.scorer_supplier(segment) {
                Ok(Some(s)) => s,
                _ => {
                    let agg_results = agg_collectors.into_iter().map(|c| c.finish()).collect();
                    return (collector, 0, agg_results);
                }
            };

            let mut scorer = match supplier.scorer() {
                Ok(s) => s,
                Err(_) => {
                    let agg_results = agg_collectors.into_iter().map(|c| c.finish()).collect();
                    return (collector, 0, agg_results);
                }
            };

            if top_k == 0 {
                loop {
                    let doc_id = scorer.doc_id();
                    if doc_id == NO_MORE_DOCS { break; }
                    total_hits += 1;
                    for agg_collector in &mut agg_collectors {
                        agg_collector.collect(doc_id);
                    }
                    scorer.next();
                }
            } else {
                loop {
                    let doc_id = scorer.doc_id();
                    if doc_id == NO_MORE_DOCS { break; }
                    let score = scorer.score();
                    collector.collect(doc_id, seg_id, score);
                    total_hits += 1;
                    for agg_collector in &mut agg_collectors {
                        agg_collector.collect(doc_id);
                    }
                    scorer.set_min_competitive_score(collector.min_score());
                    scorer.next();
                }
            }

            let agg_results = agg_collectors.into_iter().map(|c| c.finish()).collect();
            (collector, total_hits, agg_results)
        };

        // Execute: parallel for multi-segment, sequential for single
        let segment_results: Vec<(TopDocsCollector, u64, Vec<crate::agg::AggResult>)> =
            if self.segments.len() <= 1 {
                self.segments.iter().map(search_one).collect()
            } else {
                self.segments.par_iter().map(search_one).collect()
            };

        // Merge collectors and agg results
        let mut merged = TopDocsCollector::new(top_k);
        let mut total_hits: u64 = 0;
        let num_aggs = agg_factories.len();
        let mut agg_partial_results: Vec<Vec<crate::agg::AggResult>> =
            (0..num_aggs).map(|_| Vec::new()).collect();

        for (collector, hits, agg_results) in segment_results {
            merged.merge(collector);
            total_hits += hits;
            for (i, result) in agg_results.into_iter().enumerate() {
                if i < num_aggs {
                    agg_partial_results[i].push(result);
                }
            }
        }

        let mut aggregations = std::collections::HashMap::new();
        for (i, (name, factory)) in agg_factories.iter().enumerate() {
            let partial = std::mem::take(&mut agg_partial_results[i]);
            let merged_agg = factory.merge_results(partial);
            aggregations.insert(name.clone(), merged_agg);
        }

        let score_docs = merged.into_sorted_results();
        let hits = score_docs
            .into_iter()
            .skip(from)
            .map(|sd| {
                let source = match source_filter {
                    crate::search::SourceFilter::Disabled => None,
                    crate::search::SourceFilter::Enabled => self.get_source(sd.segment_id, sd.doc_id),
                    filter => {
                        self.get_source(sd.segment_id, sd.doc_id)
                            .and_then(|v| crate::search::filter_source(&v, filter))
                    }
                };
                let fields = requested_fields.as_ref().map(|names| {
                    self.retrieve_fields(sd.segment_id, sd.doc_id, names)
                });
                SearchHit {
                    doc_id: sd.doc_id,
                    segment_id: sd.segment_id,
                    score: sd.score,
                    source,
                    fields,
                    sort_values: None,
                    explanation: None,
                    inner_hits: None,
                }
            })
            .collect();

        SearchResults { hits, total_hits: crate::search::TotalHits::exact(total_hits), aggregations }
    }

    /// Search for a term in a field, returning the top-K results with BM25 scoring.
    ///
    /// Segments are searched in parallel via rayon when there are 2+ segments.
    pub fn search_term(&self, field_id: FieldId, term: &str, top_k: usize) -> SearchResults {
        // Compute index-level statistics for BM25 (sequential, fast)
        let total_docs = self.total_docs();
        let doc_freq: u32 = self.segments.iter().map(|s| s.doc_freq(field_id, term)).sum();

        if doc_freq == 0 {
            return SearchResults {
                hits: Vec::new(),
                total_hits: crate::search::TotalHits::exact(0),
                aggregations: std::collections::HashMap::new(),
            };
        }

        let avg_field_length = self.avg_field_length_by_id(field_id);
        let weight = Bm25Weight::new(total_docs, doc_freq, avg_field_length);

        let search_one = |segment: &SegmentReader| -> (TopDocsCollector, u64) {
            let mut collector = TopDocsCollector::new(top_k);
            let mut total_hits: u64 = 0;

            let postings = match segment.postings(field_id, term) {
                Some(p) => p,
                None => return (collector, 0),
            };
            let norms = match segment.norms(field_id) {
                Some(n) => n,
                None => return (collector, 0),
            };

            let seg_weight = weight.clone();
            let mut scorer = Bm25Scorer::new(seg_weight, postings, norms);
            let seg_id = segment.segment_id();

            while scorer.doc_id() != NO_MORE_DOCS {
                let score = scorer.score();
                collector.collect(scorer.doc_id(), seg_id, score);
                total_hits += 1;
                scorer.next();
            }

            (collector, total_hits)
        };

        let segment_results: Vec<(TopDocsCollector, u64)> = if self.segments.len() <= 1 {
            self.segments.iter().map(search_one).collect()
        } else {
            self.segments.par_iter().map(search_one).collect()
        };

        let mut merged = TopDocsCollector::new(top_k);
        let mut total_hits: u64 = 0;
        for (collector, hits) in segment_results {
            merged.merge(collector);
            total_hits += hits;
        }

        self.build_results(merged, total_hits, 0, &crate::search::SourceFilter::Enabled, &None, None)
    }

    /// Compute average field length by field ID (internal helper).
    fn avg_field_length_by_id(&self, field_id: FieldId) -> f32 {
        let mut total_length = 0.0f64;
        let mut total_docs = 0u32;

        for segment in &self.segments {
            if let Some(norms) = segment.norms(field_id) {
                for i in 0..norms.doc_count() {
                    total_length += norms.norm(DocId::new(i)) as f64;
                }
                total_docs += norms.doc_count();
            }
        }

        if total_docs == 0 {
            return 0.0;
        }
        (total_length / total_docs as f64) as f32
    }

    /// Apply rescore: re-score the top hits with an additional query and
    /// combine scores. See [[feature-search-rescore]].
    pub fn apply_rescore(
        &self,
        results: &mut SearchResults,
        rescore_query: &dyn crate::query::Query,
        window_size: usize,
        query_weight: f32,
        rescore_query_weight: f32,
        score_mode: crate::search::RescoreScoreMode,
    ) {
        use luci_core::{ScoreMode, Scorer as _, NO_MORE_DOCS};
        use std::cmp::Ordering;

        let weight = match rescore_query.create_weight(self, ScoreMode::Complete) {
            Ok(w) => w,
            Err(_) => return,
        };

        // Only rescore the top window_size hits
        let rescore_count = results.hits.len().min(window_size);

        // Group hits by segment, process each segment with one scorer
        // advancing forward through sorted doc_ids
        for segment in &self.segments {
            let seg_id = segment.segment_id();

            let supplier = match weight.scorer_supplier(segment) {
                Ok(Some(s)) => s,
                _ => continue,
            };
            let mut scorer = match supplier.scorer() {
                Ok(s) => s,
                Err(_) => continue,
            };

            // Collect indices of hits in this segment, sorted by doc_id
            let mut seg_hits: Vec<usize> = (0..rescore_count)
                .filter(|&i| results.hits[i].segment_id == seg_id)
                .collect();
            seg_hits.sort_by_key(|&i| results.hits[i].doc_id);

            for idx in seg_hits {
                let doc = results.hits[idx].doc_id;

                // Advance scorer forward to this doc
                if scorer.doc_id() == NO_MORE_DOCS {
                    break;
                }
                if scorer.doc_id() < doc {
                    scorer.advance(doc);
                }

                let rescore_score = if scorer.doc_id() == doc {
                    if let Some(tp) = scorer.two_phase() {
                        if tp.matches() { scorer.score() } else { 0.0 }
                    } else {
                        scorer.score()
                    }
                } else {
                    0.0
                };

                results.hits[idx].score = score_mode.combine(
                    results.hits[idx].score, rescore_score, query_weight, rescore_query_weight,
                );
            }
        }

        // Re-sort by new combined score
        results.hits[..rescore_count].sort_by(|a, b| {
            b.score.partial_cmp(&a.score)
                .unwrap_or(Ordering::Equal)
                .then_with(|| a.doc_id.cmp(&b.doc_id))
        });
    }

    /// Resolve inner hits for each result hit that came from a nested query.
    ///
    /// For each hit, re-runs the inner query scoped to that parent's child
    /// block, extracts nested _source from the parent's JSON.
    /// See [[feature-search-inner-hits]].
    fn resolve_inner_hits(
        &self,
        results: &mut SearchResults,
        specs: &[crate::query::nested::InnerHitSpec],
    ) {
        use luci_core::{Scorer as _, NO_MORE_DOCS};
        use std::cmp::Ordering;

        for hit in &mut results.hits {
            let segment = match self.segments.iter().find(|s| s.segment_id() == hit.segment_id) {
                Some(s) => s,
                None => continue,
            };
            let parent_bitset = match segment.parent_bitset() {
                Some(b) => b,
                None => continue,
            };
            let parent_doc = hit.doc_id.as_u32() as usize;

            // Children are stored after parent: [parent, child_0, child_1, ...]
            // Find child range: (parent_doc+1 ... next_parent or end)
            let child_start = parent_doc + 1;
            let child_end = (child_start..parent_bitset.len())
                .find(|&i| parent_bitset[i])
                .unwrap_or(parent_bitset.len());

            if child_start >= child_end {
                continue; // No children for this parent
            }

            // Get parent source for nested _source extraction
            let parent_source = self.get_source(hit.segment_id, hit.doc_id);

            let mut inner_hits_map = std::collections::HashMap::new();

            for spec in specs {
                let supplier = match spec.weight.scorer_supplier(segment) {
                    Ok(Some(s)) => s,
                    _ => continue,
                };
                let mut scorer = match supplier.scorer() {
                    Ok(s) => s,
                    Err(_) => continue,
                };

                // Advance to child range and collect matches
                scorer.advance(DocId::new(child_start as u32));
                let mut matches: Vec<(u32, f32)> = Vec::new();
                while scorer.doc_id() != NO_MORE_DOCS {
                    let doc = scorer.doc_id().as_u32();
                    if doc >= child_end as u32 { break; }
                    matches.push((doc, scorer.score()));
                    scorer.next();
                }

                if matches.is_empty() { continue; }

                // Sort by score desc, apply from/size
                matches.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(Ordering::Equal));
                let total = matches.len();
                let page: Vec<_> = matches.into_iter()
                    .skip(spec.config.from)
                    .take(spec.config.size)
                    .collect();

                // Build inner hits JSON
                let inner_hit_docs: Vec<serde_json::Value> = page.iter().map(|(child_doc, score)| {
                    let offset = (*child_doc as usize).saturating_sub(child_start);
                    let nested_source = parent_source.as_ref()
                        .and_then(|ps| ps.get(&spec.path))
                        .and_then(|arr| arr.as_array())
                        .and_then(|a| a.get(offset))
                        .cloned();
                    serde_json::json!({
                        "_nested": { "field": spec.path, "offset": offset },
                        "_score": score,
                        "_source": nested_source,
                    })
                }).collect();

                inner_hits_map.insert(spec.name.clone(), serde_json::json!({
                    "hits": {
                        "total": { "value": total, "relation": "eq" },
                        "hits": inner_hit_docs,
                    }
                }));
            }

            if !inner_hits_map.is_empty() {
                hit.inner_hits = Some(inner_hits_map);
            }
        }
    }

    /// Retrieve a document's source JSON from its segment.
    fn get_source(&self, segment_id: SegmentId, doc_id: DocId) -> Option<serde_json::Value> {
        let segment = self.segments.iter().find(|s| s.segment_id() == segment_id)?;
        let store = segment.doc_store();
        let bytes = store.get(doc_id.as_u32())?;
        serde_json::from_slice(&bytes).ok()
    }

    /// Retrieve typed field values from the columnar store.
    ///
    /// Each value is wrapped in a JSON array (ES `fields` API compat).
    /// Fields with null values or missing columns (e.g. text fields) are
    /// omitted. See [[feature-search-fields-retrieval]].
    fn retrieve_fields(
        &self,
        segment_id: SegmentId,
        doc_id: DocId,
        field_names: &[String],
    ) -> serde_json::Map<String, serde_json::Value> {
        let segment = match self.segments.iter().find(|s| s.segment_id() == segment_id) {
            Some(s) => s,
            None => return serde_json::Map::new(),
        };

        let doc = doc_id.as_u32();
        let mut result = serde_json::Map::new();

        for name in field_names {
            let field_id = match segment.header().fields.iter()
                .find(|f| f.field_name == *name)
                .map(|f| f.field_id)
            {
                Some(id) => id,
                None => continue,
            };

            let col = match segment.column(field_id) {
                Some(c) => c,
                None => continue,
            };

            use crate::columnar::writer::ColumnType;
            let value = match col.col_type() {
                ColumnType::Keyword => match col.keyword_value(doc) {
                    Some(s) => serde_json::json!([s]),
                    None => continue,
                },
                ColumnType::F64 | ColumnType::ConstantF64 => match col.f64_value(doc) {
                    Some(n) => serde_json::json!([n]),
                    None => continue,
                },
                ColumnType::I64 | ColumnType::BitpackedI64 | ColumnType::ConstantI64 => match col.i64_value(doc) {
                    Some(n) => serde_json::json!([n]),
                    None => continue,
                },
                ColumnType::Bool => match col.bool_value(doc) {
                    Some(b) => serde_json::json!([b]),
                    None => continue,
                },
                _ => continue,
            };

            result.insert(name.clone(), value);
        }

        result
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::segment::builder::SegmentBuilder;
    use luci_analysis::Token;
    use luci_mapping::{FieldType,Mapping};

    fn make_tokens(terms: &[&str]) -> Vec<Token> {
        terms
            .iter()
            .enumerate()
            .map(|(i, t)| Token::new(*t, 0, t.len(), i as u32))
            .collect()
    }

    fn test_schema() -> Mapping {
        Mapping::builder()
            .field("body", FieldType::Text)
            .build()
    }

    fn build_segment(id: u64, docs: &[(&[&str], &str)]) -> SegmentReader {
        let schema = test_schema();
        let mut builder = SegmentBuilder::new(SegmentId::new(id), &schema);
        for (terms, source) in docs {
            builder.add_document(
                &[(FieldId::new(0), make_tokens(terms))],
                source.as_bytes(),
            );
        }
        SegmentReader::open(builder.build()).unwrap()
    }

    #[test]
    fn single_segment_search() {
        let reader = build_segment(
            1,
            &[
                (&["hello", "world"], r#"{"body":"hello world"}"#),
                (&["hello", "luci"], r#"{"body":"hello luci"}"#),
                (&["goodbye", "world"], r#"{"body":"goodbye world"}"#),
            ],
        );

        let searcher = Searcher::new(vec![reader]);
        let results = searcher.search_term(FieldId::new(0), "hello", 10);

        assert_eq!(results.total_hits.value, 2);
        assert_eq!(results.hits.len(), 2);
        // Both results should have "hello" in the source
        for hit in &results.hits {
            let source = hit.source.as_ref().unwrap();
            assert!(source["body"].as_str().unwrap().contains("hello"));
        }
    }

    #[test]
    fn multi_segment_search() {
        let seg1 = build_segment(
            1,
            &[
                (&["hello", "world"], r#"{"body":"hello world"}"#),
                (&["foo", "bar"], r#"{"body":"foo bar"}"#),
            ],
        );
        let seg2 = build_segment(
            2,
            &[
                (&["hello", "luci"], r#"{"body":"hello luci"}"#),
                (&["baz", "qux"], r#"{"body":"baz qux"}"#),
            ],
        );

        let searcher = Searcher::new(vec![seg1, seg2]);
        let results = searcher.search_term(FieldId::new(0), "hello", 10);

        assert_eq!(results.total_hits.value, 2);
        assert_eq!(results.hits.len(), 2);
    }

    #[test]
    fn correct_bm25_ranking() {
        // Doc with higher TF should rank higher (same field length)
        let reader = build_segment(
            1,
            &[
                (&["search"], r#"{"body":"search"}"#),
                (&["search", "search", "search"], r#"{"body":"search search search"}"#),
                (&["search", "other"], r#"{"body":"search other"}"#),
            ],
        );

        let searcher = Searcher::new(vec![reader]);
        let results = searcher.search_term(FieldId::new(0), "search", 10);

        assert_eq!(results.total_hits.value, 3);
        // Doc 1 (tf=3) should score highest
        assert_eq!(results.hits[0].doc_id, DocId::new(1));
    }

    #[test]
    fn top_k_limiting() {
        let reader = build_segment(
            1,
            &[
                (&["term"], r#"{"id":0}"#),
                (&["term"], r#"{"id":1}"#),
                (&["term"], r#"{"id":2}"#),
                (&["term"], r#"{"id":3}"#),
                (&["term"], r#"{"id":4}"#),
            ],
        );

        let searcher = Searcher::new(vec![reader]);
        let results = searcher.search_term(FieldId::new(0), "term", 2);

        assert_eq!(results.total_hits.value, 5);
        assert_eq!(results.hits.len(), 2);
    }

    #[test]
    fn score_descending_order() {
        let reader = build_segment(
            1,
            &[
                (&["a", "b", "c", "d", "search"], r#"{"id":0}"#),
                (&["search", "search"], r#"{"id":1}"#),
                (&["search"], r#"{"id":2}"#),
            ],
        );

        let searcher = Searcher::new(vec![reader]);
        let results = searcher.search_term(FieldId::new(0), "search", 10);

        for i in 0..results.hits.len() - 1 {
            assert!(
                results.hits[i].score >= results.hits[i + 1].score,
                "results should be in descending score order"
            );
        }
    }

    #[test]
    fn source_returned() {
        let reader = build_segment(
            1,
            &[(&["hello"], r#"{"body":"hello"}"#)],
        );

        let searcher = Searcher::new(vec![reader]);
        let results = searcher.search_term(FieldId::new(0), "hello", 10);

        assert_eq!(results.hits.len(), 1);
        let source = results.hits[0].source.as_ref().unwrap();
        assert_eq!(source["body"], "hello");
    }

    #[test]
    fn missing_term() {
        let reader = build_segment(
            1,
            &[(&["hello"], r#"{"body":"hello"}"#)],
        );

        let searcher = Searcher::new(vec![reader]);
        let results = searcher.search_term(FieldId::new(0), "nonexistent", 10);

        assert_eq!(results.total_hits.value, 0);
        assert!(results.hits.is_empty());
    }

    #[test]
    fn missing_field() {
        let reader = build_segment(
            1,
            &[(&["hello"], r#"{"body":"hello"}"#)],
        );

        let searcher = Searcher::new(vec![reader]);
        let results = searcher.search_term(FieldId::new(99), "hello", 10);

        assert_eq!(results.total_hits.value, 0);
        assert!(results.hits.is_empty());
    }

    #[test]
    fn empty_index() {
        let searcher = Searcher::new(vec![]);
        let results = searcher.search_term(FieldId::new(0), "hello", 10);
        assert_eq!(results.total_hits.value, 0);
        assert!(results.hits.is_empty());
    }
}
