#!/usr/bin/env python3
# Copyright 2022 The ChromiumOS Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import os
import typing
from typing import Generator, List, Literal, Optional, Tuple

from impl.common import (
    CROSVM_ROOT,
    TOOLS_ROOT,
    Triple,
    argh,
    chdir,
    cmd,
    is_kiwi_repo,
    run_main,
)
from impl.presubmit import Check, CheckContext, run_checks, Group

python = cmd("python3")
mypy = cmd("mypy").with_color_env("MYPY_FORCE_COLOR")
black = cmd("black").with_color_arg(always="--color", never="--no-color")
mdformat = cmd("mdformat")
lucicfg = cmd("third_party/depot_tools/lucicfg")

# All supported platforms as a type and a list.
Platform = Literal["x86_64", "aarch64", "mingw64", "armhf", "riscv64"]
PLATFORMS: Tuple[Platform, ...] = typing.get_args(Platform)


def platform_is_supported(platform: Platform):
    "Returns true if the platform is available as a target in rustup."
    triple = Triple.from_shorthand(platform)
    installed_toolchains = cmd("rustup target list --installed").lines()
    return str(triple) in installed_toolchains


####################################################################################################
# Check methods
#
# Each check returns a Command (or list of Commands) to be run to execute the check. They are
# registered and configured in the CHECKS list below.
#
# Some check functions are factory functions that return a check command for all supported
# platforms.


def check_python_tests(_: CheckContext):
    "Runs unit tests for python dev tooling."
    PYTHON_TESTS = [
        "tests.cl_tests",
        "impl.common",
    ]
    return [python.with_cwd(TOOLS_ROOT).with_args("-m", file) for file in PYTHON_TESTS]


def check_python_types(context: CheckContext):
    "Run mypy type checks on python dev tooling."
    return [mypy("--pretty", file) for file in context.all_files]


def check_python_format(context: CheckContext):
    "Runs the black formatter on python dev tooling."
    return black.with_args(
        "--check" if not context.fix else None,
        *context.modified_files,
    )


def check_markdown_format(context: CheckContext):
    "Runs mdformat on all markdown files."
    if "blaze" in mdformat("--version").stdout():
        raise Exception(
            "You are using google's mdformat. "
            + "Please update your PATH to ensure the pip installed mdformat is available."
        )
    return mdformat.with_args(
        "--wrap 100",
        "--check" if not context.fix else "",
        *context.modified_files,
    )


def check_rust_format(context: CheckContext):
    "Runs rustfmt on all modified files."
    if context.nightly_fmt:
        rustfmt = cmd(
            cmd("rustup +nightly which rustfmt"),
            "--config imports_granularity=item,group_imports=StdExternalCrate",
        )
    else:
        rustfmt = cmd(cmd("rustup which rustfmt"))
    return rustfmt.with_color_flag().with_args(
        "--check" if not context.fix else "",
        *context.modified_files,
    )


def check_cargo_doc(_: CheckContext):
    "Runs cargo-doc and verifies that no warnings are emitted."
    return cmd("./tools/cargo-doc").with_env("RUSTDOCFLAGS", "-D warnings").with_color_flag()


def check_doc_tests(_: CheckContext):
    "Runs cargo doc tests. These cannot be run via nextest and run_tests."
    return cmd(
        "cargo test",
        "--doc",
        "--workspace",
        "--features=all-x86_64",
    ).with_color_flag()


def check_mdbook(_: CheckContext):
    "Runs cargo-doc and verifies that no warnings are emitted."
    return cmd("mdbook build docs/book/")


def check_crosvm_tests(platform: Platform):
    def check(_: CheckContext):
        if not platform_is_supported(platform):
            return None
        dut = None
        if os.access("/dev/kvm", os.W_OK):
            if platform == "x86_64":
                dut = "--dut=vm"
            elif platform == "aarch64":
                dut = "--dut=vm"
        return cmd("./tools/run_tests --verbose --platform", platform, dut).with_color_flag()

    check.__doc__ = f"Runs all crosvm tests for {platform}."

    return check


def check_crosvm_unit_tests(platform: Platform):
    def check(_: CheckContext):
        if not platform_is_supported(platform):
            return None
        command = cmd("./tools/run_tests --verbose --platform", platform).with_color_flag()
        if platform == "riscv64":
            command = command.with_args("--no-default-features")
        return command

    check.__doc__ = f"Runs crosvm unit tests for {platform}."

    return check


def check_crosvm_build(platform: Platform, features: Optional[str]):
    def check(_: CheckContext):
        return cmd(
            "./tools/run_tests --no-run --verbose --platform",
            platform,
            f"--features={features}" if features is not None else None,
        ).with_color_flag()

    check.__doc__ = f"Builds crosvm for {platform} with features {features}."

    return check


def check_clippy(platform: Platform):
    def check(context: CheckContext):
        if not platform_is_supported(platform):
            return None
        return cmd(
            "./tools/clippy --platform",
            platform,
            "--fix" if context.fix else None,
        ).with_color_flag()

    check.__doc__ = f"Runs clippy for {platform}."

    return check


def check_infra_configs(context: CheckContext):
    "Validate luci configs by sending them to luci-config."
    # TODO: Validate config files. Requires authentication with luci inside docker.
    return [lucicfg("fmt --dry-run", file) for file in context.modified_files]


def check_infra_tests(_: CheckContext):
    "Run recipe.py tests, all of them, regardless of which files were modified."
    recipes = cmd("infra/recipes.py").with_path_env("third_party/depot_tools")
    return recipes("test run")


def custom_check(name: str, can_fix: bool = False):
    "Custom checks are written in python in tools/custom_checks. This is a wrapper to call them."

    def check(context: CheckContext):
        return cmd(
            TOOLS_ROOT / "custom_checks",
            name,
            *context.modified_files,
            "--fix" if can_fix and context.fix else None,
        )

    check.__name__ = name.replace("-", "_")
    check.__doc__ = f"Runs tools/custom_check {name}"
    return check


####################################################################################################
# Checks configuration
#
# Configures which checks are available and on which files they are run.
# Check names default to the function name minus the check_ prefix

CHECKS: List[Check] = [
    Check(
        check_rust_format,
        files=["**.rs"],
        exclude=["system_api/src/bindings/*"],
        can_fix=True,
    ),
    Check(
        check_mdbook,
        files=["docs/**/*"],
    ),
    Check(
        check_cargo_doc,
        files=["**.rs", "**Cargo.toml"],
        priority=True,
    ),
    Check(
        check_doc_tests,
        files=["**.rs", "**Cargo.toml"],
        priority=True,
    ),
    Check(
        check_python_tests,
        files=["tools/**.py"],
        python_tools=True,
        priority=True,
    ),
    Check(
        check_python_types,
        files=["tools/**.py"],
        exclude=["tools/windows/*", "tools/contrib/memstats_chart/*"],
        python_tools=True,
    ),
    Check(
        check_python_format,
        files=["**.py"],
        python_tools=True,
        exclude=["infra/recipes.py"],
        can_fix=True,
    ),
    Check(
        check_markdown_format,
        files=["**.md"],
        exclude=[
            "infra/README.recipes.md",
            "docs/book/src/appendix/memory_layout.md",
        ],
        can_fix=True,
    ),
    *(
        Check(
            check_crosvm_build(platform, features="default"),
            custom_name=f"crosvm_build_default_{platform}",
            files=["**.rs"],
            priority=True,
        )
        for platform in PLATFORMS
    ),
    *(
        Check(
            check_crosvm_tests(platform),
            custom_name=f"crosvm_tests_{platform}",
            files=["**.rs"],
            priority=True,
        )
        for platform in PLATFORMS
    ),
    *(
        Check(
            check_crosvm_unit_tests(platform),
            custom_name=f"crosvm_unit_tests_{platform}",
            files=["**.rs"],
            priority=True,
        )
        for platform in PLATFORMS
    ),
    *(
        Check(
            check_clippy(platform),
            custom_name=f"clippy_{platform}",
            files=["**.rs"],
            can_fix=True,
            priority=True,
        )
        for platform in PLATFORMS
    ),
    Check(
        custom_check("check-copyright-header"),
        files=["**.rs", "**.py", "**.c", "**.h", "**.policy", "**.sh"],
        exclude=[
            "infra/recipes.py",
            "hypervisor/src/whpx/whpx_sys/*.h",
            "third_party/vmm_vhost/*",
            "net_sys/src/lib.rs",
            "system_api/src/bindings/*",
        ],
        python_tools=True,
        can_fix=True,
    ),
    Check(
        custom_check("check-rust-features"),
        files=["**Cargo.toml"],
    ),
    Check(
        custom_check("check-rust-lockfiles"),
        files=["**Cargo.toml"],
    ),
    Check(
        custom_check("check-line-endings"),
    ),
    Check(
        custom_check("check-file-ends-with-newline"),
        exclude=[
            "**.h264",
            "**.vp8",
            "**.vp9",
            "**.ivf",
            "**.bin",
            "**.png",
            "**.min.js",
            "**.drawio",
            "**.json",
        ],
    ),
]

# We disable LUCI infra related tests because kokoro doesn't have internet connectivity that
# the tests rely on.
if not is_kiwi_repo():
    CHECKS.extend(
        [
            Check(
                check_infra_configs,
                files=["infra/config/**.star"],
                can_fix=True,
            ),
            Check(
                check_infra_tests,
                files=["infra/**.py"],
                can_fix=True,
            ),
        ]
    )

####################################################################################################
# Group configuration
#
# Configures pre-defined groups of checks. Some are configured for CI builders and others
# are configured for convenience during local development.

GROUPS: List[Group] = [
    # The default group is run if no check or group is explicitly set
    Group(
        name="default",
        doc="Checks run by default",
        checks=[
            "default_health_checks",
            # Run only one task per platform to prevent blocking on the build cache.
            "crosvm_tests_x86_64",
            "crosvm_unit_tests_aarch64",
            "crosvm_unit_tests_mingw64",
            "clippy_armhf",
        ],
    ),
    Group(
        name="quick",
        doc="Runs a quick subset of presubmit checks.",
        checks=[
            "default_health_checks",
            "crosvm_unit_tests_x86_64",
            "clippy_aarch64",
        ],
    ),
    Group(
        name="all",
        doc="Run checks of all builders.",
        checks=[
            "health_checks",
            *(f"linux_{platform}" for platform in PLATFORMS),
        ],
    ),
    # Convenience groups for local usage:
    Group(
        name="clippy",
        doc="Runs clippy for all platforms",
        checks=[f"clippy_{platform}" for platform in PLATFORMS],
    ),
    Group(
        name="unit_tests",
        doc="Runs unit tests for all platforms",
        checks=[f"crosvm_unit_tests_{platform}" for platform in PLATFORMS],
    ),
    Group(
        name="format",
        doc="Runs all formatting checks (or fixes)",
        checks=[
            "rust_format",
            "markdown_format",
            "python_format",
        ],
    ),
    Group(
        name="default_health_checks",
        doc="Health checks to run by default",
        checks=[
            # Check if lockfiles need updating first. Otherwise another step may do the update.
            "rust_lockfiles",
            "copyright_header",
            "file_ends_with_newline",
            "line_endings",
            "markdown_format",
            "mdbook",
            "cargo_doc",
            "python_format",
            "python_types",
            "rust_features",
            "rust_format",
        ],
    ),
    # The groups below are used by builders in CI:
    Group(
        name="health_checks",
        doc="Checks run on the health_check builder",
        checks=[
            "default_health_checks",
            "doc_tests",
            "python_tests",
        ]
        + (["infra_configs", "infra_tests"] if not is_kiwi_repo() else []),
    ),
    *(
        Group(
            name=f"linux_{platform}",
            doc=f"Checks run on the linux-{platform} builder",
            checks=[
                f"crosvm_tests_{platform}",
                f"clippy_{platform}",
                f"crosvm_build_default_{platform}",
            ],
        )
        for platform in PLATFORMS
    ),
]

# Turn both lists into dicts for convenience
CHECKS_DICT = dict((c.name, c) for c in CHECKS)
GROUPS_DICT = dict((c.name, c) for c in GROUPS)


def validate_config():
    "Validates the CHECKS and GROUPS configuration."
    for group in GROUPS:
        for check in group.checks:
            if check not in CHECKS_DICT and check not in GROUPS_DICT:
                raise Exception(f"Group {group.name} includes non-existing item {check}.")

    def find_in_group(check: Check):
        for group in GROUPS:
            if check.name in group.checks:
                return True
        return False

    for check in CHECKS:
        if not find_in_group(check):
            raise Exception(f"Check {check.name} is not included in any group.")

    all_names = [c.name for c in CHECKS] + [g.name for g in GROUPS]
    for name in all_names:
        if all_names.count(name) > 1:
            raise Exception(f"Check or group {name} is defined multiple times.")


def get_check_names_in_group(group: Group) -> Generator[str, None, None]:
    for name in group.checks:
        if name in GROUPS_DICT:
            yield from get_check_names_in_group(GROUPS_DICT[name])
        else:
            yield name


@argh.arg("--list-checks", default=False, help="List names of available checks and exit.")
@argh.arg("--fix", default=False, help="Asks checks to fix problems where possible.")
@argh.arg("--no-delta", default=False, help="Run on all files instead of just modified files.")
@argh.arg("--no-parallel", default=False, help="Do not run checks in parallel.")
@argh.arg("--nightly-fmt", default=False, help="Use nightly rust for rustfmt")
@argh.arg(
    "checks_or_groups",
    help="List of checks or groups to run. Defaults to run the `default` group.",
)
def main(
    list_checks: bool = False,
    fix: bool = False,
    no_delta: bool = False,
    no_parallel: bool = False,
    nightly_fmt: bool = False,
    *checks_or_groups: str,
):
    chdir(CROSVM_ROOT)
    validate_config()

    if not checks_or_groups:
        checks_or_groups = ("default",)

    # Resolve and validate the groups and checks provided
    check_names: List[str] = []
    for check_or_group in checks_or_groups:
        if check_or_group in CHECKS_DICT:
            check_names.append(check_or_group)
        elif check_or_group in GROUPS_DICT:
            check_names += list(get_check_names_in_group(GROUPS_DICT[check_or_group]))
        else:
            raise Exception(f"No such check or group: {check_or_group}")

    # Remove duplicates while preserving order
    check_names = list(dict.fromkeys(check_names))

    if list_checks:
        for check in check_names:
            print(check)
        return

    check_list = [CHECKS_DICT[name] for name in check_names]

    run_checks(
        check_list,
        fix=fix,
        run_on_all_files=no_delta,
        nightly_fmt=nightly_fmt,
        parallel=not no_parallel,
    )


def usage():
    groups = "\n".join(f"  {group.name}: {group.doc}" for group in GROUPS)
    checks = "\n".join(f"  {check.name}: {check.doc}" for check in CHECKS)
    return f"""\
Runs checks on the crosvm codebase.

Basic usage, to run a default selection of checks:

    ./tools/presubmit

Some checkers can fix issues they find (e.g. formatters, clippy, etc):

    ./tools/presubmit --fix


Various groups of presubmit checks can be run via:

    ./tools/presubmit group_name

Available groups are:
{groups}

You can also provide the names of specific checks to run:

    ./tools/presubmit check1 check2

Available checks are:
{checks}
"""


if __name__ == "__main__":
    run_main(main, usage=usage())