# -*- coding: utf-8 -*-
import sys
import argparse
import os
import logging
import xml.etree.ElementTree as ET


class Twerk(object):

    def __init__(self):
        parser = argparse.ArgumentParser(prog='twerk', description='Command-line utility tool for various QE tasks', epilog='See: \'https://git-aws.internal.justin.tv/qa/twerk\' for more info.')
        parser.add_argument('command', help='Subcommand to run')

        args = parser.parse_args(sys.argv[1:2])
        if not hasattr(self, args.command):
            print 'Unrecognized command'
            parser.print_help()
            exit(1)
        # use dispatch pattern to invoke method with same name
        getattr(self, args.command)()

    def get_tests(self):
        output_file = 'test_results.xml'
        parser = argparse.ArgumentParser(description='Merges JUnit files into one single file \'%s\'.  If an error is encountered with a file while merging, then the file won\'t be merged.' % output_file)
        parser.add_argument('files', type=str, nargs='+', help='List of files or directories to merge.  If a directory is provided then the directory will be recursively scanned for .xml files.')
        args = parser.parse_args()

        failures = 0
        tests = 0
        errors = 0
        time = 0.0
        subtrees = []
        xml_files = []

        for entry in args.files:
            if os.path.isdir(entry):
                for subdir, dirs, files in os.walk(entry):
                    for dir_file in files:
                        if not dir_file.endswith('.xml'):
                            continue
                        xml_file = os.path.join(subdir, dir_file)
                        xml_files.append(xml_file)

        for file_name in xml_files:
            if not file_name.endswith('.xml'):
                continue
            try:
                tree = ET.parse(file_name)
                test_suite = tree.getroot()
                failures += int(test_suite.attrib['failures'])
                tests += int(test_suite.attrib['tests'])
                errors += int(test_suite.attrib['errors'])
                time += float(test_suite.attrib['time'])
                subtrees.append(test_suite)
            except:
                logging.warning('Failed to merge: \'%s\'' % (file_name))

        new_root = ET.Element('testsuite')
        new_root.attrib['failures'] = '%s' % failures
        new_root.attrib['tests'] = '%s' % tests
        new_root.attrib['errors'] = '%s' % errors
        new_root.attrib['time'] = '%s' % time
        suites = ET.Element('testsuites')
        for subtree in subtrees:
            suites.append(subtree)
        new_root.append(suites)
        new_tree = ET.ElementTree(new_root)
        new_tree.write(output_file)


def main():
    Twerk()
