#!/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 re
import sys
import json
from datetime import datetime
from pathlib import Path
from typing import Dict, Generator, List, cast

from impl.check_code_hygiene import has_crlf_line_endings
from impl.common import (
    CROSVM_ROOT,
    TOOLS_ROOT,
    argh,
    chdir,
    cmd,
    cwd_context,
    parallel,
    run_main,
)
from impl.health_check import Check, CheckContext, run_checks

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")


def check_python_tests(_: CheckContext):
    "No matter which python files have changed, run all available python tests."
    PYTHON_TESTS = [
        *TOOLS_ROOT.glob("tests/*.py"),
        TOOLS_ROOT / "impl/common.py",
    ]
    parallel(*cmd("python3").foreach(PYTHON_TESTS)).fg(quiet=True)


def check_python_types(context: CheckContext):
    "Run mypy on all python files to type-check."
    parallel(*mypy("--pretty").foreach(context.all_files)).fg(quiet=True)


def check_python_format(context: CheckContext):
    parallel(*black("--check" if not context.fix else None).foreach(context.modified_files)).fg(
        quiet=not context.fix
    )


def check_crlf_line_endings(_: CheckContext):
    "Checks for crlf line endingings."
    crlf_endings = has_crlf_line_endings()
    if crlf_endings:
        print("Error: Following files have crlf(dos) line encodings")
        print(*crlf_endings)
        raise Exception("Files with crlf line endings.")


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."
        )
    parallel(
        *mdformat("--wrap 100", "--check" if not context.fix else "").foreach(
            context.modified_files
        )
    ).fg(quiet=not context.fix)


def check_rust_clippy(_: CheckContext):
    "Runs clippy on the whole project, no matter which rs files were touched."
    cmd("./tools/clippy --locked").with_color_flag().fg(quiet=True)


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


def check_rust_lockfiles(_: CheckContext):
    "Verifies that none of the Cargo.lock files require updates."
    lockfiles = [Path("Cargo.lock"), *Path("common").glob("*/Cargo.lock")]
    for path in lockfiles:
        with cwd_context(path.parent):
            if not cmd("cargo update --workspace --locked").success():
                print(f"{path} is not up-to-date.")
                print()
                print("You may need to rebase your changes and run `cargo update --workspace`")
                print("(or ./tools/run_tests) to ensure the Cargo.lock file is current.")
                raise Exception("Cargo.lock out of date")


# These crosvm features are currently not built upstream. Do not add to this list.
KNOWN_DISABLED_FEATURES = [
    "default-no-sandbox",
    "direct",
    "libvda",
    "whpx",
]


def check_rust_features(_: CheckContext):
    "Verifies that all cargo features are included in the list of features we compile upstream."
    metadata = json.loads(cmd("cargo metadata --format-version=1").stdout())
    crosvm_metadata = next(p for p in metadata["packages"] if p["name"] == "crosvm")
    features = cast(Dict[str, List[str]], crosvm_metadata["features"])

    def collect_features(feature_name: str) -> Generator[str, None, None]:
        yield feature_name
        for feature in features[feature_name]:
            name = feature.split("/")[0]
            if name in features:
                yield from collect_features(name)

    all_platform_features = set(
        (
            *collect_features("all-x86_64"),
            *collect_features("all-aarch64"),
            *collect_features("all-armhf"),
            *collect_features("all-mingw64"),
            *collect_features("all-msvc64"),
        )
    )
    disabled_features = [
        feature
        for feature in features
        if feature not in all_platform_features and feature not in KNOWN_DISABLED_FEATURES
    ]
    if disabled_features:
        raise Exception(
            f"The features {', '.join(disabled_features)} are not enabled in upstream crosvm builds."
        )


LICENSE_HEADER_RE = (
    r".*Copyright (?P<year>20[0-9]{2})(?:-20[0-9]{2})? The ChromiumOS Authors\n"
    r".*Use of this source code is governed by a BSD-style license that can be\n"
    r".*found in the LICENSE file\.\n"
    r"( *\*/\n)?"  # allow the end of a C-style comment before the blank line
    r"\n"
)

NEW_LICENSE_HEADER = [
    f"Copyright {datetime.now().year} The ChromiumOS Authors",
    "Use of this source code is governed by a BSD-style license that can be",
    "found in the LICENSE file.",
]


def new_licence_header(file_suffix: str):
    if file_suffix in (".py", "", ".policy", ".sh"):
        prefix = "#"
    else:
        prefix = "//"
    return "\n".join(f"{prefix} {line}" for line in NEW_LICENSE_HEADER) + "\n\n"


def check_copyright_header(context: CheckContext):
    "Checks copyright header. Can 'fix' them if needed by adding the header."
    license_re = re.compile(LICENSE_HEADER_RE, re.MULTILINE)
    for file in context.modified_files:
        header = file.open("r").read(512)
        license_match = license_re.search(header)
        if license_match:
            continue
        # Generated files do not need a copyright header.
        if "generated by" in header:
            continue
        if context.fix:
            print(f"Adding copyright header: {file}")
            contents = file.read_text()
            file.write_text(new_licence_header(file.suffix) + contents)
        else:
            raise Exception(f"Bad copyright header: {file}")


def check_infra_configs(context: CheckContext):
    "Validate luci configs by sending them to luci-config."
    for file in context.modified_files:
        if context.fix:
            lucicfg("fmt", file).fg()
            lucicfg("generate", file).fg()
        lucicfg("fmt --dry-run", file).fg(quiet=True)
        # TODO: Validate config files. Requires authentication with luci inside docker.


def check_infra_tests(context: 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")
    if context.fix:
        recipes("test train --py3-only").fg()
    recipes("test run --py3-only").fg(quiet=True)


def check_file_ends_with_newline(context: CheckContext):
    "Checks if files end with a newline."
    for file_path in context.modified_files:
        with file_path.open("rb") as file:
            # Skip empty files
            file.seek(0, os.SEEK_END)
            if file.tell() == 0:
                continue
            # Check last byte of the file
            file.seek(-1, os.SEEK_END)
            file_end = file.read(1)
            if file_end.decode("utf-8") != "\n":
                if context.fix:
                    file_path.write_text(file_path.read_text() + "\n")
                else:
                    raise Exception(f"File does not end with a newline {file_path}")


# List of all checks and on which files they should run.
CHECKS: List[Check] = [
    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,
    ),
    Check(
        check_rust_format,
        files=["**.rs"],
        exclude=["system_api/src/bindings/*"],
        can_fix=True,
    ),
    Check(
        check_rust_lockfiles,
        files=["**Cargo.toml"],
    ),
    Check(
        check_rust_features,
        files=["Cargo.toml"],
    ),
    Check(
        check_rust_clippy,
        files=["**.rs", "**Cargo.toml"],
    ),
    Check(
        check_python_tests,
        files=["tools/**.py"],
        python_tools=True,
    ),
    Check(
        check_python_types,
        files=["tools/**.py"],
        exclude=["tools/windows/*"],
        python_tools=True,
    ),
    Check(
        check_python_format,
        files=["**.py"],
        python_tools=True,
        exclude=["infra/recipes.py"],
        can_fix=True,
    ),
    Check(
        check_infra_configs,
        files=["infra/config/**.star"],
        can_fix=True,
    ),
    Check(
        check_infra_tests,
        files=["infra/**.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_file_ends_with_newline,
        exclude=[
            "**.h264",
            "**.vp8",
            "**.vp9",
            "**.ivf",
            "**.bin",
            "**.png",
            "**.min.js",
            "**.drawio",
            "infra/**.json",
        ],
    ),
    Check(check_crlf_line_endings),
]

CHECKS_DICT = dict((c.name, c) for c in CHECKS)


@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("--all", default=False, help="Run on all files instead of just modified files.")
@argh.arg(
    "checks",
    choices=[*CHECKS_DICT.keys(), []],
    help="Optional list of checks to run. Defaults to run all checks.",
)
def main(
    list_checks: bool = False,
    fix: bool = False,
    all: bool = False,
    nightly: bool = False,
    *checks: str,
):
    """
    Run health checks on crosvm. This includes formatting, linters and other various checks.
    """
    chdir(CROSVM_ROOT)

    if not checks:
        checks_list = [*CHECKS_DICT.values()]
    else:
        checks_list = [CHECKS_DICT[check] for check in checks]

    if list_checks:
        for check in checks_list:
            print(check.name)
        return
    success = run_checks(checks_list, fix=fix, run_on_all_files=all, nightly=nightly)

    sys.exit(0 if success else 1)


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