import os
import random
import subprocess
import sys

MEGA = 2**20

def main():
    
    if len(sys.argv) < 6:
        print 'ERROR: not enough arguments provided'
        return -1
    
    # Assign input arguments    
    write_path = sys.argv[1]
    min_chunks_per_file = int(sys.argv[2])
    max_chunks_per_file = int(sys.argv[3])
    chunk_size = int(sys.argv[4])
    max_write_amount = int(sys.argv[5])
    
    # Do some basic input validation for S&G
    if min_chunks_per_file > max_chunks_per_file:
        print 'ERROR: specified minimum chunks per file is larger than '\
              'specified maximum chunks per file'
        return -1
        
    if chunk_size > max_write_amount:
        print 'ERROR: specified chunk size is larger than specified maximum '\
              'write amount'
    
    amount_written = 0
    file_counter = 0
    chunk_idx_counter = 0
    total_written = 0
    keep_running = True
    
    while keep_running:
        
        # Create the file name and randomly determine the number of blocks
        file_name = '%08X' % file_counter
        file_counter += 1
        chunks_per_file = random.randint(min_chunks_per_file, max_chunks_per_file)
        file_size = chunks_per_file * chunk_size
        
        # Make sure we don't write more than intended
        if (total_written + file_size) > max_write_amount:
            chunks_per_file = (max_write_amount - total_written) / chunk_size
            file_size = chunks_per_file * chunk_size
            keep_running = False
        
        # Print out status
        sys.stdout.write('Creating file "%s" with %5d chunks - %7d MB ... ' % 
                         (file_name, chunks_per_file, file_size))
        sys.stdout.flush()
        
        # Make the "file" directory
        path = os.path.join(write_path, file_name)
        os.mkdir(path)
        
        # Create each of the actual chunk files
        for chunk in xrange(chunks_per_file):
            chunk_name = '%016X' % chunk_idx_counter
            chunk_idx_counter += 1
            chunk_path = os.path.join(path, chunk_name)
            command = 'head -c %d /dev/zero > %s' % (chunk_size * MEGA, chunk_path)
            subprocess.call(command, shell=True)
            
        sys.stdout.write('DONE\n')
    
        total_written += file_size
        print '   Total written: %d MB' % total_written    
    
if __name__ == '__main__':
    sys.exit(main())