import argparse
import xml.etree.ElementTree as ET
from termcolor import colored


def parse_args():
    parser = argparse.ArgumentParser(description='Parse Unity test results XML files and display test results.')
    parser.add_argument('editmode_xml_file', type=str, help='Path to Unity EditMode test results XML file')
    parser.add_argument('playmode_xml_file', type=str, help='Path to Unity PlayMode test results XML file')
    return parser.parse_args()


def get_test_suites(xml_file, mode):
    tree = ET.parse(xml_file)
    root = tree.getroot()

    test_suites = []
    # get all test_suite elements with type assembly
    for test_suite in root.iter('test-suite'):
        if test_suite.attrib['type'] == 'Assembly':
            test_suites.append(test_suite)

    return test_suites

def get_tests_cases(test_suite):
    test_cases = []
    for test_case in test_suite.iter('test-case'):
        test_cases.append({
            'name': test_case.attrib['name'],
            'classname': test_case.attrib['classname'],
            'result': test_case.attrib['result'],
            'duration': test_case.attrib['duration'],
        })
    return test_cases


def print_test_results(test_suites, mode):
    print(f'\n================ {mode} Tests ================')
    all_tests = []
    for test_suite in test_suites:
        test_suite_name = test_suite.attrib['name']
        print("\n" + colored(test_suite_name, 'cyan'))

        tests = get_tests_cases(test_suite)
        all_tests += tests
        if tests:
            passed_tests = []
            failed_tests = []
            for test in tests:
                test_name = test['name']
                test_classname = test['classname']
                test_result = test['result']
                test_duration = test['duration']
                test_duration_str = '{:.3f}s'.format(float(test_duration))
                if test_result == 'Passed':
                    print(colored('✓ PASS', 'green'), test_name, '({} {})'.format(test_classname, test_duration_str))
                    passed_tests.append(test)
                elif test_result == 'Failed':
                    print(colored('✗ FAIL', 'red'), test_name, '({} {})'.format(test_classname, test_duration_str))
                    failed_tests.append(test)
                else:
                    print(colored('⚠ SKIP', 'yellow'), test_name, '({} {})'.format(test_classname, test_duration_str))
    return all_tests

def get_combined_results(edit_results, play_results):
    combined_results = edit_results + play_results
    return combined_results

def print_summary(combined_results):
    passed_tests = []
    failed_tests = []
    for test in combined_results:
        test_result = test['result']
        if test_result == 'Passed':
            passed_tests.append(test)
        elif test_result == 'Failed':
            failed_tests.append(test)

    print('\n=================== Summary ===================')
    print(colored('✓ Passed tests', 'green') + ': {} / {}'.format(len(passed_tests), len(combined_results)))
    print(colored('✗ Failed tests', 'red') + ': {} / {}\n'.format(len(failed_tests), len(combined_results)))

    if len(failed_tests) > 0:
        # exit with error
        exit(1)

if __name__ == '__main__':
    args = parse_args()

    edit_suites = get_test_suites(args.editmode_xml_file, 'EditMode')
    play_suites = get_test_suites(args.playmode_xml_file, 'PlayMode')

    edit_tests = print_test_results(edit_suites, 'EditMode')
    play_tests = print_test_results(play_suites, 'PlayMode')

    combined_results = get_combined_results(edit_tests, play_tests)
    print_summary(combined_results)