#!/usr/bin/python3


"""
The script runs program massively on all files from a directory, concurrently and remembering its progress.
To set program to be run edit function make_command().
"""

# --------  CONFIG


def make_command(input_files, output_files, args):
    """ Function creating command in form of standard argv, e.g. ['ls', '~/Documents', '-l', '-a']. """
    
    PROGRAM = '/home/dekakaruk/Źródła/supermatrix/bin/tools/architect2/mwefinder/mwefinder'
    return [PROGRAM, '--corpus'] + input_files + ['--outputs'] + output_files + ['--quiet', 'true'] + args


# --------  /CONFIG










import argparse
import os
import sys
import subprocess
import shutil
import itertools
import time
import math
import tempfile

DEVNULL = open(os.devnull, 'wb')




def load_args():
    """ Read command-line aruments. """
    parser = argparse.ArgumentParser(description='Run program for all files in the directory, recursively.')
    
    parser.add_argument('--in',  help='Input directory. If index is used it still should be provided, for it is used to compute output paths.', dest='input')
    parser.add_argument('--out', help='Output directory.', dest='output')
    parser.add_argument('--index', help='Index of all files (not directories!) to be converted. In each line of index shall be one full path to file.')
    parser.add_argument('--interval', help='After how many files the progress should be saved.', type = int, default = 100)
    parser.add_argument('--threads', help='How many threads should be run at one time.', type = int, default = 3)
    parser.add_argument('--prefix', help='Path and prefix of all backup files', default = './convert_dir_')
    parser.add_argument('args', help='Arguments for underlying program. If there be any nonpositional parameters, the whole list must be preceded by \' -- \'.', nargs=argparse.REMAINDER)
    
    args = parser.parse_args()
    if '--' in args.args:
        args.args.remove('--')
    return args


def log(msg):
    """ Print msg on stdout, preceded by current time. """
    print('[' + time.strftime('%c') + ']: ' + msg)
    sys.stdout.flush()


def replace_root(path, old_root, new_root):
    """ Replace old prefix of path with new one. """
    path = os.path.normpath(path)
    old_root = os.path.normpath(old_root)
    new_root = os.path.normpath(new_root)
    return new_root + path[len(old_root):]
    

def chunks(list1, list2, n):
    """ Iterate both lists at once, yield parts of size n. """
    for i in range(0, min(len(list1), len(list2)), n):
        yield list1[i:i+n], list2[i:i+n]


def convert(inputs, args):
    """ Convert input files. """
    ni = []
    for f in inputs:
        if not os.path.isfile(f):
            print('Error: "' + f + '" does not exist or is not a regular file!', file=sys.stderr)
            sys.stderr.flush()
        else:
            ni.append(f)
    inputs = ni
    
    outputs = [replace_root(i, args.input, args.output) for i in inputs]
    
    files_per_thread = math.ceil(args.interval / args.threads)
    
    processes = []
    for ins, outs in chunks(inputs, outputs, files_per_thread):
        if not ins:
            continue
        command = make_command(list(ins), outs, args.args)
        f = tempfile.TemporaryFile('w+')
        processes.append((subprocess.Popen(command, stderr=f, stdout=DEVNULL), f, command))
    
    for p, f, c in processes:
        ret = p.wait()
        f.seek(0)
        t = f.read()
        if ret or t:
            print("Error during command " + ' '.join(c) + ":\n", file=sys.stderr)
            print("[return code: " + str(ret) +"]\n", file=sys.stderr)
            print(t, file=sys.stderr)
            print("---\n", file=sys.stderr)
            
        f.close()
    sys.stderr.flush()
    




INDEX_FILE = 'index'  # file with index of all files (not directories!) to process
PROGRESS_FILE = 'progress'  # file in which progress will be remembered
PROGRESS_FILE_NEW = 'progress_new'  # temporary file used to safely swap progress
BAD_OLD_FLAG_FILE = 'old_bad'  # if such file exists PROGRESS_FILE may be incorrect and therefore PROGRESS_FILE_NEW shall be used


def save_progress(current, index, prefix):
    """ Save current progress. """
    log('Saving progress (' + str(current) + ' files, ' + str(round(current * 100 / len(index), 3)) + '%)')
    
    with open(prefix + PROGRESS_FILE_NEW, 'w') as f:
        f.write(str(current))
    
    open(prefix + BAD_OLD_FLAG_FILE, 'a').close()
    
    shutil.copyfile(prefix + PROGRESS_FILE_NEW, prefix + PROGRESS_FILE)
    
    os.remove(prefix + BAD_OLD_FLAG_FILE)
    os.remove(prefix + PROGRESS_FILE_NEW)
    
    log('Progress succesfully saved\n')


def build_index(input):
    """ Create index of all files in input directory. """
    return [os.path.join(root, f) for root, dirs, files in os.walk(input) for f in files]


def write_index(index, path):
    """ Save index to file. """
    with open(path, 'w') as f:
        for file in index:
            f.write(file + '\n')


def load_index(path):
    """ Load index from file. """
    lst = []
    with open(path) as f:
        for line in f:
            pth = line.strip()
            if pth:
                lst.append(pth)
    return lst


def load_progress(args):
    """ Load progress if it is possible or start ab ovo. Return pair (current_file_number, index). """
    
    prefix = args.prefix
    
    current_file = 0
    index = []
    
    if os.path.isfile(prefix + PROGRESS_FILE) or (os.path.isfile(prefix + PROGRESS_FILE_NEW) and os.path.isfile(prefix + BAD_OLD_FLAG_FILE)):
        log("Loading progress from backup. If 'tis not your will run 'rm " + prefix + "*'.\n" )
        
        if os.path.isfile(prefix + BAD_OLD_FLAG_FILE):
            shutil.copyfile(prefix + PROGRESS_FILE_NEW, prefix + PROGRESS_FILE)
        
        with open(prefix + PROGRESS_FILE, 'r') as f:
            current_file = eval(f.read())
    
    if args.index:
        index = load_index(args.index)
    elif os.path.isfile(prefix + INDEX_FILE):
        index = load_index(prefix + INDEX_FILE)
    else:
        index = build_index(args.input)
        write_index(index, prefix + INDEX_FILE)
    
    return current_file, index


def clean_up(prefix):
    """ Remove backup files. """
    if os.path.exists(prefix + INDEX_FILE):
        os.remove(prefix + INDEX_FILE)
    if os.path.exists(prefix + PROGRESS_FILE):
        os.remove(prefix + PROGRESS_FILE)
    if os.path.exists(prefix + PROGRESS_FILE_NEW):
        os.remove(prefix + PROGRESS_FILE_NEW)
    if os.path.exists(prefix + BAD_OLD_FLAG_FILE):
        os.remove(prefix + BAD_OLD_FLAG_FILE)
    


def process(args, index, current_file):
    """ Do everything. """
    
    log('Starting conversion\n')
    
    while current_file < len(index):
        convert(index[current_file : current_file + args.interval], args)
        current_file += args.interval
        
        if current_file < len(index):
            save_progress(current_file, index, args.prefix)
    
    log('All is converted\n')



if __name__ == '__main__':
    args = load_args()
    
    log('Starting\n')
    
    current_file, index = load_progress(args)
    process(args, index, current_file)
    clean_up(args.prefix)
    
    log('All is done\n')
    
    
