use std::sync::Arc;

use anyhow::{Context, Result};
use tantivy::{
    DocAddress, Searcher, SegmentOrdinal, SegmentReader, TERMINATED, TantivyDocument,
    index::SegmentId,
    query::{EnableScoring, Query, Scorer},
};

pub trait SegmentDocProvider {
    fn fetch_doc(&self, addr: DocAddress) -> Result<Arc<TantivyDocument>>;
}

impl SegmentDocProvider for Searcher {
    fn fetch_doc(&self, addr: DocAddress) -> Result<Arc<TantivyDocument>> {
        Ok(self.doc(addr).map(Arc::new)?)
    }
}

// A lightweight wrapper that scans documents in a specific segment
// of a Tantivy index based on a given query.
//
// The `SegQueryIter ` internally holds a Tantivy `Scorer` which
// iterates through matching document IDs in the segment.
pub struct SegQueryIter {
    seg_id: SegmentId,
    // The ordinal (index) of the current segment within the index searcher.
    seg_ord: SegmentOrdinal,
    // Document provider used to retrieve documents.
    doc_provider: Arc<dyn SegmentDocProvider>,
    segment_reader: Arc<SegmentReader>,
    // A scorer that advances through matching document IDs.
    // Scorer guarantee to return document IDs in ascending order.
    scorer: Box<dyn Scorer>,
}

impl SegQueryIter {
    pub fn new(
        seg_ord: SegmentOrdinal,
        doc_provider: Arc<dyn SegmentDocProvider>,
        seg_reader: Arc<SegmentReader>,
        query: &dyn Query,
    ) -> Result<Self> {
        let weight = query.weight(EnableScoring::Disabled {
            schema: seg_reader.schema(),
            searcher_opt: None,
        })?;
        let scorer = weight.scorer(seg_reader.as_ref(), 1.0)?;
        Ok(Self {
            seg_id: seg_reader.segment_id(),
            seg_ord,
            doc_provider,
            segment_reader: seg_reader,
            scorer,
        })
    }

    #[cfg(test)]
    pub(crate) fn segment_reader_for_test(&self) -> &Arc<SegmentReader> {
        &self.segment_reader
    }

    pub fn next_doc_id(&mut self) -> Option<u32> {
        let mut doc_id = self.scorer.doc();

        if let Some(alive_bitset) = self.segment_reader.alive_bitset() {
            while doc_id != TERMINATED && alive_bitset.is_deleted(doc_id) {
                doc_id = self.scorer.advance();
            }
        }

        if doc_id == TERMINATED {
            return None;
        }
        // Make scorer point to next doc id.
        let _ = self.scorer.advance();

        Some(doc_id)
    }

    pub fn next_doc_addr(&mut self) -> Option<DocAddress> {
        self.next_doc_id().map(|doc_id| DocAddress {
            segment_ord: self.seg_ord,
            doc_id,
        })
    }

    pub fn next_doc(&mut self) -> Result<Option<Arc<TantivyDocument>>> {
        match self.next_doc_addr() {
            None => Ok(None),
            Some(doc_addr) => {
                let doc = self.doc_provider.fetch_doc(doc_addr).with_context(|| {
                    format!(
                        "Failed to get doc {} from segment {}",
                        doc_addr.doc_id, self.seg_id
                    )
                })?;
                Ok(Some(doc))
            }
        }
    }
}