#! /usr/bin/env python
import os
import re
from abjad.tools import systemtools


def fix_test_case_block_comments():
    total_test_modules = 0
    total_broken_modules = 0
    total_broken_comments = 0
    # TODO: extend to check for r""" too
    open_comment_re = re.compile(r"(\s*)(r''')(.+)")
    close_comment_re = re.compile(r"(\s*)([}|>]+)(.*)(''')")
    for directory, subdirectory_names, file_names in os.walk('.'):
        test_modules = []
        for file_name in file_names:
            if file_name.startswith('test_') and file_name.endswith('.py'):
                test_modules.append(file_name)
        total_test_modules += len(test_modules)
        for test_module in test_modules:
            total_test_modules += 1
            full_module_name = os.path.join(directory, test_module)
            contained_broken_comment = False
            fp = file(full_module_name, 'r')
            new_lines = []
            lines = fp.readlines()
            new_lines = []
            previous_line = ''
            for line in lines:
                open_match = open_comment_re.match(line)
                close_match = close_comment_re.match(line)
                if previous_line.startswith('def '):
                    new_lines.append(line)
                elif open_match:
                    total_broken_comments += 1
                    contained_broken_comment = True
                    whitespace, open_quote, expr = open_match.groups()
                    new_line = whitespace + open_quote + '\n'
                    new_lines.append(new_line)
                    new_line = whitespace + expr + '\n'
                    new_lines.append(new_line)
                elif close_match:
                    groups = close_match.groups()
                    whitespace, close_bracket, chord_duration, close_quote = \
                        groups
                    new_line = \
                        whitespace + close_bracket + chord_duration + '\n'
                    new_lines.append(new_line)
                    new_line = whitespace + close_quote + '\n'
                    new_lines.append(new_line)
                else:
                    new_lines.append(line)
                previous_line = line
            fp.close()
            if contained_broken_comment:
                total_broken_modules += 1
                new_file_string = ''.join(new_lines)
                fp = open(full_module_name, 'w')
                fp.write(new_file_string)
                fp.close()

    print 'Total test modules:       {}'.format(total_test_modules)
    print 'Total fixed test modules: {}'.format(total_broken_modules)
    print 'Total fixed comments:     {}'.format(total_broken_comments)


if __name__ == '__main__':
    systemtools.IOManager.clear_terminal()
    fix_test_case_block_comments()
    print
