//! PrefixQuery: match documents containing terms with a given prefix.
//!
//! Enumerates matching terms via `TermDict::prefix_iter`, then builds
//! a disjunction of posting lists.
//!
//! See [[query-dsl#Term-Level Queries]] and [[architecture-query-execution#Step 6]].

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

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

pub struct PrefixQuery {
    pub field: String,
    pub value: String,
}

impl Query for PrefixQuery {
    fn create_weight(
        &self,
        _searcher: &Searcher,
        _score_mode: ScoreMode,
    ) -> Result<Box<dyn Weight>> {
        Ok(Box::new(PrefixWeight {
            field: self.field.clone(),
            prefix: self.value.clone(),
        }))
    }
}

struct PrefixWeight {
    field: String,
    prefix: String,
}

impl Weight for PrefixWeight {
    fn scorer_supplier(
        &self,
        reader: &SegmentReader,
    ) -> Result<Option<Box<dyn ScorerSupplier>>> {
        let field_id = match reader.header().fields.iter()
            .find(|f| f.field_name == self.field)
            .map(|f| f.field_id)
        {
            Some(id) => id,
            None => return Ok(None),
        };

        // Find all terms with this prefix
        let terms = reader.terms_with_prefix(field_id, &self.prefix);
        if terms.is_empty() {
            return Ok(None);
        }

        let cost: u64 = terms.iter().map(|(_, df)| *df as u64).sum();

        Ok(Some(Box::new(PrefixScorerSupplier {
            field_id,
            terms,
            cost,
            segment_data: reader as *const SegmentReader,
        })))
    }
}

struct PrefixScorerSupplier {
    field_id: luci_core::FieldId,
    terms: Vec<(String, u32)>,
    cost: u64,
    segment_data: *const SegmentReader,
}

unsafe impl Send for PrefixScorerSupplier {}

impl ScorerSupplier for PrefixScorerSupplier {
    fn cost(&self) -> u64 {
        self.cost
    }

    fn scorer(self: Box<Self>) -> Result<Box<dyn Scorer>> {
        let reader = unsafe { &*self.segment_data };
        let mut scorers: Vec<Box<dyn Scorer>> = Vec::new();

        for (term, _) in &self.terms {
            if let Some(postings) = reader.postings(self.field_id, term) {
                scorers.push(Box::new(SimplePostingScorer::new(postings)));
            }
        }

        if scorers.is_empty() {
            return Ok(Box::new(EmptyScorer));
        }
        if scorers.len() == 1 {
            return Ok(scorers.pop().unwrap());
        }
        Ok(Box::new(DisjunctionScorer::new(scorers)))
    }
}

struct SimplePostingScorer<'a> {
    postings: crate::inverted::postings::PostingListReader<'a>,
    current: DocId,
}

impl<'a> SimplePostingScorer<'a> {
    fn new(mut postings: crate::inverted::postings::PostingListReader<'a>) -> Self {
        let current = match postings.next() {
            Some((id, _)) => id,
            None => NO_MORE_DOCS,
        };
        Self { postings, current }
    }
}

impl Scorer for SimplePostingScorer<'_> {
    fn doc_id(&self) -> DocId { self.current }
    fn next(&mut self) -> DocId {
        self.current = match self.postings.next() {
            Some((id, _)) => id,
            None => NO_MORE_DOCS,
        };
        self.current
    }
    fn advance(&mut self, target: DocId) -> DocId {
        while self.current < target && self.current != NO_MORE_DOCS {
            self.next();
        }
        self.current
    }
    fn score(&mut self) -> f32 { 1.0 }
    fn two_phase(&mut self) -> Option<&mut dyn TwoPhaseIterator> { None }
}

struct EmptyScorer;
impl Scorer for EmptyScorer {
    fn doc_id(&self) -> DocId { NO_MORE_DOCS }
    fn next(&mut self) -> DocId { NO_MORE_DOCS }
    fn advance(&mut self, _: DocId) -> DocId { NO_MORE_DOCS }
    fn score(&mut self) -> f32 { 0.0 }
    fn two_phase(&mut self) -> Option<&mut dyn TwoPhaseIterator> { None }
}
