#!/usr/bin/env python3
import argparse
import re
from typing import TextIO
import xml.etree.ElementTree as ET

from log_parser import ctest_log_parser, log_reader, GTEST_MARK, YUNIT_MARK
from mute_utils import mute_target, remove_failure, update_suite_info, MuteTestCheck


def find_targets_to_remove(log_fp):
    target_with_tests = set()
    for target, reason, buf in ctest_log_parser(log_fp):
        if reason != "Failed":
            continue

        for line in buf:
            if line.startswith((GTEST_MARK, YUNIT_MARK)):
                target_with_tests.add(target)
                break
    return target_with_tests


shard_suffix_re = re.compile(r"_\d+$")


def strip_shardname(testcase):
    name = testcase.get('classname')
    classname = shard_suffix_re.sub('', name)

    testcase.set('classname', classname)


def postprocess_ctest(log_fp: TextIO, ctest_junit_report, is_mute_shard, dry_run):
    to_remove = find_targets_to_remove(log_fp)
    tree = ET.parse(ctest_junit_report)
    root = tree.getroot()
    n_remove_failures = n_skipped = 0

    for testcase in root.findall("testcase"):
        target = testcase.attrib["classname"]

        strip_shardname(testcase)

        if is_mute_shard(target):
            if mute_target(testcase):
                print(f"mute {target}")
                testcase.set("status", "run")  # CTEST specific
                n_remove_failures += 1
                n_skipped += 1
        elif target in to_remove:
            print(f"set {target} as passed")
            n_remove_failures += 1
            remove_failure(testcase)

    update_suite_info(root, n_remove_failures, n_skipped=n_skipped)

    print(f"{'(dry-run) ' if dry_run else ''}update {ctest_junit_report}")

    if not dry_run:
        tree.write(ctest_junit_report, xml_declaration=True, encoding="UTF-8")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--dry-run", action="store_true", default=False)
    parser.add_argument("--filter-file", required=False)
    parser.add_argument("--decompress", action="store_true", default=False, help="decompress ctest log")
    parser.add_argument("ctest_log", type=str)
    parser.add_argument("ctest_junit_report")
    args = parser.parse_args()

    log = log_reader(args.ctest_log, args.decompress)
    is_mute_shard = MuteTestCheck(args.filter_file)
    postprocess_ctest(log, args.ctest_junit_report, is_mute_shard, args.dry_run)


if __name__ == "__main__":
    main()