#!/usr/bin/env python3
#
# Copyright (C) 2022 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.

import argparse
from collections.abc import Iterator
from contextlib import contextmanager
from pathlib import Path
import shutil
import subprocess
import sys
from tempfile import TemporaryDirectory
from typing import Optional, TextIO

import context

from android_rust import build_platform
from android_rust.paths import *
from android_rust.toolchains import CLANG_TOOLCHAIN_HOST, ClangToolchain, RustToolchain
from android_rust.utils import (
    TERM_RED,
    ExtantPath,
    ResolvedPath,
    ScriptException,
    archive_create,
    archive_extract,
    extend_suffix,
    get_prebuilt_binary_paths,
    print_colored,
    print_fs_tree,
    is_archive,
    reify_singleton_patterns)

from test_compiler import initialize_clang_prebuilt

#
# Constants
#

BOLT_INSTRUMENTATION_SUBJECTS: list[str] = [
    "lib/librustc_driver-*.so",
    "lib/libstd-*.so",
    "lib/libLLVM.so.*",
]

BOLT_SKIP_FUNCS: dict[str, list[str]] = {
    "librustc_driver": ["blake3_hash_many_sse2/1",],
}

#
# Program logic
#


def parse_args() -> argparse.Namespace:
    """Parses arguments and returns the parsed structure."""
    parser = argparse.ArgumentParser("Run BOLT on a Rust toolchain archive")
    parser.add_argument(
        "archive_path", type=ExtantPath, help="Path to the prebuilt archive to instrument/optimize")

    parser.add_argument(
        "--build-name",
        default="bolted",
        help="Desired filename for the output archive (w/o extension)")
    parser.add_argument(
        "--dist",
        "-d",
        dest="dist_path",
        type=ResolvedPath,
        default=DIST_PATH_DEFAULT,
        help="Where to place distributable artifacts")
    parser.add_argument(
        "--clang-prebuilt",
        type=ExtantPath,
        default=CLANG_TOOLCHAIN_HOST,
        help="Path to Clang toolchain that will be used to merge profile data")

    parser.add_argument(
        "--strip-only", action="store_true", help="Don't invoke BOLT, just strip the binaries")

    action_group = parser.add_mutually_exclusive_group()
    action_group.add_argument(
        "--profile-generate",
        type=ExtantPath,
        nargs="?",
        const=OUT_PATH_PROFILES,
        help="Instrument the toolchain for BOLT profile generation with an absolute path")
    action_group.add_argument(
        "--profile-generate-relative",
        type=Path,
        nargs="?",
        const=Path("out/profiles"),
        help="Instrument the toolchain for BOLT profile generation with a relative path")
    action_group.add_argument(
        "--profile-use",
        type=ExtantPath,
        nargs="?",
        const=OUT_PATH_PROFILES,
        help="Path to archive or directory with a bolt/ subdir containing BOLT profiles")

    return parser.parse_args()


@contextmanager
def handle_input_path(input_path: Path) -> Iterator[Path]:
    if is_archive(input_path):
        with TemporaryDirectory() as temp_dir:
            temp_dir_path = Path(temp_dir)
            print(f"Unpacking archive into {temp_dir}")
            archive_extract(input_path, temp_dir_path)
            yield temp_dir_path
    else:
        yield input_path


@contextmanager
def expand_profiles(profile_path: Path | None) -> Iterator[Path | None]:
    if profile_path is not None:
        if is_archive(profile_path):
            with TemporaryDirectory() as temp_dir:
                temp_dir_path = Path(temp_dir)
                print(f"Unpacking profiles into {temp_dir}")
                archive_extract(profile_path, temp_dir_path)
                yield temp_dir_path
        else:
            yield profile_path
    else:
        yield None


def invoke_bolt(
        toolchain: ClangToolchain,
        obj_path: Path,
        bolt_log: TextIO,
        options: list[str] = []) -> None:
    bolt_path = toolchain.bolt()
    assert (bolt_path is not None)

    obj_path_bolted = extend_suffix(obj_path, ".bolt")

    bolt_log.write(f"BOLTing {str(obj_path)}\n")
    print(f"BOLTing {str(obj_path)}\n")
    bolt_log.flush()

    for (lib_prefix, func_names) in BOLT_SKIP_FUNCS.items():
        if obj_path.name.startswith(lib_prefix):
            options = options.copy()
            options.append(f"--skip-funcs={','.join(func_names)}")

    try:
        subprocess.run([bolt_path] + options + ["-o", obj_path_bolted, obj_path],
                       check=True,
                       stdout=bolt_log,
                       stderr=bolt_log)
    except subprocess.CalledProcessError as err:
        print(f"Failed to process object {str(obj_path)}")
        raise err

    bolt_log.write("\n")
    bolt_log.flush()

    obj_path.unlink()
    obj_path_bolted.rename(obj_path)


def process_objects(
        toolchain: ClangToolchain,
        root: Path,
        strip_only: bool,
        profile_generate: Optional[Path],
        profile_use: Optional[Path],
        bolt_log: TextIO) -> None:

    print("Processing objects")

    instrumentation_subjects = reify_singleton_patterns(root, BOLT_INSTRUMENTATION_SUBJECTS)
    optimization_subjects = get_prebuilt_binary_paths(root, toplevel_only=True)

    for obj_path in get_prebuilt_binary_paths(root, subdirs=["bin", "lib"]):
        print(f"Boltifyer is looking at object {str(obj_path)}")

        if not strip_only and obj_path in optimization_subjects:
            print(f"Object {str(obj_path)} is an optimization subject")
            options = ["--peepholes=all"]
            if obj_path in instrumentation_subjects:
                if profile_generate:
                    fdata_path = profile_generate / PROFILE_SUBDIR_BOLT / obj_path.name
                    options += [
                        "--instrument",
                        f"--instrumentation-file={fdata_path}",
                        "--instrumentation-file-append-pid",
                    ]

                elif profile_use:
                    fdata_path = profile_use / PROFILE_SUBDIR_BOLT / (obj_path.name + ".fdata")

                    if fdata_path.exists():
                        options += [
                            f"--data={fdata_path}",
                            "--reorder-blocks=ext-tsp",
                            "--reorder-functions=hfsort",
                            "--split-functions",
                            "--split-all-cold",
                            "--split-eh",
                            "--dyno-stats",
                        ]
                    else:
                        print(f"Profile data missing for {obj_path.relative_to(root)}")

            invoke_bolt(toolchain, obj_path, bolt_log, options)

        else:
            print(f"Object {str(obj_path)} is NOT an optimization subject\n")
            toolchain.strip_symbols(obj_path)


def boltify_toolchain(
        input_path: Path,
        dist_path: Path,
        build_name: str,
        strip_only: bool = False,
        profile_generate: Path | None = None,
        profile_use: Path | None = None,
        toolchain: ClangToolchain = CLANG_TOOLCHAIN_HOST) -> None:
    with (dist_path / BOLT_LOG_NAME).open("w") as bolt_log:
        with handle_input_path(input_path) as input_dir_path:
            print_fs_tree(input_dir_path)

            with expand_profiles(profile_use) as profile_use_path:
                if profile_use_path is not None:
                    print_fs_tree(profile_use_path)

                    version_file_path = profile_use_path / PROFILE_SUBDIR_BOLT / PROFILED_VERSION_FILENAME
                    if version_file_path.exists():
                        with open(version_file_path) as version_file:
                            profiled_version = version_file.read()
                            input_toolchain = RustToolchain(
                                input_dir_path, host=build_platform.system())
                            if input_toolchain.get_version() != profiled_version:
                                raise ScriptException(
                                    "Profiled compiler and input compiler versions do not match")

                try:
                    process_objects(
                        toolchain,
                        input_dir_path,
                        strip_only,
                        profile_generate,
                        profile_use_path,
                        bolt_log)
                    bolt_log.flush()
                except subprocess.CalledProcessError as err:
                    print("Error BOLTing prebuilt objects")
                    raise err

                print(f"Creating BOLTed archive")
                try:
                    archive_create(dist_path / f"rust-{build_name}", input_dir_path, overwrite=True)
                except subprocess.CalledProcessError as err:
                    print("Error creating the BOLTed archive")
                    raise err


def main() -> None:
    if not build_platform.is_linux():
        raise ScriptException("BOLT tools are only available on Linux hosts")

    args = parse_args()

    initialize_clang_prebuilt(args)

    boltify_toolchain(
        args.archive_path,
        args.dist_path,
        args.build_name,
        args.strip_only,
        args.profile_generate or args.profile_generate_relative,
        args.profile_use)

    if args.profile_use is not None and args.profile_use.name.endswith(
            ".tar.xz") and args.profile_use.parent != args.dist_path:
        shutil.copy2(args.profile_use, args.dist_path)


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

    except (ScriptException, argparse.ArgumentTypeError) as err:
        print_colored(str(err), TERM_RED)
        sys.exit(1)