#!/usr/bin/env python3
#
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License").
# You may not use this file except in compliance with the License.
# A copy of the License is located at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# or in the "license" file accompanying this file. This file is distributed
# on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
# express or implied. See the License for the specific language governing
# permissions and limitations under the License.


import argparse
import importlib
import logging
import os
import pathlib
import re
import subprocess
import sys
import uuid

DESCRIPTION = "Execute e2e and unit tests for Litani"


def get_args():
    pars = argparse.ArgumentParser(description=DESCRIPTION)
    for arg in [{
            "flags": ["--output-dir"],
            "help": "output dir for test results",
            "default": pathlib.Path(__file__).resolve().parent / "output",
            "type": pathlib.Path
        }, {
            "flags": ["--fast"],
            "help": "run fast tests only",
            "action": "store_true"
        }]:
        flags = arg.pop("flags")
        pars.add_argument(*flags, **arg)
    return pars.parse_args()


def run_cmd(cmd):
    try:
        subprocess.run([str(c) for c in cmd], check=True)
    except subprocess.CalledProcessError:
        logging.error(
            "Failed to run command '%s'", " ".join([str(c) for c in cmd]))
        sys.exit(1)


def is_slow_test(module_file):
    try:
        return importlib.import_module(str(module_file.stem)).SLOW
    except AttributeError:
        logging.error("Variable SLOW is missing from: %s", module_file.name)
        sys.exit(1)


def litani_add(litani, counter, *args, **kwargs):
    cmd = [litani, "add-job"]
    for arg in args:
        switch = re.sub("_", "-", arg)
        cmd.append(f"--{switch}")
    for arg, value in kwargs.items():
        switch = re.sub("_", "-", arg)
        cmd.append(f"--{switch}")
        if isinstance(value, list):
            cmd.extend(value)
        else:
            cmd.append(value)
    run_cmd(cmd)
    counter["added"] += 1
    print_counter(counter)


def collapse(string):
    return re.sub(r"\s+", " ", string)


def add_e2e_tests(litani, test_dir, counter, output_dir, fast):
    e2e_test_dir = test_dir / "e2e"
    # 4 jobs per test (init, add-jobs, run-build, check-run)
    # skip __init__.py and __pycache__
    counter["total"] += (len(os.listdir(e2e_test_dir / "tests")) - 2) * 4
    sys.path.insert(1, str(e2e_test_dir / "tests"))
    for test_file in (e2e_test_dir / "tests").iterdir():
        if test_file.name in ["__init__.py", "__pycache__"]:
            continue

        add_transform_jobs = False
        with open(test_file) as handle:
            for line in handle:
                if line.strip().startswith("def transform_jobs("):
                    add_transform_jobs = True
                    break
        if fast and is_slow_test(test_file):
            continue

        run_dir = output_dir / "e2e_outputs" / str(uuid.uuid4())

        timeout=10 if fast else 0

        litani_add(
            litani, counter,
            command=collapse(f"""
                {e2e_test_dir / 'run'}
                --test-file {test_file}
                --litani {litani}
                --run-dir {run_dir}
                --operation init"""),
            pipeline=f"End-to-end: {test_file.stem}",
            ci_stage="test",
            description=f"{test_file.stem}: init",
            outputs=run_dir / ".litani_cache_dir",
            cwd=run_dir,
            timeout=timeout)

        run_build_input = str(uuid.uuid4())
        litani_add(
            litani, counter,
            command=collapse(f"""
                {e2e_test_dir / 'run'}
                --test-file {test_file}
                --litani {litani}
                --run-dir {run_dir}
                --operation add-jobs"""),
            pipeline=f"End-to-end: {test_file.stem}",
            ci_stage="test",
            description=f"{test_file.stem}: add jobs",
            inputs=run_dir / ".litani_cache_dir",
            phony_outputs=run_build_input,
            outputs=f"{run_dir}/output/jobs",
            cwd=run_dir,
            timeout=timeout)

        if add_transform_jobs:
            add_jobs_output = run_build_input
            run_build_input = str(uuid.uuid4())
            litani_add(
                litani, counter,
                command=collapse(f"""
                    {e2e_test_dir / 'run'}
                    --test-file {test_file}
                    --litani {litani}
                    --run-dir {run_dir}
                    --operation transform-jobs"""),
                pipeline=f"End-to-end: {test_file.stem}",
                ci_stage="test",
                description=f"{test_file.stem}: transform jobs",
                inputs=add_jobs_output,
                phony_outputs=run_build_input,
                cwd=run_dir)

        litani_add(
            litani, counter,
            command=collapse(f"""
                {e2e_test_dir / 'run'}
                --test-file {test_file}
                --litani {litani}
                --run-dir {run_dir}
                --operation run-build"""),
            pipeline=f"End-to-end: {test_file.stem}",
            ci_stage="test",
            description=f"{test_file.stem}: run build",
            inputs=run_build_input,
            outputs=f"{run_dir}/output/run.json",
            cwd=run_dir,
            timeout=timeout)

        litani_add(
            litani, counter,
            command=collapse(f"""
                {e2e_test_dir / 'run'}
                --test-file {test_file}
                --litani {litani}
                --run-dir {run_dir}
                --operation check-run"""),
            pipeline=f"End-to-end: {test_file.stem}",
            ci_stage="report",
            description=f"{test_file.stem}: check run",
            inputs=f"{run_dir}/output/run.json",
            cwd=run_dir,
            timeout=timeout)


def add_unit_tests(litani, test_dir, root_dir, counter):
    for fyle in (test_dir / "unit").iterdir():
        if fyle.name in ["__init__.py", "__pycache__"]:
            continue
        litani_add(
            litani, counter,
            command=f"python3 -m unittest test.unit.{fyle.stem}",
            pipeline="Unit tests",
            ci_stage="test",
            description=fyle.stem,
            cwd=root_dir)
        counter["total"] += 1


def print_counter(counter):
    print("\r{added} / {total} tests added".format(**counter), end="")


def main():
    args = get_args()
    logging.basicConfig(format="\nrun-tests: %(message)s")
    test_dir = pathlib.Path(__file__).resolve().parent
    root = test_dir.parent
    litani = root / "litani"

    output_dir = args.output_dir.resolve()
    output_dir.mkdir(exist_ok=True, parents=True)
    os.chdir(output_dir)

    run_cmd([
        litani, "init",
        "--project", "Litani Test Suite",
        "--output-prefix", ".",
        "--output-symlink", "latest"])

    counter = {
        "added": 0,
        "total": 0,
    }

    add_unit_tests(litani, test_dir, root, counter)
    add_e2e_tests(
        litani, test_dir, counter, output_dir, args.fast)
    print()

    run_cmd([litani, "run-build"])


if __name__ == "__main__":
    main()
