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

from collections import OrderedDict
from os import path, getcwd
import argparse
import glob
import json
import os
import platform
import sys
import tempfile

DEFAULT_TOOLS = [
    "cargo",
    "rustfmt",
    "cargo-fmt",
    "clippy",
    "src",
    "rust-demangler",
    "llvm-tools",
    "rust-analyzer",
    "rustdoc",
]

# TODO: http://fxbug.dev/310271984 - Remove the duplicated triples that are
# missing "-unknown".
FUCHSIA_TRIPLE_TO_GN_ARCH = {
    "aarch64-fuchsia": "arm64",
    "aarch64-unknown-fuchsia": "arm64",
    "riscv64gc-unknown-fuchsia": "riscv64",
    "x86_64-unknown-fuchsia": "x64",
    "x86_64-fuchsia": "x64",
}
FUCHSIA_TRIPLES = list(FUCHSIA_TRIPLE_TO_GN_ARCH.keys())
FUCHSIA_RUNTIME_TARGETS = [
    ("aarch64-fuchsia",),
    ("aarch64-unknown-fuchsia",),
    ("riscv64gc-unknown-fuchsia", "riscv64-unknown-fuchsia"),
    ("x86_64-fuchsia",),
    ("x86_64-unknown-fuchsia",),
]


CMAKE_SETTINGS = """
cflags = "{cflags}"
cxxflags = "{cxxflags}"
ldflags = "{ldflags}"
"""

BOOTSTRAP_CONFIG = """
# Use our own compiler as stage0 for bootstrapping, rather than downloading
cargo = "{stage0}/bin/cargo"
rustc = "{stage0}/bin/rustc"
rustfmt = "{stage0}/bin/rustfmt"
"""

BUILD_CONFIG = """
change-id = 123711

[llvm]
optimize = true
thin-lto = {llvm_thinlto}
ccache = {gomacc_quote}
ninja = true
targets = "AArch64;ARM;RISCV;WebAssembly;X86"
{cmake_settings}
use-libcxx = true
# TODO(https://github.com/rust-lang/rust/issues/94983)
# We shouldn't have to set this to false to enable statically linking libc++
static-libstdcpp = false

[build]
target = [{target}]
{host_triple}
build-dir = "fuchsia-build"
submodules = false
docs = false
extended = true
cargo-native-static = true
tools = [{tools}]
{bootstrap}

[install]
prefix = "{prefix}"
sysconfdir = "etc"

[rust]
optimize = true
codegen-units-std = 1
channel = "{channel}"
lld = false
llvm-libunwind = "in-tree"
jemalloc = {jemalloc}
{rust_thinlto}

# Enable full debuginfo for libstd.
debuginfo-level-std = 2

[dist]
"""

TARGET_CONFIG = """
[target.{triple}]
cc = "{cc}"
cxx = "{cxx}"
ar = "{ar}"
ranlib = "{ranlib}"
linker = "{linker}"
profiler = {profiler}
llvm-has-rust-patches = {llvm_has_rust_patches}
"""

# This uses % formatting to avoid interactions with the brace syntax used in bash.
TOOL_WRAPPER_SCRIPT = """#!/bin/bash

fail() {
    cat >&2 <<EOF
FATAL: Detected system library path '$1' in the linker search directory.
       This leads to interactions between Fuchsia sysroot linker scripts
       and system libraries that cause horribly confusing linker errors,
       so we're going to error here instead.
EOF
    exit 1
}

# Check for system paths.
# For now we only look for paths likely to be hit on Linux.
shopt -s extglob
for arg in "$@"; do
    case "$arg" in
        ?(-L)/lib) fail "$arg"; break;;
        ?(-L)/usr/lib) fail "$arg"; break;;
        ?(-L)/usr/lib/*) fail "$arg"; break;;
    esac
done

exec %(linker)s %(flags)s "$@"
"""


def _cflags(sysroot):
    return "--sysroot=%s" % sysroot


def _cxxflags(sysroot):
    return "--sysroot=%s -stdlib=libc++" % sysroot


def _ldflags(args, sysroot):
    return "--sysroot=%s" % sysroot + (
        " -L%s" % path.join(args.clang_prefix, "lib")
        if "apple" in args.host
        else " -static-libstdc++"
    )


def clang_tool(args, tool, flags=None):
    tool_binary = path.join(args.clang_prefix, "bin", tool) + (
        ".exe" if "windows" in args.host else ""
    )

    # Create a wrapper script to include necessary args (e.g. --sysroot) as
    # Cargo does not properly select RUSTFLAGS when `--target` is passed.
    if flags:
        fd, name = tempfile.mkstemp()
        with os.fdopen(fd, "w") as f:
            f.write(TOOL_WRAPPER_SCRIPT % dict(linker=tool_binary, flags=flags))
        os.chmod(name, 0o775)
        return name
    return tool_binary


def generate_config_toml(args):
    if args.thinlto is None:
        args.thinlto = "linux" in args.host and args.channel in ["nightly", "dev"]

    targets = args.targets.split(",")
    if "windows" in args.host:
        # don't build fuchsia targets on windows
        pass
    elif "apple" in args.host:
        # TODO(https://fxbug.dev/126336): fix riscv64gc-unknown-fuchsia build on mac
        targets += [t for t in FUCHSIA_TRIPLES if t != "riscv64gc-unknown-fuchsia"]
    else:
        targets += FUCHSIA_TRIPLES
    tools = args.tools.split(",")

    def string_list(items):
        return ", ".join('"%s"' % item for item in items)

    cmake_fragment = (
        'build-config = { CMAKE_SYSROOT="%s" }' % args.host_sysroot
        if "windows" in args.host
        else CMAKE_SETTINGS.format(
            cflags=_cflags(args.host_sysroot),
            cxxflags=_cxxflags(args.host_sysroot),
            ldflags=_ldflags(args, args.host_sysroot),
        )
    )

    config_toml = BUILD_CONFIG.format(
        prefix=args.prefix,
        jemalloc=str(not "windows" in args.host).lower(),
        gomacc_quote=(
            ('"%s"' % path.join(args.goma_dir, "gomacc")) if args.goma_dir else "false"
        ),
        cmake_settings=cmake_fragment,
        llvm_thinlto=str(args.thinlto).lower(),
        rust_thinlto='lto = "thin"' if args.thinlto else "",
        target=string_list(targets),
        channel=args.channel,
        tools=string_list(tools),
        # For windows specify that host is msvc not mingw
        host_triple=(
            'build = "x86_64-pc-windows-msvc"' if "windows" in args.host else ""
        ),
        bootstrap=BOOTSTRAP_CONFIG.format(stage0=args.stage0) if args.stage0 else "",
    )
    for target in targets:
        cflags = None
        cxxflags = None
        ldflags = None
        if args.host == target.lower():
            cflags = _cflags(args.host_sysroot)
            cxxflags = _cxxflags(args.host_sysroot)
            ldflags = _ldflags(args, args.host_sysroot)
        cc = clang_tool(args, "clang", cflags)
        cxx = clang_tool(args, "clang++", cxxflags)
        linker = clang_tool(args, "clang++", ldflags)
        ar = clang_tool(args, "llvm-ar")
        if "fuchsia" in target:
            linker = clang_tool(args, "ld.lld")
        elif "windows" in target:
            cc = cxx = clang_tool(args, "clang-cl")
            linker = clang_tool(args, "lld-link")
            ar = clang_tool(args, "llvm-lib")

        config_toml += TARGET_CONFIG.format(
            triple=target,
            cc=cc,
            cxx=cxx,
            ar=ar,
            ranlib=clang_tool(args, "llvm-ranlib"),
            linker=linker,
            # Enabling the profiler for wasm caused the build to fail, so we
            # disable it for now.
            profiler="false" if "wasm" in target else "true",
            llvm_has_rust_patches="false" if args.llvm_is_vanilla else "true",
        )
    print(config_toml.replace("\\", "/"))


def generate_env(args):
    env = OrderedDict()
    env["RUST_BACKTRACE"] = 1

    if args.revision:
        env["CFG_VERSION"] = args.revision

    # Tell bootstrap binary not to link to system libs.
    env["LZMA_API_STATIC"] = 1

    # Configure Fuchsia targets.
    for target, gn_arch in list(FUCHSIA_TRIPLE_TO_GN_ARCH.items()):
        triple_lower = target.replace("-", "_").lower()
        triple_upper = triple_lower.upper()
        clang_target = target.replace("riscv64gc", "riscv64")
        env["CC_%s" % triple_lower] = clang_tool(args, "clang")
        env["CXX_%s" % triple_lower] = clang_tool(args, "clang++")
        env["CFLAGS_%s" % triple_lower] = "--target=%s --sysroot=%s -I%s" % (
            clang_target,
            path.join(args.sdk_dir, "arch", gn_arch, "sysroot"),
            path.join(args.sdk_dir, "pkg", "fdio", "include"),
        )
        env["CXXFLAGS_%s" % triple_lower] = "--target=%s --sysroot=%s -I%s" % (
            clang_target,
            path.join(args.sdk_dir, "arch", gn_arch, "sysroot"),
            path.join(args.sdk_dir, "pkg", "fdio", "include"),
        )
        env["LDFLAGS_%s" % triple_lower] = "--target=%s --sysroot=%s -L%s" % (
            clang_target,
            path.join(args.sdk_dir, "arch", gn_arch, "sysroot"),
            path.join(args.sdk_dir, "arch", gn_arch, "lib"),
        )
        env["CARGO_TARGET_%s_AR" % triple_upper] = clang_tool(args, "llvm-ar")
        env["CARGO_TARGET_%s_RUSTFLAGS" % triple_upper] = " ".join(
            [
                "-Clink-arg=--sysroot=%s"
                % path.join(args.sdk_dir, "arch", gn_arch, "sysroot"),
                "-Lnative=%s/lib" % path.join(args.sdk_dir, "arch", gn_arch, "sysroot"),
                "-Lnative=%s" % path.join(args.sdk_dir, "arch", gn_arch, "lib"),
                "-Cpanic=abort",
                "-Cforce-unwind-tables=yes",
            ]
        )

    target_sysroots = {
        "aarch64-unknown-linux-gnu": args.linux_arm64_sysroot or args.linux_sysroot,
        "x86_64-unknown-linux-gnu": args.linux_amd64_sysroot or args.linux_sysroot,
        "riscv64-unknown-linux-gnu": args.linux_riscv64_sysroot or args.linux_sysroot,
    }
    if args.darwin_sysroot:
        target_sysroots.update({"aarch64-apple-darwin": args.darwin_sysroot})
        target_sysroots.update({"x86_64-apple-darwin": args.darwin_sysroot})
    if args.win_sysroot:
        target_sysroots.update({"x86_64-pc-windows-msvc": args.win_sysroot})

    # Configure non-Fuchsia target sysroots.
    for target in args.targets.split(","):
        clang_target = target.replace("riscv64gc", "riscv64")
        # These are the only targets we define a sysroot for.
        if not ("linux" in target or "apple" in target):
            continue
        sysroot = target_sysroots[clang_target]
        assert sysroot, f"Sysroot path not supplied for target {target}"
        arch = clang_target.split("-")[0]
        if "linux" in target:
            assert any(
                p.startswith(f"{arch}-")
                for p in os.listdir(os.path.join(sysroot, "lib"))
            ), f"Sysroot does not appear to support {arch}"

        triple_lower = target.replace("-", "_").lower()
        triple_upper = triple_lower.upper()
        env["CC_%s" % triple_lower] = clang_tool(args, "clang", _cflags(sysroot))
        env["CXX_%s" % triple_lower] = clang_tool(args, "clang++", _cxxflags(sysroot))
        env["CFLAGS_%s" % triple_lower] = " ".join(
            ["--target=" + clang_target, "--sysroot=" + sysroot]
            # TODO(https://fxbug.dev/69260): Find a longer-term solution.
            + (["-mno-outline-atomics"] if "aarch64" in target else [])
        )
        env["CXXFLAGS_%s" % triple_lower] = "--target=%s --sysroot=%s" % (
            clang_target,
            sysroot,
        )
        env["LDFLAGS_%s" % triple_lower] = "--target=%s --sysroot=%s" % (
            clang_target,
            sysroot,
        ) + (" -fuse-ld=lld" if "apple" not in target else "")
        env["CARGO_TARGET_%s_AR" % triple_upper] = clang_tool(args, "llvm-ar")
        env["CARGO_TARGET_%s_RUSTFLAGS" % triple_upper] = " ".join(
            [
                "-Clink-arg=--target=%s" % clang_target,
                "-Clink-arg=--sysroot=%s" % sysroot,
                # In the cases where we link libc++ we need to let the driver (clang++) use
                # its knowledge of other dependencies, like libunwind. This only happens
                # when building for host (linking to libLLVM), but we specify this flag for
                # all non-Fuchsia targets for now.
                "-Cdefault-linker-libraries=yes",
            ]
            + (["-Clink-arg=-fuse-ld=lld"] if "apple" not in target else [])
        )

    # FIXME(http://b/326089386): Rust appears to have an issue generating the
    # right target triple for riscv32imc when compiling code with clang.
    # Explicitly specify the target to get the build to work while we figure
    # out how to fix upstream rust.
    env["CFLAGS_riscv32imc_unknown_none_elf"] = (
        "-ffunction-sections -fdata-sections --target=riscv32"
    )

    for key, val in list(env.items()):
        if args.eval:
            print('export {}="{}"'.format(key, val).replace("\\", "/"))
        else:
            print("{}={}".format(key, val).replace("\\", "/"))


# pylint: disable=unused-argument
def generate_runtimes_spec(args):
    """
    Unused argument `args` required to preserve usage pattern:
    ```
    parser_env.set_defaults(func=generate_env)
    parser_runtime.set_defaults(func=generate_runtimes_spec)

    args = parser.parse_args()
    args.func(args)
    ```
    """
    runtimes = []
    for dynlink_flags in [[], ["-Cprefer-dynamic"]]:
        for targets in FUCHSIA_RUNTIME_TARGETS:
            if "apple" in args.host and targets[0] == "riscv64gc-unknown-fuchsia":
                # TODO(https://fxbug.dev/126336): fix riscv64gc-unknown-fuchsia build on mac
                continue
            runtime = []
            if dynlink_flags:
                for lib in ["libstd", "libtest"]:
                    # The runtimes module will expand this wildcard
                    # (it must match exactly one file.)
                    pattern = "rustlib/%s/lib/%s-*.so" % (targets[0], lib)
                    matches = glob.glob(os.path.join(args.prefix, "lib", pattern))
                    if "libtest" == lib and len(matches) == 0:
                        # TODO: http://fxbug.dev/318746783 - libtest might be missing due to recent
                        # upstream changes. Hopefully we can remove this branch once things are
                        # settled.
                        continue
                    runtime.append(
                        {
                            "dist": pattern,
                            "name": lib,
                        }
                    )
            runtimes.append(
                {
                    "target": targets,
                    "rustflags": dynlink_flags,
                    "runtime": runtime,
                }
            )
    print(json.dumps(runtimes))


def main():
    systems = {
        "Linux": "unknown-linux-gnu",
        "Darwin": "apple-darwin",
        "Windows": "pc-windows",
    }
    host = f"{platform.machine()}-{systems[platform.system()]}".replace(
        # Convert the arch returned by uname on Macs to the LLVM spelling.
        "arm64",
        "aarch64",
    )
    default_targets = [
        "aarch64-unknown-linux-gnu",
        "x86_64-unknown-linux-gnu",
        "thumbv6m-none-eabi",
        "thumbv7m-none-eabi",
        # FIXME(https://fxbug.dev/325488864): Enable once the build for this target is working again.
        # "wasm32-unknown-unknown",
    ]
    if "apple" in host:
        default_targets.append("aarch64-apple-darwin")
        default_targets.append("x86_64-apple-darwin")
    else:
        # TODO(https://fxbug.dev/42079174): Enable for mac hosts.
        default_targets.append("riscv32imc-unknown-none-elf")
        default_targets.append("riscv64gc-unknown-linux-gnu")
    if "windows" in host:
        default_targets = ["x86_64-pc-windows-msvc"]

    parser = argparse.ArgumentParser()
    parser.set_defaults(host=host)

    subparsers = parser.add_subparsers()
    parser_config = subparsers.add_parser("config_toml")
    parser_env = subparsers.add_parser("environment")
    parser_runtime = subparsers.add_parser("runtime")

    for subcommand in [parser_config, parser_env]:
        subcommand.add_argument(
            "--stage0",
            help="path to stage0 toolchain, will download from upstream by default",
        )
        subcommand.add_argument(
            "--targets",
            default=",".join(default_targets),
            help="comma-separated list of targets to build for, in addition to Fuchsia",
        )
        subcommand.add_argument(
            "--clang-prefix", help="path to clang toolchain", required=True
        )

    parser_config.add_argument("--prefix", help="installation prefix", required=True)
    parser_config.add_argument(
        "--host-sysroot", help="path to host sysroot", required=True
    )
    parser_config.add_argument(
        "--llvm-is-vanilla",
        action="store_true",
        help="set if using upstream LLVM",
        required=False,
    )
    parser_config.add_argument("--goma-dir", help="path to goma", required=False)
    parser_config.add_argument(
        "--thinlto",
        action="store_true",
        help="enable thinlto (enabled by default on linux nightly builds)",
    )
    parser_config.add_argument(
        "--channel",
        help="what compiler channel to build: dev, nightly, beta, or stable",
        default="nightly",
    )
    parser_config.add_argument(
        "--tools",
        help="comma-separated list of tools to build",
        default=",".join(DEFAULT_TOOLS),
    )
    parser_config.set_defaults(func=generate_config_toml)

    parser_env.add_argument("--source", help="rust source path", default=getcwd())
    parser_env.add_argument("--sdk-dir", help="path to Fuchsia SDK", required=True)
    parser_env.add_argument("--linux-sysroot", help="path to Linux sysroot")
    parser_env.add_argument(
        "--linux-amd64-sysroot",
        help="path to Linux amd64 sysroot (overrides --linux-sysroot)",
    )
    parser_env.add_argument(
        "--linux-arm64-sysroot",
        help="path to Linux arm64 sysroot (overrides --linux-sysroot)",
    )
    parser_env.add_argument(
        "--linux-riscv64-sysroot",
        help="path to Linux riscv64 sysroot (overrides --linux-sysroot)",
    )
    parser_env.add_argument(
        "--darwin-sysroot", help="path to Darwin sysroot, if building on Mac"
    )
    parser_env.add_argument(
        "--win-sysroot", help="path to Windows sysroot, if building on Windows"
    )
    parser_env.add_argument("--revision", help="git commit of Rust repository")
    parser_env.add_argument(
        "--eval",
        default=False,
        action="store_true",
        help="format output for evaluating in a terminal",
    )

    parser_env.set_defaults(func=generate_env)
    parser_runtime.set_defaults(func=generate_runtimes_spec)
    # TODO: http://fxbug.dev/318746783 - Maybe remove this option. It's
    # currently defaulted to reduce churn on the local workflow, but that could
    # be a bit error-prone.
    parser_runtime.add_argument(
        "--prefix", help="installation prefix", default="install/fuchsia-rust"
    )

    args = parser.parse_args()
    args.func(args)

    return 0


if __name__ == "__main__":
    sys.exit(main())