#!/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)):
            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.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').readlines()
        after = open(out_config, 'r').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)).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').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').write(contents)
            else:
                open(os.path.join(outdir, family_arch_dir, config),
                     'w').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').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').write(contents)


# Arches relevant to this family.
def family_arches(family: str) -> List[str]:
    return [os.path.basename(p)
            for p in glob.glob('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('chromeos/config/*')
                if os.path.isdir(p)]
    for family in families:
        for arch in family_arches(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(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())