#!/usr/bin/env python3
#
# Copyright (C) 2017 The Android Open Source Project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# pylint: disable=not-callable

"""Update the prebuilt clang from the build server."""

import argparse
import contextlib
import glob
import inspect
import logging
import os
from pathlib import Path
import shutil
import subprocess
import sys
import utils

import paths

PGO_PROFILE_PATTERN = 'pgo-*.tar.xz'
BOLT_PROFILE_PATTERN = 'bolt-*.tar.xz'

def logger():
    """Returns the module level logger."""
    return logging.getLogger(__name__)


class ArgParser(argparse.ArgumentParser):
    def __init__(self) -> None:
        super(ArgParser, self).__init__(
            description=inspect.getdoc(sys.modules[__name__]))

        self.add_argument(
            'build', metavar='BUILD',
            help='Build number to pull from the build server.')

        self.add_argument(
            '-b', '--bug', help='Bug to reference in commit message.')

        self.add_argument(
            '-br', '--branch', help='Branch to fetch from (or automatic).')

        self.add_argument(
            '--use-current-branch', action='store_true',
            help='Do not repo start a new branch for the update.')

        self.add_argument(
            '--skip-fetch',
            '-sf',
            action='store_true',
            default=False,
            help='Skip the fetch, and only do the extraction step')

        self.add_argument(
            '--skip-cleanup',
            '-sc',
            action='store_true',
            default=False,
            help='Skip the cleanup, and leave intermediate files')

        self.add_argument(
            '--skip-update-profiles',
            '-sp',
            action='store_true',
            default=False,
            help='Skip updating PGO and BOLT profiles')

        self.add_argument(
            '--overwrite', action='store_true',
            help='Remove/overwrite any existing prebuilt directories.')

        self.add_argument(
            '--no-validity-check', action='store_true',
            help='Skip validity checks on the prebuilt binaries.')

        host_choices = ['darwin-x86', 'linux-x86', 'windows-x86']
        self.add_argument(
            '--host', metavar='HOST_OS',
            choices=host_choices,
            help=f'Update prebuilts only for HOST_OS (one of {host_choices}).')

        self.add_argument(
            '--repo-upload', action='store_true',
            help='Upload prebuilts CLs to gerrit using \'repo upload\'')

        self.add_argument(
            '--hashtag', metavar='HASHTAGS',
            help='Extra hashtags (comma separated) during \'repo upload\'')


def fetch_artifact(branch, target, build, pattern):
    fetch_artifact_path = '/google/data/ro/projects/android/fetch_artifact'
    cmd = [fetch_artifact_path, f'--branch={branch}',
           f'--target={target}', f'--bid={build}', pattern]
    utils.check_call(cmd)


def extract_clang_info(clang_dir):
    version_file_path = os.path.join(clang_dir, 'AndroidVersion.txt')
    with open(version_file_path) as version_file:
        # e.g. for contents: ['7.0.1', 'based on r326829']
        contents = [l.strip() for l in version_file.readlines()]
        full_version = contents[0]
        major_version = full_version.split('.')[0]
        revision = contents[1].split()[-1]
        return full_version, major_version, revision


def symlink_to_linux_resource_dir(install_dir):
    # Assume we're in a Darwin (non-linux) prebuilt dir.  Find the Clang version
    # string.  Pick the longest string, if there's more than one.
    version_dirs = os.listdir(os.path.join(install_dir, 'lib', 'clang'))
    if version_dirs:
        version_dirs.sort(key=len)
    version_dir = version_dirs[-1]

    symlink_dir = os.path.join(install_dir, 'lib', 'clang', version_dir,
                               'lib')
    link_src = os.path.join('/'.join(['..'] * 6), 'linux-x86', symlink_dir,
                            'linux')
    link_dst = 'linux'

    # 'cd' to symlink_dir and create a symlink from link_dst to link_src
    prebuilt_dir = os.getcwd()
    os.chdir(symlink_dir)
    os.symlink(link_src, link_dst)
    os.chdir(prebuilt_dir)


def validity_check(host, install_dir, clang_version_major):
    # Make sure the official toolchain (non llvm-next) is built with PGO
    # profiles.
    if host == 'linux-x86':
      realClangPath = os.path.join(install_dir, 'bin', 'clang-' + clang_version_major)
      strings = utils.check_output([realClangPath, '--version'])
      llvm_next = strings.find('ANDROID_LLVM_NEXT') != -1

      if not llvm_next:
        has_pgo = ('+pgo' in strings) and ('-pgo' not in strings)
        if not has_pgo:
          logger().error('The Clang binary is not built with PGO profiles.')
          return False
        has_bolt = ('+bolt' in strings) and ('-bolt' not in strings)
        if not has_bolt:
          logger().error('The Clang binary is not built with BOLT profiles.')
          return False
        has_lto = ('+lto' in strings) and ('-lto' not in strings)
        if not has_lto:
          logger().error('The Clang binary is not built with LTO.')
          return False
        has_mlgo = ('+mlgo' in strings) and ('-mlgo' not in strings)
        if not has_mlgo:
          logger().error('The Clang binary is not built with MLGO support.')
          return False

    # Check that all the files listed in remote_toolchain_inputs are valid
    if host == 'linux-x86':
      with open(os.path.join(install_dir, 'bin', 'remote_toolchain_inputs')) as inputs_file:
        files = [line.strip() for line in inputs_file.readlines()]
        fail = False
        for f in files:
          if not os.path.exists(os.path.join(install_dir, 'bin', f)):
            logger().error(f'remote_toolchain_inputs malformed, {f} does not exist')
            fail = True
        if fail:
          return False

    return True


def format_bug(bug):
    """Formats a bug for use in a commit message.

    Bugs might be a number, in which case they're a buganizer reference to be
    formatted. If not, assume the user knows what they're doing and just return
    the string as-is.
    """
    if bug.isnumeric():
        return f'http://b/{bug}'
    return bug


def update_clang(host, build_number, use_current_branch, download_dir, bug,
                 manifest, overwrite, do_validity_check, is_testing):
    prebuilt_dir = paths.PREBUILTS_DIR / 'clang' / 'host' / host
    os.chdir(prebuilt_dir)

    if not use_current_branch:
        branch_name = f'update-clang-{build_number}'
        utils.unchecked_call(
            ['repo', 'abandon', branch_name, '.'])
        utils.check_call(
            ['repo', 'start', branch_name, '.'])

    package = f'{download_dir}/clang-{build_number}-{host}.tar.xz'

    # Handle legacy versions of packages (like those from aosp/llvm-r365631).
    if not os.path.exists(package) and host == 'windows-x86':
        package = f'{download_dir}/clang-{build_number}-windows-x86-64.tar.xz'

    build_info_file = f'{download_dir}/BUILD_INFO-{host}'
    manifest_file = f'{download_dir}/{manifest}'

    utils.extract_tarball(prebuilt_dir, package)

    extract_subdir = 'clang-' + build_number
    clang_version_full, clang_version_major, svn_revision = extract_clang_info(extract_subdir)

    # Install into clang-<svn_revision>.  Suffixes ('a', 'b', 'c' etc.), if any,
    # are included in the svn_revision.
    install_subdir = 'clang-' + svn_revision
    install_clang_directory(extract_subdir, install_subdir, overwrite)

    # Linux prebuilts need to include a few libraries from the linux_musl artifacts
    if host == 'linux-x86':
        musl_install_subdir = install_subdir + '/musl'
        musl_package = f'{download_dir}/clang-{build_number}-linux_musl-x86.tar.xz'
        if os.path.exists(extract_subdir):
            shutil.rmtree(extract_subdir)
        utils.extract_tarball(prebuilt_dir, musl_package, [
            "--wildcards",
            "*/lib/libclang.so*",
            "*/lib/*/libc++.so*",
            "*/lib/libc_musl.so",
            "*/lib/aarch64-unknown-linux-musl/libc++.a",
            "*/lib/aarch64-unknown-linux-musl/libc++abi.a",
            "*/lib/x86_64-unknown-linux-musl/libc++.a",
            "*/lib/x86_64-unknown-linux-musl/libc++abi.a",
            ])
        install_clang_directory(extract_subdir, musl_install_subdir, overwrite)

        for triple in ('aarch64-unknown-linux-musl', 'x86_64-unknown-linux-musl'):
            # Move archives.
            src_dir = Path(musl_install_subdir) / 'lib' / triple
            dest_dir = Path(install_subdir) / 'lib' / 'clang' / clang_version_major / 'lib' / triple
            dest_dir.mkdir(exist_ok=True)  # The x86_64 triple will already exist.
            for name in ('libc++.a', 'libc++abi.a'):
                shutil.move(src_dir / name, dest_dir / name)

        with open(paths.KLEAF_VERSIONS_BZL) as f:
            kleaf_versions_lines = f.read().splitlines()
        new_version_line = '    "{}",'.format(svn_revision)
        list_end_idx = kleaf_versions_lines.index("]")
        if new_version_line not in kleaf_versions_lines:
            kleaf_versions_lines.insert(list_end_idx, new_version_line)
        with open(paths.KLEAF_VERSIONS_BZL, "w") as f:
            f.write("\n".join(kleaf_versions_lines))
        utils.check_call(['git', 'add', paths.KLEAF_VERSIONS_BZL])

    # Some platform tests (e.g. system/bt/profile/sdp) build directly with
    # coverage instrumentation and rely on the driver to pick the correct
    # profile runtime.  Symlink the Linux resource dir from the Linux toolchain
    # into the Darwin toolchain so the runtime is found by the Darwin Clang
    # driver.
    if host == 'darwin-x86':
        symlink_to_linux_resource_dir(install_subdir)

    if do_validity_check:
        if not validity_check(host, install_subdir, clang_version_major):
            sys.exit(1)

    shutil.copy(build_info_file, str(prebuilt_dir / install_subdir / 'BUILD_INFO'))
    shutil.copy(manifest_file, str(prebuilt_dir / install_subdir))

    utils.check_call(['git', 'add', install_subdir])

    # If there is no difference with the new files, we are already done.
    diff = utils.unchecked_call(['git', 'diff', '--cached', '--quiet'])
    if diff == 0:
        logger().info('Bypassed commit with no diff')
        return

    message_lines = [
        f'Update prebuilt Clang to {svn_revision} ({clang_version_full}).',
        '',
        f'clang {clang_version_full} (based on {svn_revision}) from build {build_number}.'
    ]
    if is_testing:
        message_lines.append('Note: This prebuilt is from testing branch.')
    if bug is not None:
        message_lines.append('')
        message_lines.append(f'Bug: {format_bug(bug)}')
    message_lines.append('Test: N/A')
    message = '\n'.join(message_lines)
    utils.check_call(['git', 'commit', '-m', message])


def install_clang_directory(extract_subdir: str, install_subdir: str, overwrite: bool):
    if os.path.exists(install_subdir):
        if overwrite:
            logger().info('Removing/overwriting existing path: %s',
                          install_subdir)
            shutil.rmtree(install_subdir)
        else:
            logger().info('Cannot remove/overwrite existing path: %s',
                          install_subdir)
            sys.exit(1)
    os.rename(extract_subdir, install_subdir)


def update_profiles(download_dir, build_number, bug):
    profiles_dir = paths.PREBUILTS_DIR / 'clang' / 'host' / 'linux-x86' / 'profiles'

    with contextlib.chdir(profiles_dir):
        # First, delete the old profiles.
        for f in glob.glob(PGO_PROFILE_PATTERN):
            os.remove(f)
        for f in glob.glob(BOLT_PROFILE_PATTERN):
            os.remove(f)

        # Replace with the downloaded new profiles.
        shutil.copy(glob.glob(f'{download_dir}/{PGO_PROFILE_PATTERN}')[0], '.')
        shutil.copy(glob.glob(f'{download_dir}/{BOLT_PROFILE_PATTERN}')[0], '.')

        utils.check_call(['git', 'add', '.'])
        message_lines = [f'Check in profiles from build {build_number}']
        if bug is not None:
            message_lines.append('')
            message_lines.append(f'Bug: {format_bug(bug)}')
        message_lines.append('Test: N/A')
        message = '\n'.join(message_lines)
        utils.check_call(['git', 'commit', '-m', message])


def main():
    args = ArgParser().parse_args()
    logging.basicConfig(level=logging.DEBUG)

    do_fetch = not args.skip_fetch
    do_cleanup = not args.skip_cleanup

    if do_fetch or args.repo_upload:
        utils.check_gcertstatus()

    download_dir = os.path.realpath('.download')
    if do_fetch:
        if os.path.isdir(download_dir):
            shutil.rmtree(download_dir)
        os.makedirs(download_dir)

    os.chdir(download_dir)

    targets_map = {'darwin-x86': 'darwin_mac',
                   'linux-x86': 'linux',
                   'windows-x86': 'windows_x86_64'}
    hosts = [args.host] if args.host else targets_map.keys()
    targets = [targets_map[h] for h  in hosts]
    if 'linux-x86' in hosts:
        targets.append('linux_musl')

    build_info = 'BUILD_INFO'
    clang_pattern = 'clang-*.tar.xz'
    manifest = f'manifest_{args.build}.xml'

    branch = args.branch
    if branch is None:
        output = utils.check_output(['/google/bin/releases/android/ab/ab.par',
                                     'get',
                                     '--raw', # prevent color text
                                     '--bid', args.build,
                                     '--target', 'linux'])
        # Example output is:
        #   aosp-llvm-toolchain linux 6732143 complete True
        branch = output.split()[0]

    logger().info('Using branch: %s', branch)
    is_testing = (branch == 'aosp-llvm-toolchain-testing')

    try:
        if do_fetch:
            fetch_artifact(branch, targets[0], args.build, manifest)
            for host in hosts:
                target = targets_map[host]
                fetch_artifact(branch, target, args.build, build_info)
                os.rename(f'{download_dir}/{build_info}', f'{download_dir}/{build_info}-{host}')
            for target in targets:
                fetch_artifact(branch, target, args.build, clang_pattern)

            if not args.skip_update_profiles and 'linux-x86' in hosts:
                fetch_artifact(branch, 'linux', args.build, PGO_PROFILE_PATTERN)
                fetch_artifact(branch, 'linux', args.build, BOLT_PROFILE_PATTERN)

        for host in hosts:
            update_clang(host, args.build, args.use_current_branch,
                         download_dir, args.bug, manifest, args.overwrite,
                         not args.no_validity_check, is_testing)

        if not args.skip_update_profiles and 'linux-x86' in hosts:
            update_profiles(download_dir, args.build, args.bug)

        if args.repo_upload:
            topic = f'clang-prebuilt-{args.build}'
            if is_testing:
                topic = topic.replace('prebuilt', 'testing-prebuilt')

            for host in hosts:
                utils.prebuilt_repo_upload(host, topic, args.hashtag, is_testing)
    finally:
        if do_cleanup:
            shutil.rmtree(download_dir)

    return 0


if __name__ == '__main__':
    main()