#!/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, Union

from impl.common import (
    CROSVM_ROOT,
    TOOLS_ROOT,
    Triple,
    argh,
    chdir,
    cmd,
    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)

ClippyOnlyPlatform = Literal["android"]
CLIPPY_ONLY_PLATFORMS: Tuple[ClippyOnlyPlatform, ...] = typing.get_args(ClippyOnlyPlatform)


def platform_is_supported(platform: Union[Platform, ClippyOnlyPlatform]):
    "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 = [
        # Disabled due to b/309148074
        # "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."
    rustfmt = cmd(cmd("rustup +nightly which rustfmt").stdout())
    # Windows doesn't accept very long arguments: https://learn.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-createprocessa#:~:text=The%20maximum%20length%20of%20this%20string%20is%2032%2C767%20characters%2C%20including%20the%20Unicode%20terminating%20null%20character.%20If%20lpApplicationName%20is%20NULL%2C%20the%20module%20name%20portion%20of%20lpCommandLine%20is%20limited%20to%20MAX_PATH%20characters.
    return list(
        rustfmt.with_color_flag()
        .with_args("--check" if not context.fix else "")
        .foreach(context.modified_files, batch_size=10)
    )


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
        return cmd("./tools/run_tests --verbose --platform", platform).with_color_flag()

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

    return check


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

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

    return check


def check_clippy(platform: Union[Platform, ClippyOnlyPlatform]):
    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 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/*",
            "tools/contrib/cros_tracing_analyser/*",
        ],
        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_build(platform, features="", no_default_features=True),
            custom_name=f"crosvm_build_no_default_{platform}",
            files=["**.rs"],
            priority=True,
        )
        # TODO: b/260607247 crosvm does not compile with no-default-features on mingw64
        for platform in PLATFORMS
        if platform != "mingw64"
    ),
    *(
        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, *CLIPPY_ONLY_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",
            "**.dtb",
            "**.dtbo",
        ],
    ),
]

####################################################################################################
# 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),
            *(f"clippy_{platform}" for platform in CLIPPY_ONLY_PLATFORMS),
        ],
    ),
    # Convenience groups for local usage:
    Group(
        name="clippy",
        doc="Runs clippy for all platforms",
        checks=[f"clippy_{platform}" for platform in PLATFORMS + CLIPPY_ONLY_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",
        ],
    ),
    Group(
        name="android-aarch64",
        doc="Checks run on the android-aarch64 builder",
        checks=[
            "clippy_android",
        ],
    ),
    *(
        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}",
            ]
            # TODO: b/260607247 crosvm does not compile with no-default-features on mingw64
            + ([f"crosvm_build_no_default_{platform}"] if platform != "mingw64" else []),
        )
        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(
    "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,
    *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,
        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())