#!/usr/bin/env python

"""
Remove the DONE label from a dag file from a <stage> and after this <stage>
"""


if __name__ == '__main__':
    import sys
    FILENAME = sys.argv[1]
    STAGE = sys.argv[2]
    # Make a backup
    BACKUP_FILE = '%s.backup' % FILENAME
    try:
        BACKUP = open(BACKUP_FILE)
        BACKUP.close()
        print 'Remove/Rename old backup first [%s] since I will not delete it' % BACKUP_FILE
        sys.exit(1)
    except IOError:
        pass
    print 'Backup file: %s' % BACKUP_FILE
    import shutil
    shutil.copyfile(FILENAME, BACKUP_FILE)
    # Parse and filter
    INPUT_FILE = open(BACKUP_FILE)
    OUTPUT = open(FILENAME, 'w')
    STATE = 0
    for line in INPUT_FILE:
        if STATE == 0:
            if STAGE in line:
                STATE = 1
        if STATE == 1:
            line = line.replace(' DONE', '')
        OUTPUT.write(line)
    INPUT_FILE.close()
    OUTPUT.close()
