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 _suffix_predicate(*suffixes):
    def predicate(file):
        return file.endswith(suffixes)

    return predicate


def _prefix_predicate(*prefixes):
    def predicate(file):
        return file.startswith(prefixes)

    return predicate


is_examples = _prefix_predicate("examples/", "tests/poc/")
is_poc = _prefix_predicate("tests/poc")
is_0day = _prefix_predicate("tests/0day")
is_pipecase = _prefix_predicate("tests/pipeline")
is_owners = _prefix_predicate("OWNERS")
is_license = _prefix_predicate("LICENSE")
is_ut = _prefix_predicate("tests/unit_tests")
is_markdown = _suffix_predicate(".md")
is_image = _suffix_predicate(".jpg", ".png")
is_txt = _suffix_predicate(".txt")
is_json = _suffix_predicate(".json")


def is_no_suffix(file):
    return Path(file).suffix == ''


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,
]


def skip_ci(files, skip_conds):
    return all(any(cond(file) for cond in skip_conds) for file in files)


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)
    return skip_ci(file_list, SKIP_CONDS)


UT_FILTER_CONDS = [is_ut, is_markdown]


def filter_exec_ut(raw_txt_file):
    file_list = read_files_from_txt(raw_txt_file)
    if all(any(cond(f) for cond in UT_FILTER_CONDS) for f in file_list):
        return True, file_list
    return False, None


def acquire_exitcode(command):
    return os.WEXITSTATUS(os.system(command))  # pylint: disable=no-member


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


class UTTest:
    GROUP_COUNT = 4

    def __init__(self):
        self.base_dir = Path(__file__).absolute().parents[1]
        self.test_dir = self.base_dir / 'tests'

    def _collect_all_ut_files(self):
        all_ut_files = []
        unit_tests_dir = self.test_dir / "unit_tests"
        if unit_tests_dir.exists():
            for py_file in unit_tests_dir.rglob("*.py"):
                all_ut_files.append(str(py_file))
        return all_ut_files

    def _resolve_files_to_run(self, all_ut_files, raw_txt_file):
        if raw_txt_file is not None and os.path.exists(raw_txt_file):
            filtered_results = filter_exec_ut(raw_txt_file)
            print(f"[LOG run_ut] filter_exec_ut returned: {filtered_results[0]}")

            if filtered_results[0]:
                filtered_files = filtered_results[1]
                full_path = [str(self.base_dir / file) for file in filtered_files]
                existing_ut_files = [file for file in full_path if os.path.exists(file) and file.endswith(".py")]
                print("[LOG run_ut] changed UT/MD files count: {len(existing_ut_files)}")
                return existing_ut_files
            else:
                print("[LOG run_ut] no UT/MD changes, will run all unit_tests")
                return all_ut_files
        else:
            print("[LOG run_ut] no raw_txt_file, will run all unit_tests")
            return all_ut_files

    def _assign_files_for_rank(self, sorted_files, rank):
        total = len(sorted_files)
        if rank == 0:
            print(f"[LOG run_ut] rank=0, running ALL {total} files")
            return " ".join(sorted_files)
        else:
            group_files = sorted_files[rank - 1 :: self.GROUP_COUNT]
            print(
                f"[LOG run_ut] rank={rank}, total={total}, interleaved group size={len(group_files)}, files={group_files}"
            )
            return " ".join(group_files)

    def _execute_pytest(self, ut_files_str):
        print(f"[LOG run_ut] final ut_files_str: {ut_files_str}")
        command = f"pytest -x --log-cli-level=INFO {ut_files_str}"
        code = acquire_exitcode(command)
        if code == 0:
            print("UT test success")
        else:
            print("UT fail")
            sys.exit(1)

    def run_ut(self, rank=0, raw_txt_file=None):
        print(f"[LOG run_ut] rank={rank}, raw_txt_file={raw_txt_file}")

        all_ut_files = self._collect_all_ut_files()
        print(f"[LOG run_ut] total all_ut_files in unit_tests: {len(all_ut_files)}")

        files_to_run = self._resolve_files_to_run(all_ut_files, raw_txt_file)
        sorted_files = sorted(files_to_run)
        ut_files_str = self._assign_files_for_rank(sorted_files, rank)
        self._execute_pytest(ut_files_str)


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 = str(parent_dir / "modify.txt")

    skip_signal = choose_skip_ci(modified_txt_file)
    if skip_signal:
        print("Skipping CI")
    else:
        ut = UTTest()
        ut.run_ut(options.rank, modified_txt_file)
        print("UT test finish")