//! Top-K document collector with min-heap.
//!
//! Collects the highest-scoring documents up to a configurable limit,
//! using a min-heap to efficiently maintain the top-K set.
//!
//! See [[architecture-query-execution#Aggregation Collection]].

use std::cmp::Ordering;
use std::collections::BinaryHeap;

use luci_core::{DocId, SegmentId};

/// A scored document result (internal to the search pipeline).
#[derive(Clone, Debug)]
pub(crate) struct ScoreDoc {
    pub(crate) doc_id: DocId,
    pub(crate) segment_id: SegmentId,
    pub(crate) score: f32,
}

/// Wrapper for min-heap ordering (lowest score at top).
struct MinScoreDoc(ScoreDoc);

impl PartialEq for MinScoreDoc {
    fn eq(&self, other: &Self) -> bool {
        self.0.score == other.0.score
    }
}

impl Eq for MinScoreDoc {}

impl PartialOrd for MinScoreDoc {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for MinScoreDoc {
    fn cmp(&self, other: &Self) -> Ordering {
        // Reverse ordering: we want lowest score at the top of the BinaryHeap
        // (BinaryHeap is a max-heap, so reversing makes it a min-heap).
        other
            .0
            .score
            .partial_cmp(&self.0.score)
            .unwrap_or(Ordering::Equal)
            .then_with(|| other.0.doc_id.cmp(&self.0.doc_id))
    }
}

/// Collects the top-K documents by score.
pub struct TopDocsCollector {
    heap: BinaryHeap<MinScoreDoc>,
    limit: usize,
}

impl TopDocsCollector {
    /// Create a collector that keeps the top `limit` documents.
    pub fn new(limit: usize) -> Self {
        Self {
            heap: BinaryHeap::with_capacity(limit + 1),
            limit,
        }
    }

    /// Offer a scored document. It will be kept if it's in the top-K.
    ///
    /// Uses replace-top via `peek_mut()` for competitive documents — a single
    /// sift-down instead of pop (sift-up) + push (sift-down).
    /// See [[optimization-scoring-throughput#Phase 2]].
    pub fn collect(&mut self, doc_id: DocId, segment_id: SegmentId, score: f32) {
        if self.heap.len() < self.limit {
            self.heap.push(MinScoreDoc(ScoreDoc {
                doc_id,
                segment_id,
                score,
            }));
            return;
        }
        // Fast path: non-competitive (most common case).
        // peek() returns None when limit=0 (agg-only queries).
        if let Some(min) = self.heap.peek() {
            if score <= min.0.score {
                return;
            }
        } else {
            return;
        }
        // Replace top and sift down (1 sift instead of 2)
        let mut top = self.heap.peek_mut().unwrap();
        *top = MinScoreDoc(ScoreDoc {
            doc_id,
            segment_id,
            score,
        });
        // PeekMut::drop calls sift_down automatically
    }

    /// Current minimum score in the top-K heap.
    /// Returns 0.0 if the heap is not yet full.
    pub fn min_score(&self) -> f32 {
        if self.heap.len() < self.limit {
            0.0
        } else {
            self.heap.peek().map(|m| m.0.score).unwrap_or(0.0)
        }
    }

    /// Merge another collector's results into this one.
    pub fn merge(&mut self, other: TopDocsCollector) {
        for item in other.heap {
            self.collect(item.0.doc_id, item.0.segment_id, item.0.score);
        }
    }

    /// Drain the collector and return results sorted by score descending.
    pub fn into_sorted_results(self) -> Vec<ScoreDoc> {
        let mut results: Vec<ScoreDoc> = self.heap.into_iter().map(|m| m.0).collect();
        results.sort_by(|a, b| {
            b.score
                .partial_cmp(&a.score)
                .unwrap_or(Ordering::Equal)
                .then_with(|| a.doc_id.cmp(&b.doc_id))
        });
        results
    }
}

// --- TopFieldCollector: sort by field values ---
// See [[feature-sort-by-field]].

use crate::search::{SortField, SortValue, compare_sort_values_cascade};

/// A collected document with its sort values.
#[derive(Clone, Debug)]
pub(crate) struct FieldDoc {
    pub(crate) doc_id: DocId,
    pub(crate) segment_id: SegmentId,
    pub(crate) score: f32,
    pub(crate) sort_values: Vec<SortValue>,
}

/// Wrapper for min-heap ordering by sort fields.
/// The "worst" doc (that should be evicted first) sits at the top.
struct MinFieldDoc {
    doc: FieldDoc,
    /// Shared reference to sort spec for comparison. We store a raw pointer
    /// because BinaryHeap requires Ord which can't borrow the collector.
    /// SAFETY: The pointer is valid for the lifetime of the TopFieldCollector.
    sort_fields: *const Vec<SortField>,
}

// SAFETY: MinFieldDoc is only used within TopFieldCollector which owns
// the sort_fields Vec and outlives all MinFieldDoc instances.
unsafe impl Send for MinFieldDoc {}

impl PartialEq for MinFieldDoc {
    fn eq(&self, other: &Self) -> bool {
        self.cmp(other) == Ordering::Equal
    }
}
impl Eq for MinFieldDoc {}

impl PartialOrd for MinFieldDoc {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for MinFieldDoc {
    fn cmp(&self, other: &Self) -> Ordering {
        // Reverse: we want the "worst" doc at the top of the max-heap
        // so it gets evicted first. "Worst" = sorts last in the desired order.
        let sort_fields = unsafe { &*self.sort_fields };
        let natural = compare_sort_values_cascade(
            &self.doc.sort_values,
            &other.doc.sort_values,
            sort_fields,
        );
        // Reverse so the max-heap acts as a min-heap for our sort order
        natural.reverse()
            .then_with(|| other.doc.doc_id.cmp(&self.doc.doc_id))
    }
}

/// Collects the top-K documents ordered by field sort values.
pub struct TopFieldCollector {
    heap: BinaryHeap<MinFieldDoc>,
    limit: usize,
    sort_fields: Vec<SortField>,
    /// Cursor for search_after pagination. Docs sorting <= this are skipped.
    search_after: Option<Vec<SortValue>>,
}

impl TopFieldCollector {
    pub fn new(limit: usize, sort_fields: Vec<SortField>) -> Self {
        Self {
            heap: BinaryHeap::with_capacity(limit + 1),
            limit,
            sort_fields,
            search_after: None,
        }
    }

    /// Set the search_after cursor. Docs sorting at or before this
    /// position are skipped. See [[feature-search-after]].
    pub fn set_search_after(&mut self, cursor: Vec<SortValue>) {
        self.search_after = Some(cursor);
    }

    /// Quick check: would a doc with this primary sort value be competitive?
    /// Returns true if the heap isn't full or the value beats the worst entry
    /// on the primary sort field. Used to avoid Vec allocation for non-competitive docs.
    #[inline]
    pub fn is_competitive_primary(&self, primary: &SortValue) -> bool {
        if self.heap.len() < self.limit {
            return true;
        }
        if let Some(top) = self.heap.peek() {
            if let Some(worst_primary) = top.doc.sort_values.first() {
                let cmp = primary.compare(worst_primary, &self.sort_fields[0]);
                // Competitive if this doc sorts BEFORE the worst
                return cmp == Ordering::Less;
            }
        }
        false
    }

    /// Quick check for keyword primary sort: compare &str without allocating.
    #[inline]
    pub fn is_competitive_keyword(&self, value: Option<&str>) -> bool {
        if self.heap.len() < self.limit {
            return true;
        }
        if let Some(top) = self.heap.peek() {
            if let Some(worst) = top.doc.sort_values.first() {
                let sv = match value {
                    Some(s) => SortValue::Str(String::new()), // placeholder, compare below
                    None => SortValue::Null,
                };
                // Direct comparison without allocating
                let natural = match (value, worst) {
                    (None, SortValue::Null) => std::cmp::Ordering::Equal,
                    (None, _) => match self.sort_fields[0].missing {
                        crate::search::MissingValue::Last => match self.sort_fields[0].order {
                            crate::search::SortOrder::Asc => std::cmp::Ordering::Greater,
                            crate::search::SortOrder::Desc => std::cmp::Ordering::Less,
                        },
                        crate::search::MissingValue::First => match self.sort_fields[0].order {
                            crate::search::SortOrder::Asc => std::cmp::Ordering::Less,
                            crate::search::SortOrder::Desc => std::cmp::Ordering::Greater,
                        },
                    },
                    (Some(_), SortValue::Null) => match self.sort_fields[0].missing {
                        crate::search::MissingValue::Last => match self.sort_fields[0].order {
                            crate::search::SortOrder::Asc => std::cmp::Ordering::Less,
                            crate::search::SortOrder::Desc => std::cmp::Ordering::Greater,
                        },
                        crate::search::MissingValue::First => match self.sort_fields[0].order {
                            crate::search::SortOrder::Asc => std::cmp::Ordering::Greater,
                            crate::search::SortOrder::Desc => std::cmp::Ordering::Less,
                        },
                    },
                    (Some(s), SortValue::Str(w)) => s.cmp(w.as_str()),
                    _ => std::cmp::Ordering::Equal,
                };
                let cmp = match self.sort_fields[0].order {
                    crate::search::SortOrder::Asc => natural,
                    crate::search::SortOrder::Desc => natural.reverse(),
                };
                return cmp == Ordering::Less;
            }
        }
        false
    }

    /// Offer a document with its sort values.
    pub fn collect(&mut self, doc_id: DocId, segment_id: SegmentId, score: f32, sort_values: Vec<SortValue>) {
        // search_after filter: skip docs at or before the cursor
        if let Some(ref after) = self.search_after {
            let cmp = compare_sort_values_cascade(&sort_values, after, &self.sort_fields);
            if cmp != Ordering::Greater {
                return;
            }
        }

        let doc = FieldDoc { doc_id, segment_id, score, sort_values };

        if self.heap.len() < self.limit {
            self.heap.push(MinFieldDoc {
                doc,
                sort_fields: &self.sort_fields as *const Vec<SortField>,
            });
            return;
        }
        // Fast path: non-competitive
        if let Some(top) = self.heap.peek() {
            let cmp = compare_sort_values_cascade(
                &doc.sort_values,
                &top.doc.sort_values,
                &self.sort_fields,
            );
            // If doc sorts after or equal to the worst in heap, skip
            if cmp != Ordering::Less {
                return;
            }
        } else {
            return;
        }
        // Replace top (worst) with this better doc
        let mut top = self.heap.peek_mut().unwrap();
        *top = MinFieldDoc {
            doc,
            sort_fields: &self.sort_fields as *const Vec<SortField>,
        };
    }

    /// Merge another collector's results into this one.
    pub fn merge(&mut self, other: TopFieldCollector) {
        for item in other.heap {
            self.collect(item.doc.doc_id, item.doc.segment_id, item.doc.score, item.doc.sort_values);
        }
    }

    /// Drain and return results in sort order (best first).
    pub fn into_sorted_results(self) -> Vec<FieldDoc> {
        let sort_fields = self.sort_fields;
        let mut results: Vec<FieldDoc> = self.heap.into_iter().map(|m| m.doc).collect();
        results.sort_by(|a, b| {
            compare_sort_values_cascade(&a.sort_values, &b.sort_values, &sort_fields)
                .then_with(|| a.doc_id.cmp(&b.doc_id))
        });
        results
    }
}

// --- CollapsingCollector: deduplicate by field value ---
// See [[feature-search-collapse]].

use std::collections::HashMap;

/// A collapse group key.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub(crate) enum CollapseKey {
    Null,
    Keyword(String),
    Numeric(i64),
}

impl CollapseKey {
    pub fn to_json(&self) -> serde_json::Value {
        match self {
            CollapseKey::Null => serde_json::Value::Null,
            CollapseKey::Keyword(s) => serde_json::json!(s),
            CollapseKey::Numeric(n) => serde_json::json!(n),
        }
    }
}

/// A collected doc with its collapse key.
#[derive(Clone, Debug)]
pub(crate) struct CollapsedDoc {
    pub(crate) doc_id: DocId,
    pub(crate) segment_id: SegmentId,
    pub(crate) score: f32,
    pub(crate) collapse_key: CollapseKey,
}

/// Collects the highest-scoring document per unique collapse value.
pub struct CollapsingCollector {
    groups: HashMap<CollapseKey, CollapsedDoc>,
    limit: usize,
}

impl CollapsingCollector {
    pub fn new(limit: usize) -> Self {
        Self {
            groups: HashMap::new(),
            limit,
        }
    }

    /// Offer a scored document with its collapse key.
    pub fn collect(&mut self, doc_id: DocId, segment_id: SegmentId, score: f32, key: CollapseKey) {
        let entry = self.groups.entry(key.clone()).or_insert(CollapsedDoc {
            doc_id, segment_id, score, collapse_key: key,
        });
        if score > entry.score {
            entry.doc_id = doc_id;
            entry.segment_id = segment_id;
            entry.score = score;
        }
    }

    /// Merge another collector's groups into this one.
    pub fn merge(&mut self, other: CollapsingCollector) {
        for (key, doc) in other.groups {
            self.collect(doc.doc_id, doc.segment_id, doc.score, key);
        }
    }

    /// Total number of unique groups (pre-collapse total_hits).
    pub fn group_count(&self) -> usize {
        self.groups.len()
    }

    /// Drain and return the top-K groups sorted by score descending.
    pub fn into_sorted_results(self) -> Vec<CollapsedDoc> {
        let mut results: Vec<CollapsedDoc> = self.groups.into_values().collect();
        results.sort_by(|a, b| {
            b.score
                .partial_cmp(&a.score)
                .unwrap_or(Ordering::Equal)
                .then_with(|| a.doc_id.cmp(&b.doc_id))
        });
        results.truncate(self.limit);
        results
    }
}

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

    #[test]
    fn collect_under_limit() {
        let mut collector = TopDocsCollector::new(10);
        collector.collect(DocId::new(0), SegmentId::new(1), 1.0);
        collector.collect(DocId::new(1), SegmentId::new(1), 2.0);
        collector.collect(DocId::new(2), SegmentId::new(1), 0.5);

        let results = collector.into_sorted_results();
        assert_eq!(results.len(), 3);
        assert_eq!(results[0].doc_id, DocId::new(1)); // highest score
        assert_eq!(results[1].doc_id, DocId::new(0));
        assert_eq!(results[2].doc_id, DocId::new(2)); // lowest score
    }

    #[test]
    fn top_k_limiting() {
        let mut collector = TopDocsCollector::new(2);
        collector.collect(DocId::new(0), SegmentId::new(1), 1.0);
        collector.collect(DocId::new(1), SegmentId::new(1), 3.0);
        collector.collect(DocId::new(2), SegmentId::new(1), 2.0);
        collector.collect(DocId::new(3), SegmentId::new(1), 0.5);

        let results = collector.into_sorted_results();
        assert_eq!(results.len(), 2);
        assert_eq!(results[0].doc_id, DocId::new(1)); // score 3.0
        assert_eq!(results[1].doc_id, DocId::new(2)); // score 2.0
    }

    #[test]
    fn score_descending_order() {
        let mut collector = TopDocsCollector::new(5);
        for i in 0..5 {
            collector.collect(DocId::new(i), SegmentId::new(1), i as f32);
        }
        let results = collector.into_sorted_results();
        for i in 0..results.len() - 1 {
            assert!(results[i].score >= results[i + 1].score);
        }
    }

    #[test]
    fn empty_collector() {
        let collector = TopDocsCollector::new(10);
        let results = collector.into_sorted_results();
        assert!(results.is_empty());
    }

    #[test]
    fn multi_segment() {
        let mut collector = TopDocsCollector::new(3);
        collector.collect(DocId::new(0), SegmentId::new(1), 1.0);
        collector.collect(DocId::new(0), SegmentId::new(2), 2.0);
        collector.collect(DocId::new(1), SegmentId::new(1), 3.0);
        collector.collect(DocId::new(1), SegmentId::new(2), 0.5);

        let results = collector.into_sorted_results();
        assert_eq!(results.len(), 3);
        assert_eq!(results[0].score, 3.0);
        assert_eq!(results[0].segment_id, SegmentId::new(1));
    }
}
