import os
import sys
from pathlib import Path
import argparse


def read_files_from_txt(txt_file):
    with open(txt_file, "r", encoding='utf-8') as f:
        return [line.strip() for line in f.readlines()]


def is_examples(file):
    return file.startswith("examples/") or file.startswith("tests/poc/")


def is_poc(file):
    return file.startswith("tests/poc")


def is_0day(file):
    return file.startswith("tests/0day")


def is_pipecase(file):
    return file.startswith("tests/pipeline")


def is_markdown(file):
    return file.endswith(".md")


def is_image(file):
    return file.endswith(".jpg") or file.endswith(".png")


def is_txt(file):
    return file.endswith(".txt")


def is_json(file):
    return file.endswith(".json")


def is_owners(file):
    return file.startswith("OWNERS")


def is_license(file):
    return file.startswith("LICENSE")


def is_ut(file):
    return file.startswith("tests/unit_tests")


def is_no_suffix(file):
    return os.path.splitext(file)[1] == ''


def skip_ci(files, skip_conds):
    for file in files:
        if not any(condition(file) for condition in skip_conds):
            return False
    return True


def choose_skip_ci(raw_txt_file):
    if not os.path.exists(raw_txt_file):
        return False

    file_list = read_files_from_txt(raw_txt_file)
    skip_conds = [
        is_examples,
        is_pipecase,
        is_markdown,
        is_image,
        is_txt,
        is_owners,
        is_license,
        is_no_suffix,
        is_poc,
        is_0day,
        is_json,
    ]

    return skip_ci(file_list, skip_conds)


def filter_exec_ut(raw_txt_file):
    file_list = read_files_from_txt(raw_txt_file)
    filter_conds = [is_ut, is_markdown]
    for file in file_list:
        if not any(condition(file) for condition in filter_conds):
            return False, None
    return True, file_list


def acquire_exitcode(command):
    exitcode = os.system(command)
    real_code = os.WEXITSTATUS(exitcode)
    return real_code


# =============================
# UT test, run with pytest
# =============================


class UTTest:
    def __init__(self):
        self.base_dir = Path(__file__).absolute().parents[1]
        self.test_dir = os.path.join(self.base_dir, 'tests')
        self.ut_files = os.path.join(self.base_dir, self.test_dir, "unit_tests")

    def run_ut(self, rank=0, raw_txt_file=None):
        if raw_txt_file is not None and os.path.exists(raw_txt_file):
            filtered_results = filter_exec_ut(raw_txt_file)

            if filtered_results[0]:
                filtered_files = filtered_results[1]
                full_path = [os.path.join(self.base_dir, file) for file in filtered_files]
                exsit_ut_files = [file for file in full_path if os.path.exists(file) and file.endswith(".py")]
                if rank == 0:
                    self.ut_files = " ".join(exsit_ut_files)
                else:
                    # Split UT nodes into 4 groups and test including to rank index
                    sorted_files = sorted(exsit_ut_files)
                    total = len(sorted_files)
                    group_size = (total + 3) // 4
                    start = (rank - 1) * group_size
                    end = min(start + group_size, total)
                    self.ut_files = " ".join(sorted_files[start:end])

        command = f"pytest -x --log-cli-level=INFO {self.ut_files}"
        code = acquire_exitcode(command)
        if code == 0:
            print("UT test success")
        else:
            print("UT failed")
            sys.exit(1)


def run_tests(rank, raw_txt_file):
    ut = UTTest()
    if filter_exec_ut(raw_txt_file)[0]:
        ut.run_ut(rank, raw_txt_file)
    else:
        ut.run_ut()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Control needed ut cases')
    # 4 ut nodes in total currently
    parser.add_argument('--rank', default=0, type=int, help='Index of current ut nodes')
    options = parser.parse_args()

    parent_dir = Path(__file__).absolute().parents[2]
    modified_txt_file = os.path.join(parent_dir, "modify.txt")

    skip_signal = choose_skip_ci(modified_txt_file)
    if skip_signal:
        print("Skipping CI")
    else:
        run_tests(options.rank, modified_txt_file)