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

import argparse
import json
import os
import subprocess
import sys
from itertools import chain, product, starmap
from pathlib import Path
from typing import Dict, Iterable, List, NamedTuple

from impl.common import CROSVM_ROOT, Triple, verbose
from impl.test_config import DO_NOT_BUILD_RISCV64

USAGE = """\
Build crosvm with release (optimized) profile.

To target local machine:

    $ ./tools/build_release

To cross-compile for aarch64, armhf or windows you can use:

    $ ./tools/build_release --platform=aarch64
    $ ./tools/build_release --platform=armhf
    $ ./tools/build_release --platform=mingw64
"""

# We only need PGO for main binary, but for consistency, only exclude incompatible parts from PGO
PGO_EXCLUDE = ["crosvm_control"]


class Executable(NamedTuple):
    """Container for info about an executable generated by cargo build/test."""

    binary_path: Path
    crate_name: str
    cargo_target: str
    kind: str
    is_test: bool
    is_fresh: bool

    @property
    def name(self):
        return f"{self.crate_name}:{self.cargo_target}"


def cargo(
    cargo_command: str,
    cwd: Path,
    flags: List[str],
    env: Dict[str, str],
) -> Iterable[Executable]:
    """
    Executes a cargo command and returns the list of test binaries generated.

    The build log will be hidden by default and only printed if the build
    fails. In VERBOSE mode the output will be streamed directly.

    Note: Exits the program if the build fails.
    """
    message_format = "json-diagnostic-rendered-ansi" if sys.stdout.isatty() else "json"
    cmd = [
        "cargo",
        cargo_command,
        f"--message-format={message_format}",
        *flags,
    ]
    if verbose():
        print("$", " ".join(cmd))
    process = subprocess.Popen(
        cmd,
        cwd=cwd,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        env=env,
    )

    messages: List[str] = []

    # Read messages as cargo is running.
    assert process.stdout
    for line in iter(process.stdout.readline, ""):
        # any non-json line is a message to print
        if not line.startswith("{"):
            if verbose():
                print(line.rstrip())
            messages.append(line.rstrip())
            continue
        json_line = json.loads(line)

        # 'message' type lines will be printed
        if json_line.get("message"):
            message = json_line.get("message").get("rendered")
            if verbose():
                print(message)
            messages.append(message)

        # Collect info about test executables produced
        elif json_line.get("executable"):
            yield Executable(
                Path(json_line.get("executable")),
                crate_name=json_line.get("package_id", "").split(" ")[0],
                cargo_target=json_line.get("target").get("name"),
                kind=json_line.get("target").get("kind")[0],
                is_test=json_line.get("profile", {}).get("test", False),
                is_fresh=json_line.get("fresh", False),
            )

    if process.wait() != 0:
        if not verbose():
            for message in messages:
                print(message)
        sys.exit(-1)


def main():
    parser = argparse.ArgumentParser(usage=USAGE)
    parser.add_argument(
        "--build-target",
        "--platform",
        "-p",
        help=(
            "Override the cargo triple to build. Shorthands are available: (x86_64, armhf, "
            + "aarch64, mingw64, msvc64)."
        ),
    )
    parser.add_argument(
        "--json",
        action="store_true",
        help="Output in JSON instead of human readable format.",
    )
    parser.add_argument("--strip", action="store_true", help="Strip output binaries")
    parser.add_argument(
        "--build-profile", help="Select compile profile, default to release.", default="release"
    )
    pgo_group = parser.add_mutually_exclusive_group()
    pgo_group.add_argument(
        "--profile-generate",
        help="Target directory to generate profile when running, must use absolute path",
    )
    pgo_group.add_argument(
        "--profile-use", help="Profile file used for PGO, must use absolute path"
    )
    parser.add_argument("cargo_arguments", nargs="*", help="Extra arguments pass to cargo")

    args = parser.parse_args()

    if args.profile_generate and (
        not os.path.isabs(args.profile_generate) or not os.path.isdir(args.profile_generate)
    ):
        raise ValueError("--profile-generate argument is not an absolute path to a folder")
    if args.profile_use and (
        not os.path.isabs(args.profile_use) or not os.path.isfile(args.profile_use)
    ):
        raise ValueError("--profile-use argument is not an absolute path to a file")

    build_target = Triple.from_shorthand(args.build_target) if args.build_target else None
    build_target = build_target or Triple.host_default()

    exclude_args = [
        f"--exclude={x}" for x in PGO_EXCLUDE if args.profile_generate or args.profile_use
    ]
    if build_target == Triple.from_shorthand("riscv64"):
        exclude_args += ["--exclude=" + s for s in DO_NOT_BUILD_RISCV64]

    features = build_target.feature_flag
    cargo_args = [
        "--profile",
        args.build_profile,
        "--features=" + features,
        f"--target={build_target}",
        "--workspace",
        *exclude_args,
        *args.cargo_arguments,
    ]

    build_env = os.environ.copy()
    build_env.update(build_target.get_cargo_env())
    build_env.setdefault("RUSTFLAGS", "")
    build_env["RUSTFLAGS"] += " -D warnings"
    if args.strip:
        build_env["RUSTFLAGS"] += " -C strip=symbols"
    if args.profile_generate:
        build_env["RUSTFLAGS"] += " -C profile-generate=" + args.profile_generate
    if args.profile_use:
        build_env["RUSTFLAGS"] += " -C profile-use=" + args.profile_use

    executables = list(cargo("build", CROSVM_ROOT, cargo_args, build_env))

    if args.json:
        result = {}
        for exe in executables:
            assert exe.cargo_target not in result
            result[exe.cargo_target] = str(exe.binary_path)
        print(json.dumps(result))
    else:
        print("Release binaries:")
        for exe in executables:
            print(f"Name: {exe.cargo_target}")
            print(f"Path: {str(exe.binary_path)}")
            print()


if __name__ == "__main__":
    main()