use std::{cmp::Ordering, sync::Arc};

use anyhow::Result;
use binary_heap_plus::{BinaryHeap, FnComparator};
use tantivy::TantivyDocument;

use crate::reader::seg_query_iter::SegQueryIter;

pub struct HeapEntry {
    pub doc: Arc<TantivyDocument>,
    pub scanner_idx: usize,
}

pub struct MergeIter<F>
where
    F: Fn(&HeapEntry, &HeapEntry) -> Ordering,
{
    scanners: Vec<SegQueryIter>,
    heap: BinaryHeap<HeapEntry, FnComparator<F>>,
}

impl<F> MergeIter<F>
where
    F: Fn(&HeapEntry, &HeapEntry) -> Ordering,
{
    pub fn new(mut scanners: Vec<SegQueryIter>, cmp: F) -> Result<Self> {
        let mut heap = BinaryHeap::new_by(cmp);

        for (i, scanner) in scanners.iter_mut().enumerate() {
            if let Some(doc) = scanner.next_doc()? {
                heap.push(HeapEntry {
                    doc,
                    scanner_idx: i,
                });
            }
        }

        Ok(Self { scanners, heap })
    }

    pub fn next_doc(&mut self) -> Result<Option<Arc<TantivyDocument>>> {
        let entry = match self.heap.pop() {
            Some(e) => e,
            None => return Ok(None),
        };

        let scanner_idx = entry.scanner_idx;
        if let Some(next_doc) = self.scanners[scanner_idx].next_doc()? {
            self.heap.push(HeapEntry {
                doc: next_doc,
                scanner_idx,
            });
        }

        Ok(Some(entry.doc))
    }
}