//! dis_max query: returns the score from the best-matching sub-query,
//! plus tie_breaker * sum of other matching scores.
//!
//! See [[query-dsl]].

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

use crate::query::{Query, ScorerSupplier, Weight};
use crate::search::searcher::Searcher;
use crate::segment::reader::SegmentReader;

/// Disjunction max query.
pub struct DisMaxQuery {
    pub queries: Vec<Box<dyn Query>>,
    pub tie_breaker: f32,
}

impl Query for DisMaxQuery {
    fn create_weight(&self, searcher: &Searcher, score_mode: ScoreMode) -> Result<Box<dyn Weight>> {
        let weights: Vec<Box<dyn Weight>> = self.queries
            .iter()
            .map(|q| q.create_weight(searcher, score_mode))
            .collect::<Result<_>>()?;
        Ok(Box::new(DisMaxWeight {
            weights,
            tie_breaker: self.tie_breaker,
        }))
    }
}

struct DisMaxWeight {
    weights: Vec<Box<dyn Weight>>,
    tie_breaker: f32,
}

impl Weight for DisMaxWeight {
    fn scorer_supplier(
        &self,
        reader: &SegmentReader,
    ) -> Result<Option<Box<dyn ScorerSupplier>>> {
        let mut suppliers: Vec<Box<dyn ScorerSupplier>> = Vec::new();
        for w in &self.weights {
            if let Some(s) = w.scorer_supplier(reader)? {
                suppliers.push(s);
            }
        }
        if suppliers.is_empty() {
            return Ok(None);
        }
        Ok(Some(Box::new(DisMaxScorerSupplier {
            suppliers,
            tie_breaker: self.tie_breaker,
        })))
    }
}

struct DisMaxScorerSupplier {
    suppliers: Vec<Box<dyn ScorerSupplier>>,
    tie_breaker: f32,
}

impl ScorerSupplier for DisMaxScorerSupplier {
    fn cost(&self) -> u64 {
        self.suppliers.iter().map(|s| s.cost()).sum()
    }

    fn scorer(self: Box<Self>) -> Result<Box<dyn Scorer>> {
        let scorers: Vec<Box<dyn Scorer>> = self.suppliers
            .into_iter()
            .map(|s| s.scorer())
            .collect::<Result<_>>()?;
        Ok(Box::new(DisMaxScorer::new(scorers, self.tie_breaker)))
    }
}

/// Scorer that iterates the union of sub-scorers, computing
/// max(score) + tie_breaker * sum(other_scores) per document.
struct DisMaxScorer {
    scorers: Vec<Box<dyn Scorer>>,
    tie_breaker: f32,
    current_doc: DocId,
}

impl DisMaxScorer {
    fn new(mut scorers: Vec<Box<dyn Scorer>>, tie_breaker: f32) -> Self {
        // Advance all scorers to their first doc
        let min_doc = scorers.iter()
            .map(|s| s.doc_id())
            .filter(|&d| d != NO_MORE_DOCS)
            .min()
            .unwrap_or(NO_MORE_DOCS);
        Self {
            scorers,
            tie_breaker,
            current_doc: min_doc,
        }
    }

    /// Find the minimum doc_id across all scorers.
    fn find_min_doc(&self) -> DocId {
        self.scorers.iter()
            .map(|s| s.doc_id())
            .filter(|&d| d != NO_MORE_DOCS)
            .min()
            .unwrap_or(NO_MORE_DOCS)
    }

    /// Compute the dis_max score for the current doc.
    fn compute_score(&mut self) -> f32 {
        let mut max_score: f32 = 0.0;
        let mut sum_other: f32 = 0.0;

        for scorer in &mut self.scorers {
            if scorer.doc_id() == self.current_doc {
                let s = scorer.score();
                if s > max_score {
                    sum_other += max_score;
                    max_score = s;
                } else {
                    sum_other += s;
                }
            }
        }

        max_score + self.tie_breaker * sum_other
    }
}

impl Scorer for DisMaxScorer {
    fn doc_id(&self) -> DocId {
        self.current_doc
    }

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

        // Advance all scorers that are at current_doc
        let cur = self.current_doc;
        for scorer in &mut self.scorers {
            if scorer.doc_id() == cur {
                scorer.next();
            }
        }

        self.current_doc = self.find_min_doc();
        self.current_doc
    }

    fn advance(&mut self, target: DocId) -> DocId {
        for scorer in &mut self.scorers {
            if scorer.doc_id() < target && scorer.doc_id() != NO_MORE_DOCS {
                scorer.advance(target);
            }
        }
        self.current_doc = self.find_min_doc();
        self.current_doc
    }

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

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

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

    #[test]
    fn dis_max_score_computation() {
        // Verify score = max + tie_breaker * sum(others)
        // With scores [3.0, 1.0, 2.0] and tie_breaker 0.5:
        // max = 3.0, others = 1.0 + 2.0 = 3.0
        // score = 3.0 + 0.5 * 3.0 = 4.5

        // This is a logic test — we can't easily create mock scorers
        // without the full infrastructure, so we test via integration.
    }
}
