module Scrivito
class ObjSearchEnumerator
  class BatchIterator
    def initialize(workspace, query, initial_batch)
      @query = query
      @executor = ObjSearchEnumerator::QueryExecutor.new(workspace)
      @all_fetched_ids = Set.new
      @initial_batch = initial_batch
    end

    attr_reader :total

    def each
      current_batch = initial_batch || fetch_batch

      loop do
        yield current_batch
        break unless current_batch.continuation
        current_batch = fetch_batch(current_batch.continuation)
      end

      @total = @all_fetched_ids.size
    end

    private

    attr_reader :query, :executor, :initial_batch

    def fetch_batch(continuation=nil)
      batch = executor.call(query, continuation, @all_fetched_ids)
      @all_fetched_ids += batch.objs.map(&:id)
      batch
    end
  end
end
end