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

# pylint: disable=missing-function-docstring

"""Chrome OS kernel configuration tool

Script to merge all configs and run 'make oldconfig' on it to wade out bad
juju, then split the configs into distro-commmon and flavour-specific parts

See this page for more details:
http://dev.chromium.org/chromium-os/how-tos-and-troubleshooting/kernel-configuration
"""

import argparse
import difflib
import fnmatch
import glob
import logging
import os
import re
import shutil
import subprocess
import sys
import tempfile
from typing import Dict, List


MODES = ["oldconfig", "olddefconfig", "editconfig", "genconfig", "checkconfig"]
TOOLCHAIN_PREFIXES = {
    "x86_64": "x86_64-cros-linux-gnu",
    "armel": "armv7a-cros-linux-gnueabihf",
    "arm64": "aarch64-cros-linux-gnu",
}


def splitconfig(srcdir: str) -> Dict[str, str]:
    """Split and merge common configs

    Return dict of filename-to-contents.
    """
    allconfigs = {}

    # Parse config files.
    for config in os.listdir(srcdir):
        # Only config.*
        if not config.endswith(".config"):
            continue

        allconfigs[config] = set()

        for line in open(os.path.join(srcdir, config), encoding="utf-8"):
            m = re.match(r"#*\s*CONFIG_(\w+)[\s=](.*)$", line)
            if not m:
                continue
            option, value = m.groups()
            allconfigs[config].add((option, value))

    # Split out common config options.
    common = None
    for config in allconfigs:
        if common is None:
            common = allconfigs[config].copy()
        else:
            common &= allconfigs[config]
    for config in allconfigs:
        allconfigs[config] -= common
    allconfigs["common.config"] = common

    ret = {}
    # Generate new splitconfigs.
    for config in allconfigs:
        if allconfigs[config] is None:
            continue
        contents = "#\n# Config options generated by splitconfig\n#\n"
        for option, value in sorted(list(allconfigs[config])):
            if value == "is not set":
                contents += "# CONFIG_%s %s\n" % (option, value)
            else:
                contents += "CONFIG_%s=%s\n" % (option, value)
        ret[config] = contents
    return ret


def die(msg: str = "Aborting"):
    logging.error(msg)
    sys.exit(1)


def in_chroot() -> bool:
    return os.path.exists("/etc/cros_chroot_version")


def rerun_in_chroot():
    cmd = [
        "cros_sdk",
        "--working-dir",
        ".",
        "--",
        os.path.join(os.path.curdir, os.path.relpath(sys.argv[0])),
    ] + sys.argv[1:]
    logging.info("Re-running in chroot")
    logging.debug("Run: %s", cmd)
    sys.exit(subprocess.run(cmd, check=False).returncode)


def make_toolchain_args(arch: str) -> Dict[str, str]:
    ret = {
        "LD": "ld.lld",
    }

    if in_chroot():
        prefix = TOOLCHAIN_PREFIXES[arch]
        ccompiler = prefix + "-clang"
        # CrOS build tooling wants us to use these, and avoid a default of
        # unprefixed 'gcc' or 'clang'. Supply defaults here, but also consider
        # environment varibles below.
        ret["HOSTCC"] = "x86_64-pc-linux-gnu-clang"
        ret["HOSTCXX"] = "x86_64-pc-linux-gnu-clang++"
    else:
        prefix = "x86_64-linux-gnu"
        ccompiler = prefix + "-gcc"

    if not shutil.which(ccompiler):
        if in_chroot():
            die(
                '%s not found. Try running "`sudo cros_setup_toolchains -t %s`"'
                % (ccompiler, prefix)
            )
        else:
            die("%s not found. Try running inside chroot." % ccompiler)

    # Propagate a few environment variables as make args.
    for key in ("HOSTCC", "HOSTCXX"):
        val = os.getenv(key)
        if val:
            ret[key] = val

    ret.update(
        {
            "CC": ccompiler,
            "CXX": prefix + "-g++",
            "CROSS_COMPILE": prefix + "-",
        }
    )
    return ret


def getchar() -> str:
    # We don't want line-buffering, and shelling out to bash is much simpler
    # than messing with tty settings.
    return (
        subprocess.check_output(["bash", "-c", 'read -n 1 -s; echo "${REPLY}"'])
        .decode("utf-8")
        .strip()
    )


def editconfig_prompt(variant: str) -> bool:
    print("* %s: press <Enter> to edit, S to skip" % variant)
    return getchar() not in ("S", "s")


def filter_match(target: str, pattern: str) -> bool:
    return fnmatch.fnmatch(target, "*" + pattern + "*")


def is_editconfig_interactive(
    arch_flavour: str, prompt: bool, pattern: str
) -> bool:
    if filter_match(arch_flavour, pattern):
        if not prompt:
            return True
        return editconfig_prompt(arch_flavour)
    return False


def make_cmd(builddir: str, arch: str, target: str) -> List[str]:
    pairs = make_toolchain_args(arch)
    pairs["ARCH"] = "arm" if arch == "armel" else arch
    pairs["O"] = builddir

    make_args = [k + "=" + v for k, v in pairs.items()]

    return ["make", "-j"] + make_args + [target, "savedefconfig"]


def checkconfig(srcdir: str, outdir: str) -> int:
    error = 0
    for src_config in glob.glob(
        srcdir + "/chromeos/config/**/*.config", recursive=True
    ):
        relative = os.path.relpath(src_config, srcdir)
        out_config = os.path.join(outdir, relative)
        logging.debug("Compare %s to %s", src_config, out_config)
        before = open(src_config, "r", encoding="utf-8").readlines()
        after = open(out_config, "r", encoding="utf-8").readlines()

        if before == after:
            continue

        error = 1
        diff = difflib.context_diff(
            before, after, fromfile="a/" + relative, tofile="b/" + relative
        )
        sys.stdout.writelines(diff)
        logging.error("checkconfig failed for config %s", relative)

    if error:
        logging.error(
            """Consider running `%s olddefconfig` to normalize.

NOTE: chromeos/config changes should always be in a CL by themselves and never
squashed into the same patch as code changes. If code and config changes need
to land together, consider using Cq-Depend to make a circular dependency.""",
            sys.argv[0],
        )
    else:
        logging.info("All good!")
    return error


def build_one_arch(
    args, tmpdir: str, srcdir: str, family: str, arch: str
) -> List[subprocess.Popen]:
    procs = []
    config_base_dir = srcdir + "/chromeos/config/" + family
    config_arch_dir = config_base_dir + "/" + arch
    for flavourconfig in glob.glob(
        os.path.join(srcdir, config_arch_dir, "*.flavour.config")
    ):
        flavour = os.path.basename(flavourconfig)
        builddir = os.path.join(tmpdir, "build", family, arch, flavour)
        os.makedirs(builddir)

        # Merge base.config, common.config, <flavour>.flavour.config
        conf = "\n".join(
            open(os.path.join(srcdir, f), encoding="utf-8").read()
            for f in [
                config_base_dir + "/base.config",
                config_arch_dir + "/common.config",
                config_arch_dir + "/" + flavour,
            ]
        )
        open(os.path.join(builddir, ".config"), "w", encoding="utf-8").write(
            conf
        )

        interactive = False
        if args.mode == "genconfig":
            mode = "olddefconfig"
        elif args.mode == "editconfig":
            interactive = is_editconfig_interactive(
                family + "/" + arch + "/" + flavour, not args.yes, args.filter
            )
            mode = "menuconfig" if interactive else "olddefconfig"
        elif args.mode == "checkconfig":
            mode = "olddefconfig"
        else:
            if args.mode == "oldconfig":
                # 'oldconfig' may run into interactive prompts.
                interactive = True
            mode = args.mode

        cmd = make_cmd(builddir, arch, mode)
        if interactive:
            logging.info("Starting interactive cmd: %s", " ".join(cmd))
            ret = subprocess.run(cmd, check=False, cwd=srcdir)
            if ret.returncode:
                die('cmd "%s" failed' % " ".join(cmd))
        else:
            logging.info("Starting background cmd: %s", " ".join(cmd))
            procs += [
                subprocess.Popen(
                    cmd,
                    cwd=srcdir,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                )
            ]
    return procs


def do_splitconfig(
    tmpdir: str,
    srcdir: str,
    outdir: str,
    family: str,
    arches: List[str],
    save_configs: bool,
):
    family_cfg_dir = "chromeos/config/" + family

    if save_configs:
        keep_dir = "CONFIGS/" + family
        os.makedirs(keep_dir, exist_ok=True)

    for arch in arches:
        dest = os.path.join(tmpdir, family, arch)
        family_arch_dir = family_cfg_dir + "/" + arch
        os.makedirs(dest)
        for flavourconfig in glob.glob(
            os.path.join(srcdir, family_arch_dir, "*.flavour.config")
        ):
            flavour = os.path.basename(flavourconfig)
            builddir = os.path.join(tmpdir, "build", family, arch, flavour)
            shutil.copy(
                os.path.join(builddir, "defconfig"), os.path.join(dest, flavour)
            )
            if save_configs:
                shutil.copy(
                    os.path.join(builddir, ".config"),
                    os.path.join(keep_dir, arch + "-" + flavour),
                )
                shutil.copy(
                    os.path.join(builddir, "defconfig"),
                    os.path.join(keep_dir, arch + "-" + flavour + ".def"),
                )
        os.makedirs(os.path.join(outdir, family_arch_dir), exist_ok=True)
        # Find per-arch common items; flavour-unique options can be written out
        # immediately.
        for config, contents in splitconfig(dest).items():
            if config == "common.config":
                open(
                    os.path.join(tmpdir, family, arch + ".config"),
                    "w",
                    encoding="utf-8",
                ).write(contents)
            else:
                open(
                    os.path.join(outdir, family_arch_dir, config),
                    "w",
                    encoding="utf-8",
                ).write(contents)

    # Find cross-arch common items; common ones go to base.config, and unique
    # ones to <arch>/common.config.
    for config, contents in splitconfig(tmpdir + "/" + family).items():
        if config == "common.config":
            open(
                os.path.join(outdir, family_cfg_dir, "base.config"),
                "w",
                encoding="utf-8",
            ).write(contents)
        else:
            assert config.endswith(".config")
            # Python3.9: arch = config.removesuffix('.config')
            arch = config[: -len(".config")]
            open(
                os.path.join(outdir, family_cfg_dir, arch, "common.config"),
                "w",
                encoding="utf-8",
            ).write(contents)


# Arches relevant to this family.
def family_arches(srcdir: str, family: str) -> List[str]:
    return [
        os.path.basename(p)
        for p in glob.glob(srcdir + "/chromeos/config/" + family + "/*")
        if os.path.isdir(p)
    ]


def doit(args, tmpdir: str) -> int:
    for d in (
        os.getcwd(),
        os.path.abspath(os.path.join(os.path.dirname(sys.argv[0]), "..", "..")),
    ):
        if os.path.exists(os.path.join(d, "MAINTAINERS")) and os.path.exists(
            os.path.join(d, "Makefile")
        ):
            srcdir = d
            logging.info("Using top kernel dir: %s", srcdir)
            break
    else:
        die("This does not appear to be a kernel source directory.")

    if not in_chroot() and not args.force:
        rerun_in_chroot()
        assert False  # Should not reach this.

    # Most commands write back to source directory.
    if args.mode == "checkconfig":
        outdir = tmpdir
    else:
        outdir = srcdir

    save_configs = args.mode == "genconfig"

    procs = []

    families = [
        os.path.basename(p)
        for p in glob.glob(srcdir + "/chromeos/config/*")
        if os.path.isdir(p)
    ]
    for family in families:
        for arch in family_arches(srcdir, family):
            if not arch in TOOLCHAIN_PREFIXES:
                die("Unexpected arch: %s" % arch)
            procs += build_one_arch(args, tmpdir, srcdir, family, arch)

    for p in procs:
        ret = p.wait()
        if ret:
            assert p.stderr  # Remind pytype we captured stderr.
            logging.error(p.stderr.read().decode("utf-8"))
            die('cmd "%s" failed' % " ".join(p.args))

    logging.info("Generating splitconfigs")
    for family in families:
        do_splitconfig(
            tmpdir,
            srcdir,
            outdir,
            family,
            family_arches(srcdir, family),
            save_configs,
        )

    if args.mode == "checkconfig":
        return checkconfig(srcdir, outdir)

    return 0


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Chrome OS kernel configuration script",
        epilog="""
               Note that Kbuild will evaluate some features depending on the
               toolchain, so we try to enter the SDK chroot. This can be
               overridden, with a potentially degraded experience.
               """,
    )

    log_level_choices = ["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"]
    parser.add_argument(
        "--log_level", "-l", choices=log_level_choices, default="INFO"
    )
    parser.add_argument(
        "--force",
        "-F",
        action="store_true",
        help="Force allowing to run outside the chroot",
    )
    parser.add_argument(
        "--filter",
        "-f",
        default="",
        help="Only attempt to edit configs which match filter",
    )
    parser.add_argument(
        "--yes",
        "-y",
        action="store_true",
        help="Edit all configs which match unconditionally",
    )
    parser.add_argument("mode", choices=MODES, help="sub-command/mode")

    args = parser.parse_args()
    logging.basicConfig(
        level=args.log_level, format="%(levelname)s - %(message)s"
    )

    with tempfile.TemporaryDirectory() as tmpdir:
        return doit(args, tmpdir)


sys.exit(main())