//! Disjunction scorer: OR semantics for `bool.should` clauses.
//!
//! Uses a priority queue of sub-scorers ordered by current doc ID.
//! On each `next()`, advances to the lowest doc ID across all scorers,
//! sums their scores.
//!
//! See [[architecture-query-execution#Disjunction Execution]] and [[architecture-query-execution#Step 7]].

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

use luci_core::{DocId, Scorer, TwoPhaseIterator, NO_MORE_DOCS};

/// Wrapper for min-heap ordering by doc_id.
struct ScorerEntry {
    scorer: Box<dyn Scorer>,
}

impl PartialEq for ScorerEntry {
    fn eq(&self, other: &Self) -> bool {
        self.scorer.doc_id() == other.scorer.doc_id()
    }
}
impl Eq for ScorerEntry {}

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

impl Ord for ScorerEntry {
    fn cmp(&self, other: &Self) -> Ordering {
        // Reverse for min-heap (BinaryHeap is max-heap)
        other.scorer.doc_id().cmp(&self.scorer.doc_id())
    }
}

/// A scorer that returns docs matching ANY sub-scorer (OR semantics).
///
/// Scores are summed across all sub-scorers matching the current doc.
pub struct DisjunctionScorer {
    heap: BinaryHeap<ScorerEntry>,
    current: DocId,
    current_score: f32,
}

impl DisjunctionScorer {
    /// Create a disjunction scorer. `scorers` must be non-empty.
    pub fn new(scorers: Vec<Box<dyn Scorer>>) -> Self {
        assert!(!scorers.is_empty(), "disjunction needs at least one scorer");

        let mut heap = BinaryHeap::with_capacity(scorers.len());
        for scorer in scorers {
            if scorer.doc_id() != NO_MORE_DOCS {
                heap.push(ScorerEntry { scorer });
            }
        }

        let mut disj = Self {
            heap,
            current: NO_MORE_DOCS,
            current_score: 0.0,
        };
        disj.refill_current();
        disj
    }

    /// Collect the current minimum doc_id and sum scores from all scorers
    /// positioned at that doc.
    fn refill_current(&mut self) {
        if self.heap.is_empty() {
            self.current = NO_MORE_DOCS;
            self.current_score = 0.0;
            return;
        }

        self.current = self.heap.peek().unwrap().scorer.doc_id();
        self.current_score = 0.0;

        // Sum scores from all scorers at current doc
        // We need to peek all entries at the same doc_id
        // Collect them, score, then put back
        let mut at_current = Vec::new();
        while let Some(top) = self.heap.peek() {
            if top.scorer.doc_id() == self.current {
                let mut entry = self.heap.pop().unwrap();
                self.current_score += entry.scorer.score();
                at_current.push(entry);
            } else {
                break;
            }
        }
        // Put them all back
        for entry in at_current {
            self.heap.push(entry);
        }
    }
}

impl Scorer for DisjunctionScorer {
    fn doc_id(&self) -> DocId {
        self.current
    }

    fn next(&mut self) -> DocId {
        if self.current == NO_MORE_DOCS {
            return NO_MORE_DOCS;
        }

        // Advance all scorers that are at the current doc
        let target = self.current;
        let mut to_reinsert = Vec::new();
        while let Some(top) = self.heap.peek() {
            if top.scorer.doc_id() == target {
                let mut entry = self.heap.pop().unwrap();
                entry.scorer.next();
                if entry.scorer.doc_id() != NO_MORE_DOCS {
                    to_reinsert.push(entry);
                }
            } else {
                break;
            }
        }
        for entry in to_reinsert {
            self.heap.push(entry);
        }

        self.refill_current();
        self.current
    }

    fn advance(&mut self, target: DocId) -> DocId {
        // Advance all scorers past current position to at least target
        let mut entries = Vec::new();
        while let Some(entry) = self.heap.pop() {
            entries.push(entry);
        }
        for mut entry in entries {
            if entry.scorer.doc_id() < target {
                entry.scorer.advance(target);
            }
            if entry.scorer.doc_id() != NO_MORE_DOCS {
                self.heap.push(entry);
            }
        }
        self.refill_current();
        self.current
    }

    fn score(&mut self) -> f32 {
        self.current_score
    }

    fn two_phase(&mut self) -> Option<&mut dyn TwoPhaseIterator> {
        None
    }
}

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

    struct VecScorer {
        docs: Vec<(DocId, f32)>,
        pos: usize,
    }

    impl VecScorer {
        fn new(docs: Vec<(u32, f32)>) -> Box<dyn Scorer> {
            Box::new(Self {
                docs: docs.into_iter().map(|(id, s)| (DocId::new(id), s)).collect(),
                pos: 0,
            })
        }
        fn current(&self) -> (DocId, f32) {
            if self.pos < self.docs.len() { self.docs[self.pos] } else { (NO_MORE_DOCS, 0.0) }
        }
    }

    impl Scorer for VecScorer {
        fn doc_id(&self) -> DocId { self.current().0 }
        fn next(&mut self) -> DocId {
            if self.pos < self.docs.len() { self.pos += 1; }
            self.current().0
        }
        fn advance(&mut self, target: DocId) -> DocId {
            while self.pos < self.docs.len() && self.docs[self.pos].0 < target {
                self.pos += 1;
            }
            self.current().0
        }
        fn score(&mut self) -> f32 { self.current().1 }
        fn two_phase(&mut self) -> Option<&mut dyn TwoPhaseIterator> { None }
    }

    #[test]
    fn two_scorers_union() {
        let s1 = VecScorer::new(vec![(0, 1.0), (2, 1.0), (4, 1.0)]);
        let s2 = VecScorer::new(vec![(1, 2.0), (2, 2.0), (3, 2.0)]);

        let mut disj = DisjunctionScorer::new(vec![s1, s2]);
        assert_eq!(disj.doc_id(), DocId::new(0));
        assert_eq!(disj.score(), 1.0); // only s1

        assert_eq!(disj.next(), DocId::new(1));
        assert_eq!(disj.score(), 2.0); // only s2

        assert_eq!(disj.next(), DocId::new(2));
        assert_eq!(disj.score(), 3.0); // both: 1.0 + 2.0

        assert_eq!(disj.next(), DocId::new(3));
        assert_eq!(disj.score(), 2.0); // only s2

        assert_eq!(disj.next(), DocId::new(4));
        assert_eq!(disj.score(), 1.0); // only s1

        assert_eq!(disj.next(), NO_MORE_DOCS);
    }

    #[test]
    fn three_scorers() {
        let s1 = VecScorer::new(vec![(5, 1.0)]);
        let s2 = VecScorer::new(vec![(5, 2.0)]);
        let s3 = VecScorer::new(vec![(5, 3.0)]);

        let mut disj = DisjunctionScorer::new(vec![s1, s2, s3]);
        assert_eq!(disj.doc_id(), DocId::new(5));
        assert_eq!(disj.score(), 6.0); // 1 + 2 + 3
        assert_eq!(disj.next(), NO_MORE_DOCS);
    }

    #[test]
    fn advance_disjunction() {
        let s1 = VecScorer::new(vec![(0, 1.0), (5, 1.0), (10, 1.0)]);
        let s2 = VecScorer::new(vec![(3, 2.0), (7, 2.0), (10, 2.0)]);

        let mut disj = DisjunctionScorer::new(vec![s1, s2]);
        assert_eq!(disj.advance(DocId::new(5)), DocId::new(5));
        assert_eq!(disj.score(), 1.0);
        assert_eq!(disj.advance(DocId::new(10)), DocId::new(10));
        assert_eq!(disj.score(), 3.0); // both match at 10
    }

    #[test]
    fn single_scorer_disjunction() {
        let s1 = VecScorer::new(vec![(0, 1.0), (1, 2.0)]);
        let mut disj = DisjunctionScorer::new(vec![s1]);
        assert_eq!(disj.doc_id(), DocId::new(0));
        assert_eq!(disj.next(), DocId::new(1));
        assert_eq!(disj.next(), NO_MORE_DOCS);
    }

    #[test]
    fn non_overlapping_scorers() {
        let s1 = VecScorer::new(vec![(0, 1.0), (2, 1.0)]);
        let s2 = VecScorer::new(vec![(1, 2.0), (3, 2.0)]);

        let mut disj = DisjunctionScorer::new(vec![s1, s2]);
        let mut docs = Vec::new();
        while disj.doc_id() != NO_MORE_DOCS {
            docs.push(disj.doc_id().as_u32());
            disj.next();
        }
        assert_eq!(docs, vec![0, 1, 2, 3]);
    }
}
