# ---------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# ---------------------------------------------------------
import argparse
import contextlib
import glob
import json
import os
import re
import subprocess
import sys
from datetime import datetime
from pathlib import Path

import dotenv

# pylint: disable=consider-using-with


def normalize_test_name(test_name):
    if "[" in test_name:
        test_name = test_name.split("[")[0]
    return test_name.strip()


def location_to_test_name(location):
    test_path, _, test_func = location
    test_class_name, test_func_name = test_func.split(".", 1)
    test_class = test_path.split(os.path.sep, 3)[-1] + "::" + test_class_name
    return test_class + "::" + test_func_name


def extract_test_location(test_name):
    splitor_num = test_name.count("::")
    if splitor_num <= 1:
        return test_name, None, None
    if splitor_num == 2:
        test_class, test_func_name = test_name.rsplit("::", 1)
        m = re.match(r"(\w+)\[(\w+)]", test_func_name)
        if m:
            test_func_name, test_param = m.groups()
        else:
            test_param = None
        return test_class, test_func_name, test_param
    raise ValueError(f"Invalid test name: {test_name}")


def load_tests_from_file(input_file):
    tests_to_run = set()
    with open(input_file, "r", encoding="utf-8") as f:
        for line in f:
            if len(line) < 1 or line[0] in ["#", ";"]:
                continue
            tests_to_run.add(line.strip())
    return tests_to_run


@contextlib.contextmanager
def update_dot_env_file(env_override):
    """Update env file with env_override, and restore it after the context is exited.
    Support bool variable only for now.

    :param: env_override
    :type: dict
    :return: None
    """
    env_file = dotenv.find_dotenv(raise_error_if_not_found=True)
    print(f"Updating env file: {env_file}")
    origin_env_content = None
    try:
        with open(env_file, "r", encoding="utf-8") as f:
            origin_env_content = f.read()
            env_vars = [line.strip() for line in origin_env_content.splitlines() if line.strip()]
        for key, value in env_override.items():
            if isinstance(value, bool):
                target_line = f"{key}='true'"
                for i, line in enumerate(env_vars):
                    if line == target_line and not value:
                        env_vars[i] = f"#{target_line}"
                    elif re.match(rf"# *{target_line}", line) and value:
                        env_vars[i] = f"{target_line}"
        with open(env_file, "w", encoding="utf-8") as f:
            f.write("\n".join(env_vars))
        yield
    finally:
        if origin_env_content is not None:
            with open(env_file, "w", encoding="utf-8") as f:
                f.write(origin_env_content)


def run_simple(
    tests_to_run,
    working_dir,
    extra_params,
    *,
    is_live_and_recording,
    log_file_path,
    log_suffix=None,
    log_in_json=False,
):
    print(f"Running {len(tests_to_run)} tests under {working_dir}: ")
    for test_name in tests_to_run:
        print(test_name)

    if log_file_path and log_suffix:
        log_file_path = log_file_path.with_suffix(log_file_path.suffix + log_suffix)

    if log_in_json:
        if log_file_path is None:
            raise ValueError("log_file_path must be specified when log_in_json is True")
        stdout = None
        json_log_file_path = log_file_path.with_suffix(log_file_path.suffix + ".log")
    else:
        stdout = open(log_file_path.with_suffix(log_file_path.suffix + ".txt"), "w", encoding="utf-8")
        json_log_file_path = None

    with update_dot_env_file(
        {"AZURE_TEST_RUN_LIVE": is_live_and_recording, "AZURE_SKIP_LIVE_RECORDING": not is_live_and_recording},
    ):
        for test_class, keyword_param in reorganize_tests(tests_to_run):
            tmp_extra_params = extra_params + keyword_param
            if log_in_json:
                # use a temp json file to avoid overwriting the final log file
                temp_log_file_path = json_log_file_path.with_stem("temp")
                tmp_extra_params += ["--report-log", temp_log_file_path.as_posix()]

            subprocess.run(
                [
                    sys.executable,
                    "-m",
                    "pytest",
                    test_class,
                ]
                + tmp_extra_params,
                cwd=working_dir,
                stdout=stdout,
                check=False,
            )
            if log_in_json:
                # append temp json file to the final log file
                with open(json_log_file_path, "a", encoding="utf-8") as f:
                    f.write(temp_log_file_path.read_text())
    if stdout is not None:
        stdout.close()
        print(log_file_path.with_suffix(log_file_path.suffix + ".txt").read_text())

    return json_log_file_path


def reorganize_tests(tests_to_run):
    reorganized_tests = {}
    for test_name in tests_to_run:
        test_class, test_func_name, test_param = extract_test_location(test_name)

        if test_func_name is None:
            # Register all tests in test_class
            reorganized_tests[test_class] = None
            continue
        if test_class not in reorganized_tests:
            reorganized_tests[test_class] = {}
        if reorganized_tests[test_class] is None:
            # All tests in test_class have been registered
            continue

        if test_param is None:
            # Register all params for test_class::test_func_name
            reorganized_tests[test_class][test_func_name] = None
            continue
        if test_func_name not in reorganized_tests[test_class]:
            reorganized_tests[test_class][test_func_name] = []
        if reorganized_tests[test_class][test_func_name] is None:
            # All params for test_class::test_func_name have been registered
            continue

        reorganized_tests[test_class][test_func_name].append(test_param)

    # re-run the tests with recording mismatch in live mode
    for test_class, test_info in reorganized_tests.items():
        if test_info is None:
            yield test_class, []
            continue
        keys = []
        for test_func_name, test_params in test_info.items():
            if test_params is not None:
                # TODO: problem in using parametrize with pytest-xdist
                # keys.append(f"{test_func_name}[{'-'.join(test_params)}]")
                keys.append(f"{test_func_name}")
            else:
                keys.append(test_func_name)
        keyword_param = ["-k", " or ".join(keys)]
        yield test_class, keyword_param


def get_base_log_path(working_dir, *, create_new=True):
    log_dir = working_dir / "scripts" / "tmp"
    if not create_new:
        logs = sorted(glob.glob(str(log_dir / "pytest.*.first.log")))
        if len(logs) == 0:
            raise RuntimeError("No previous run log file found")
        return Path(logs[-1][: -len(".first.log")])
    log_file_path = log_dir / "pytest.{}".format(datetime.now().strftime("%Y%m%d%H%M%S"))
    log_file_path.parent.mkdir(parents=True, exist_ok=True)
    return log_file_path


def get_failed_tests(log_file_path):
    tests_failed_with_recording_mismatch = []
    failed_tests = []
    with open(log_file_path, "r", encoding="utf-8") as f:
        for line in f:
            node = json.loads(line)
            if "outcome" not in node:
                continue
            if node["outcome"] != "failed":
                continue
            test_name = location_to_test_name(node["location"])
            failed_tests.append(test_name)
            msg = node["longrepr"]["reprcrash"]["message"]
            if "NotFound" in msg:
                tests_failed_with_recording_mismatch.append(test_name)
    return failed_tests, tests_failed_with_recording_mismatch


def run_tests(args, extras):
    if args.file:
        tests_to_run = load_tests_from_file(args.file)
    elif args.name:
        tests_to_run = [args.name]
    elif args.skip_first_run and args.record_mismatch:
        # load failed tests from last run log
        tests_to_run = []
    else:
        raise ValueError("Must specify either --file or --name")
    skip_first_run = args.skip_first_run
    record_mismatch = args.record_mismatch

    working_dir = Path(__file__).parent.parent
    log_file_path = get_base_log_path(working_dir, create_new=not skip_first_run)

    if skip_first_run:
        json_log_file_path = log_file_path.with_suffix(log_file_path.suffix + ".first.log")
    else:
        json_log_file_path = run_simple(
            tests_to_run,
            working_dir,
            extras + ["--disable-warnings", "--disable-pytest-warnings"],
            # first run is always in playback mode
            is_live_and_recording=False,
            log_file_path=log_file_path,
            log_in_json=True,
            log_suffix=".first",
        )

    failed_tests, tests_failed_with_recording_mismatch = get_failed_tests(json_log_file_path)
    failed_tests_path = log_file_path.with_suffix(log_file_path.suffix + ".failed.txt")
    with open(failed_tests_path, "w", encoding="utf-8") as f:
        f.write("\n".join(failed_tests))

    if record_mismatch and tests_failed_with_recording_mismatch:
        print(f"Redo live mode recording for tests: {failed_tests_path}")
        run_simple(
            tests_failed_with_recording_mismatch,
            working_dir,
            extra_params=["--tb=line"],
            is_live_and_recording=True,
            log_suffix=".record",
            log_file_path=log_file_path,
        )

        print(f"Rerun playback mode for failed tests: {failed_tests_path}")
        run_simple(
            failed_tests,
            working_dir,
            extra_params=extras + ["--disable-warnings", "--disable-pytest-warnings"],
            is_live_and_recording=False,
            log_file_path=log_file_path,
            log_suffix=".final",
        )


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--file",
        "-f",
        type=str,
        help="File containing tests to run, each line is a test name",
    )
    parser.add_argument(
        "--name",
        "-n",
        type=str,
        help="Name of the test to run. Usual pytest formats are supported, e.g., 'tests/pipeline_job/' "
        "and 'tests/pipeline_job/e2etests/test_pipeline_job.py::TestPipelineJob'."
        "Test param specification is also supported, e.g., 'tests/pipeline_job/e2etests/"
        "test_pipeline_job.py::TestPipelineJob::test_pipeline_job_with_data_binding_expression[0-input_basic.yml]'",
    )
    parser.add_argument(
        "--record-mismatch",
        "-r",
        action="store_true",
        help="If specified, pytest log will be outputted to tmp/pytest_log.json, "
        "then tests failed with recording not found error will be rerun in live & recording mode."
        "Note that .env file will be updated during the process, so please revert the change manually "
        "if the script run is stopped early.",
    )
    parser.add_argument(
        "--skip-first-run",
        "-s",
        action="store_true",
        help="If specified, will skip the first run in record-mismatch mode. Failed tests will be loaded from "
        "tmp/pytest_first_run.log generated in previous record-mismatch mode first run.",
    )

    _args, _extras = parser.parse_known_args()

    run_tests(
        _args,
        _extras,
    )