pub trait ScoreBatch {
    fn next(&mut self) -> Option<(t_docId, f64)>;
    fn skip_to(&mut self, target: t_docId) -> Option<(t_docId, f64)>;
}

pub trait ScoreSource<'index> {
    type Batch: ScoreBatch;
    fn next_batch(&mut self) -> Result<Option<Self::Batch>, RQEIteratorError>;
    // Optional in practice: sources that never switch to adhoc can return None.
    fn lookup_score(&mut self, doc_id: t_docId) -> Option<f64>;
    fn num_estimated(&self) -> usize;
    fn rewind(&mut self);  // Called by TopK when CollectionStrategy::Rewind is returned
    fn build_result(&self, doc_id: t_docId, score: f64) -> RSIndexResult<'index>;
    fn collection_strategy(&mut self, heap_count: usize, k: usize) -> CollectionStrategy;
}

pub enum CollectionStrategy {
    Continue,        // Keep iterating current batch sequence
    SwitchToAdhoc,   // Rewind child, switch to adhoc brute-force mode (Vector in v1)
    SwitchToBatches, // Source rewinds itself, TopK rewinds child, restart batches
    Stop,            // Collection complete
}

/// A cursor over a score-ordered batch.
///
/// The cursor yields entries in score order.
/// For filtered intersection, it must also support doc-id navigation via `skip_to`.
pub trait ScoreBatch {
    fn next(&mut self) -> Option<(t_docId, f64)>;
    /// Forward-only skip, equivalent to `RQEIterator::skip_to` semantics:
    /// - Returns the first entry with doc_id >= target from the current position.
    /// - Returns `None` if no such entry exists (cursor becomes exhausted).
    /// - If it returns `Some(entry)`, a subsequent `next()` returns the following entry.
    /// Implementations may use auxiliary state/indexes internally.
    fn skip_to(&mut self, target: t_docId) -> Option<(t_docId, f64)>;
}

/// A source of scores for top-k collection.
///
/// Implementations provide batches of results where:
/// - Each batch is ordered by score (best score first for that source/order)
/// - Batch cursors support `skip_to(doc_id)` for filtered intersection
pub trait ScoreSource<'index> {
    type Batch: ScoreBatch;

    /// Get the next batch of results.
    ///
    /// Returns `Ok(Some(batch))` if more results are available.
    /// Returns `Ok(None)` if exhausted.
    /// Returns `Err(TimedOut)` if timeout reached.
    ///
    /// Unfiltered contract:
    /// - Sources should emit a single final score-ordered batch (up to k docs), then `None`.
    fn next_batch(&mut self) -> Result<Option<Self::Batch>, RQEIteratorError>;

    /// Lookup the score for a single document (for adhoc-BF mode).
    ///
    /// Returns `None` if the document doesn't exist in the source.
    /// Takes `&mut self` because some implementations may need to acquire locks.
    fn lookup_score(&mut self, doc_id: t_docId) -> Option<f64>;

    /// Estimated total number of results.
    fn num_estimated(&self) -> usize;

    /// Rewind to the beginning (for strategy retry).
    fn rewind(&mut self);

    /// Build an RSIndexResult from a doc_id and score.
    ///
    /// Implementations can return different result types (MetricResult, AggregateResult, etc.)
    fn build_result(&self, doc_id: t_docId, score: f64) -> RSIndexResult<'index>;

    /// Decide whether to continue, switch strategy, or stop.
    ///
    /// Called after each batch is processed. Takes `&mut self` so the source
    /// can update internal parameters before returning `Rewind`.
    fn collection_strategy(&mut self, heap_count: usize, k: usize) -> CollectionStrategy;
}

/// Strategy decision for collection phase.
pub enum CollectionStrategy {
    /// Continue with current mode.
    Continue,
    /// Rewind child and switch to adhoc brute-force mode.
    SwitchToAdhoc,
    /// Rewind child and restart batches mode.
    ///
    /// The source should adjust its internal parameters (e.g., expand numeric range)
    /// AND call `self.rewind()` BEFORE returning this variant. TopK will then call
    /// `child.rewind()` before resuming batch collection.
    SwitchToBatches,
    /// Stop collection (enough results or exhausted).
    Stop,
}

/// Execution mode for top-k collection.
pub enum TopKMode {
    /// No child filter - yield directly from source batch cursor.
    Unfiltered,
    /// Get batches from source, intersect with child.
    Batches,
    /// Iterate child, lookup scores in source.
    AdhocBF,
}

/// A top-k iterator that returns the best k results by score.
///
/// Results are yielded sorted by score, NOT by document ID.
/// This iterator can only be used at the root of a query tree.
pub struct TopKIterator<'index, S: ScoreSource<'index>> {
    source: S,
    child: Option<Box<dyn RQEIterator<'index> + 'index>>,
    mode: TopKMode,
    heap: TopKHeap<ScoredResult>,
    direct_batch: Option<S::Batch>, // Used by unfiltered direct-yield path
    k: usize,
    compare: fn(f64, f64) -> Ordering,  // Score comparison

    // Execution state
    phase: Phase,

    // Output state (for RQEIterator impl)
    current: Option<RSIndexResult<'index>>,
    last_doc_id: t_docId,
    at_eof: bool,

    // Metrics (for profiling)
    metrics: TopKMetrics,
}

enum Phase {
    NotStarted,
    Collecting,
    Yielding,
    YieldingDirect,
}

pub struct TopKMetrics {
    pub num_batches: usize,
    pub strategy_switches: usize,
    pub total_comparisons: usize,
}

fn prepare_unfiltered_direct(&mut self) -> Result<(), RQEIteratorError> {
    // In unfiltered mode, source returns a single final score-ordered batch.
    self.direct_batch = self.source.next_batch()?;
    self.phase = Phase::YieldingDirect;
    Ok(())
}

fn read_unfiltered_direct(&mut self) -> Result<Option<RSIndexResult<'index>>, RQEIteratorError> {
    let Some(batch) = self.direct_batch.as_mut() else {
        return Ok(None);
    };

    if let Some((doc_id, score)) = batch.next() {
        return Ok(Some(self.source.build_result(doc_id, score)));
    }

    // Optional sanity check: sources should be exhausted after the final batch in this mode.
    if cfg!(debug_assertions) {
        debug_assert!(matches!(self.source.next_batch()?, None));
    }
    self.direct_batch = None;
    Ok(None)
}

fn collect_batches(&mut self) -> Result<(), RQEIteratorError> {
    let child = self.child.as_mut().unwrap();

    'outer: loop {
        while let Some(mut batch) = self.source.next_batch()? {
            child.rewind();
            self.intersect_batch_with_child(&mut batch, child)?;

            match self.source.collection_strategy(self.heap.len(), self.k) {
                CollectionStrategy::Continue => {}
                CollectionStrategy::SwitchToAdhoc => {
                    child.rewind();  // Need full child iteration for score lookups
                    self.metrics.strategy_switches += 1;
                    return self.collect_adhoc();
                }
                CollectionStrategy::SwitchToBatches => {
                    // Source has already adjusted parameters and rewound itself.
                    // Just rewind the child and restart batch collection.
                    child.rewind();
                    self.metrics.strategy_switches += 1;
                    continue 'outer;
                }
                CollectionStrategy::Stop => break 'outer,
            }
        }
        // Source exhausted without requesting rewind - we're done
        break;
    }
    Ok(())
}

fn intersect_batch_with_child<B: ScoreBatch>(
    &mut self,
    batch: &mut B,
    child: &mut dyn RQEIterator
) -> Result<(), RQEIteratorError> {
    let mut batch_entry = batch.next();
    let mut child_result = child.read()?;

    while let (Some((batch_id, score)), Some(child_res)) = (batch_entry, child_result.as_ref()) {
        match batch_id.cmp(&child_res.doc_id) {
            Ordering::Equal => {
                self.heap.maybe_insert(batch_id, score);
                batch_entry = batch.next();
                child_result = child.read()?;
            }
            Ordering::Greater => {
                // Child behind - skip forward
                child_result = child.skip_to(batch_id)?.map(|o| o.into_result());
            }
            Ordering::Less => {
                // Batch behind - advance batch
                batch_entry = batch.skip_to(child_res.doc_id).or_else(|| batch.next());
            }
        }
    }
    Ok(())
}

fn collect_adhoc(&mut self) -> Result<(), RQEIteratorError> {
    // Used by sources that support per-doc score lookup (Vector in v1).
    let child = self.child.as_mut().unwrap();

    while let Some(result) = child.read()? {
        if let Some(score) = self.source.lookup_score(result.doc_id) {
            self.heap.maybe_insert(result.doc_id, score);
        }

        // Check if we should stop (heap full)
        if matches!(self.source.collection_strategy(self.heap.len(), self.k),
                    CollectionStrategy::Stop) {
            break;
        }
    }
    Ok(())
}

pub enum CollectionStrategy {
    Continue,
    SwitchToAdhoc,   // Rewind child, switch to adhoc
    SwitchToBatches, // Rewind source+child, restart batches
    Stop,
}

// In ScoreSource trait (note: &mut self to allow internal state updates):
fn collection_strategy(&mut self, heap_count: usize, k: usize) -> CollectionStrategy;

pub struct VectorScoreSource {
    index: VecSimIndex,
    query_vector: Vec<f32>,
    query_params: VecSimQueryParams,
    batch_iterator: Option<VecSimBatchIterator>,
    timeout_ctx: TimeoutCtx,
}

impl<'index> ScoreSource<'index> for VectorScoreSource {
    type Batch = VecSimScoreBatchCursor;

    fn next_batch(&mut self) -> Result<Option<Self::Batch>, RQEIteratorError> {
        if self.timeout_ctx.is_expired() {
            return Err(RQEIteratorError::TimedOut);
        }

        let batch_iter = self.batch_iterator.get_or_insert_with(|| {
            VecSimBatchIterator::new(&self.index, &self.query_vector, &self.query_params)
        });

        if !batch_iter.has_next() {
            return Ok(None);
        }

        let reply = batch_iter.next(self.compute_batch_size());
        Ok(Some(VecSimScoreBatchCursor::new(reply)))
    }

    fn lookup_score(&mut self, doc_id: t_docId) -> Option<f64> {
        // Used in adhoc-BF mode - may need to acquire locks on tiered index
        let distance = self.index.get_distance_from(doc_id, &self.query_vector);
        if distance.is_nan() { None } else { Some(distance) }
    }

    fn num_estimated(&self) -> usize {
        self.index.size().min(self.query_params.k)
    }

    fn rewind(&mut self) {
        self.batch_iterator = None;
    }

    fn build_result(&self, doc_id: t_docId, score: f64) -> RSIndexResult<'index> {
        // Returns MetricResult or AggregateResult depending on query requirements
        RSIndexResult::metric(doc_id, score)
    }

    fn collection_strategy(&mut self, heap_count: usize, k: usize) -> CollectionStrategy {
        // Heuristic: switch to adhoc if batch selectivity is low
        if heap_count >= k {
            CollectionStrategy::Stop
        } else if self.should_switch_to_adhoc(heap_count, k) {
            CollectionStrategy::SwitchToAdhoc
        } else {
            CollectionStrategy::Continue
        }
    }
}

pub struct NumericScoreSource<'index> {
    numeric_index: &'index NumericIndex,
    ranges: NumericRangeIterator,  // Yielded in score order according to `ascending`
    current_range_idx: usize,
    range_batch_size: usize,
    ascending: bool, // Set at construction from SORTBY ASC/DESC
    timeout_ctx: TimeoutCtx,
}

impl<'index> NumericScoreSource<'index> {
    pub fn new(
        numeric_index: &'index NumericIndex,
        ranges: NumericRangeIterator,
        ascending: bool, // from query's SORTBY direction
        range_batch_size: usize,
        timeout_ctx: TimeoutCtx,
    ) -> Self {
        Self {
            numeric_index,
            ranges,
            current_range_idx: 0,
            range_batch_size,
            ascending,
            timeout_ctx,
        }
    }
}

impl<'index> ScoreSource<'index> for NumericScoreSource<'index> {
    type Batch = NumericRangeBatchCursor;

    fn next_batch(&mut self) -> Result<Option<Self::Batch>, RQEIteratorError> {
        if self.timeout_ctx.is_expired() {
            return Err(RQEIteratorError::TimedOut);
        }

        // Get next subset of ranges based on current batch size
        let batch = self.ranges.next_n(self.range_batch_size)?;
        if batch.is_empty() {
            return Ok(None);
        }

        Ok(Some(NumericRangeBatchCursor::new(batch)))
    }

    fn lookup_score(&mut self, doc_id: t_docId) -> Option<f64> {
        // Numeric adhoc lookup is not used in v1 strategy selection.
        // Keep API compatibility for possible future extension.
        let _ = doc_id;
        None
    }

    fn num_estimated(&self) -> usize {
        self.ranges.total_docs_estimate()
    }

    fn rewind(&mut self) {
        self.current_range_idx = 0;
        // Optionally adjust range_batch_size for retry
    }

    fn build_result(&self, doc_id: t_docId, score: f64) -> RSIndexResult<'index> {
        RSIndexResult::numeric(doc_id, score)
    }

    fn collection_strategy(&mut self, heap_count: usize, k: usize) -> CollectionStrategy {
        if heap_count >= k {
            return CollectionStrategy::Stop;
        }

        // Check if we should expand range and retry
        let success_ratio = self.compute_success_ratio(heap_count);
        if success_ratio < self.min_success_ratio && self.can_expand_range() {
            // Adjust parameters and rewind BEFORE returning SwitchToBatches
            self.expand_range();
            self.retry_count += 1;
            self.rewind();  // Reset iteration state
            return CollectionStrategy::SwitchToBatches;
        }

        // Numeric source does not switch to adhoc in v1.
        CollectionStrategy::Continue
    }
}

// Alternative: more explicit methods
fn build_metric_result(&self, doc_id: t_docId, score: f64) -> RSIndexResult<'index>;
fn build_aggregate_result(&self, doc_id: t_docId, score: f64, child: &RSIndexResult) -> RSIndexResult<'index>;

pub struct TopKIterator<'index, S: ScoreSource, const ASC: bool, const FILTERED: bool> { ... }

#[test] fn test_heap_insert_maintains_top_k();
#[test] fn test_heap_with_ascending_comparator();
#[test] fn test_heap_with_descending_comparator();
#[test] fn test_heap_pop_order();
#[test] fn test_heap_capacity_limit();

#[test] fn test_mock_source_batch_iteration();
#[test] fn test_mock_source_lookup_score();
#[test] fn test_mock_source_rewind();
#[test] fn test_mock_source_collection_strategy_stop();
#[test] fn test_mock_source_collection_strategy_switch_to_adhoc();
#[test] fn test_mock_source_collection_strategy_switch_to_batches();

#[test] fn test_unfiltered_direct_yield_no_heap_collection();
#[test] fn test_unfiltered_consumes_single_final_batch();
#[test] fn test_unfiltered_handles_empty_source();
#[test] fn test_unfiltered_timeout_propagates();
#[test] fn test_unfiltered_yields_sorted_by_score();

#[test] fn test_batches_intersects_with_child();
#[test] fn test_batches_rewinds_child_per_batch();
#[test] fn test_batches_switch_to_adhoc();
#[test] fn test_batches_switch_to_batches_rewinds_both();
#[test] fn test_batches_handles_empty_child();
#[test] fn test_batches_handles_disjoint_batch_and_child();
#[test] fn test_batches_timeout_mid_intersection();

#[test] fn test_adhoc_iterates_child_lookups_scores();
#[test] fn test_adhoc_skips_docs_not_in_source();
#[test] fn test_adhoc_early_termination();
#[test] fn test_adhoc_handles_child_eof();

#[test] fn test_read_triggers_collection_on_first_call();
#[test] fn test_read_yields_results_after_collection();
#[test] fn test_skip_to_not_supported_returns_error();
#[test] fn test_rewind_resets_to_not_started();
#[test] fn test_at_eof_after_all_results_yielded();
#[test] fn test_current_returns_last_yielded();

#[test] fn test_vector_source_creates_batch_iterator_lazily();
#[test] fn test_vector_source_batch_sorted_by_doc_id();
#[test] fn test_vector_source_lookup_score_returns_distance();
#[test] fn test_vector_source_rewind_clears_batch_iterator();
#[test] fn test_vector_source_switch_to_adhoc_heuristic();
#[test] fn test_vector_source_timeout_check();

#[test] fn test_numeric_source_iterates_ranges_by_value();
#[test] fn test_numeric_source_ascending_vs_descending();
#[test] fn test_numeric_source_lookup_returns_field_value();
#[test] fn test_numeric_source_expand_range_on_low_success_ratio();
#[test] fn test_numeric_source_switch_to_batches_rewinds_self();
#[test] fn test_numeric_source_max_retry_limit();

#[test] fn test_vector_topk_with_inverted_index_child();
#[test] fn test_vector_topk_with_intersection_child();
#[test] fn test_numeric_topk_with_tag_filter();
#[test] fn test_numeric_topk_ascending_descending();

# Run same queries against C and Rust implementations
def test_vector_topk_parity_unfiltered():
def test_vector_topk_parity_with_filter():
def test_vector_topk_parity_mode_switching():
def test_numeric_topk_parity_ascending():
def test_numeric_topk_parity_descending():
def test_numeric_topk_parity_with_filter():

def test_rust_vector_topk_profile_shows_correct_mode():
def test_rust_numeric_topk_profile_shows_correct_mode():
def test_rust_topk_profile_shows_strategy_switches():

#[test] fn prop_topk_always_returns_k_or_fewer_results();
#[test] fn prop_topk_results_sorted_by_score();
#[test] fn prop_topk_results_subset_of_child_intersection();
#[test] fn prop_batches_and_adhoc_produce_same_results();

// TopKHeap operations
fn bench_heap_insert(c: &mut Criterion);
fn bench_heap_pop_all(c: &mut Criterion);

// TopKIterator modes
fn bench_unfiltered_10k_docs(c: &mut Criterion);
fn bench_batches_10k_docs_1k_child(c: &mut Criterion);
fn bench_adhoc_10k_child(c: &mut Criterion);

// Strategy switching overhead
fn bench_switch_batches_to_adhoc(c: &mut Criterion);
fn bench_switch_to_batches_with_rewind(c: &mut Criterion);

// Comparison with C (via FFI)
fn bench_rust_vs_c_unfiltered(c: &mut Criterion);
fn bench_rust_vs_c_batches(c: &mut Criterion);