#!/usr/bin/env fuchsia-vendored-python
# Copyright 2025 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.

"""Describe or export Bazel host test artifacts to the Ninja build directory."""

import argparse
import dataclasses
import json
import os
import shutil
import sys
import typing as T
from pathlib import Path

_SCRIPT_DIR = os.path.dirname(__file__)
_SCRIPT_DIR_PATH = Path(_SCRIPT_DIR)
sys.path.insert(0, _SCRIPT_DIR)
import build_utils
import runfiles_utils

# The main workspace name as it appears in the runfiles directory.
# For the Fuchsia build, which is still using Bazel 7.x and WORKSPACE
# files, this is "main", but will become `_main` when enabling BzlMod
_MAIN_WORKSPACE_NAME = "main"

_LIST_TESTS_COMMAND_HELP = r"""
Generate a Fuchsia-compatible `tests.json` file listing host test targets
reachable from one ore more Bazel target pattern(s). This can be called
directly after `fx set` or `fx gen`, as it does not require building anything.

The only difference from the GN-generated `tests.json` file is that each
entry's target label begins with an `@`, and does not include any toolchain
suffix. For example, the following command:

```bash
python3 build/bazel/scripts/export_host_test.py \
    list_tests \
    //build/bazel/host_tests/static_cc_test/...
```

Produces the following output:

```json
[
  {
    "environments": [
      {
        "dimensions": {
          "cpu": "x64",
          "os": "Linux"
        }
      }
    ],
    "expects_ssh": false,
    "label": "@//build/bazel/host_tests/static_cc_test:test_program",
    "name": "bazel_host_tests/build/bazel/host_tests/static_cc_test/test_program",
    "path": "bazel_host_tests/build/bazel/host_tests/static_cc_test/test_program",
    "runtime_deps": "bazel_host_tests/build/bazel/host_tests/static_cc_test/test_program.runtime_deps.json",
    "os": "linux",
    "cpu": "x64"
  }
]
```

All exported test artifacts are located under $NINJA_BUILD_DIR/bazel_host_tests/ by default. This
location can be changed with the --export-subdir=SUBDIR argument. The rest of the path must follow
the Bazel-specific execroot-relative file layout.

Use the `export` command to create the test artifacts and the `.runtime_deps.json` file.
"""

_EXPORT_COMMAND_HELP = r"""Export Bazel test artifacts to the Ninja build directory. Making them
usable outside of the Bazel output base.

For example, the following command:

```bash
python3 build/bazel/scripts/export_host_test.py \
    export \
    //build/bazel/host_tests/static_cc_test/...
```

Will populate the build directory with files that look like:

```
bazel_host_tests/build/bazel/host_tests/
    static_cc_test/
        test_program
        test_program.runtime_deps.json
        test_program.repo_mapping
        test_program.runfiles_manifest
        test_program.runfiles/
            ...
```

Where `test_program` is the test entry point that is listed in the
corresponding entry generated by the `list` command.

NOTE: In many cases, the test executable (e.g. `test_program` here) *MUST* be invoked from
      the Ninja build directory to work correctly.
"""

# Location of the file containing the Starlark cquery code to run in this script
# to locate all test targets.
_CQUERY_STARLARK_FILE_PATH = (
    _SCRIPT_DIR_PATH.parent / "starlark" / "host_test_info.cquery"
)

# The sub-directory, relative to the Ninja build directory, where all Bazel host test
# artifacts will be exported by this script. The content will be mostly symlinks into
# the Bazel output base, unless the --no-symlinks option is used.
_DEFAULT_EXPORT_SUBDIR = "bazel_host_tests"


class HostInfo(object):
    """Convenience class to expose host-specific information."""

    def __init__(self) -> None:
        self._os = build_utils.get_host_platform()
        self._cpu = build_utils.get_host_arch()

    @property
    def os(self) -> str:
        """The host operating system name, following Fuchsia conventions."""
        return self._os

    @property
    def cpu(self) -> str:
        """The host cpu architecture name, following Fuchsia conventions."""
        return self._cpu

    @property
    def test_json_environments(self) -> list[dict[str, T.Any]]:
        """The value of a tests.json "environments" key for a given host test."""
        # See https://fuchsia.dev/fuchsia-src/reference/testing/tests-json-format?hl=en#environments_field
        # and //build/testing/environments.gni.
        # There are only two host environment supported, Linux and MacOS
        # using environments[].dimensions.os values "Linux" and "Mac"
        # respectively.
        host_os = {"linux": "Linux", "mac": "Mac"}.get(self._os)
        assert host_os, f"Unsupported host operating system: {self._os}"

        return [
            {
                "dimensions": {
                    "cpu": self._cpu,
                    "os": host_os,
                }
            }
        ]


def hardlink_or_copy_file(src_path: Path, dst_path: Path) -> None:
    """Hardlink or copy a source file to a destination path.

    If the source path is a directory, ValueError is raised.

    If the source path is a symlink, a symlink is created at the destination
    that points to the same location, but using a destination-relative target path.

    If the hardlink fails, a copy operation will be performed.

    Args:
        src_path: Source file path.
        dst_path: Destination path.
    Raise:
        ValueError if src_path is a directory path.
    """
    # Resolve symlink source paths.
    if src_path.is_dir():
        raise ValueError(f"Cannot hard-link or copy directory: {src_path}")

    dst_path.parent.mkdir(parents=True, exist_ok=True)
    if dst_path.is_symlink():
        dst_path.unlink()

    if src_path.is_symlink():
        dst_path.symlink_to(
            os.path.relpath(src_path.resolve(), dst_path.parent)
        )
    else:
        try:
            dst_path.hardlink_to(src_path)
        except Exception:
            shutil.copyfile(src_path, dst_path)


def filter_bazel_bin_path(bazel_path: str) -> str:
    """Filter the bazel-out/<config-dir>/bin/ prefix of an execroot-relative path.

    Args:
        bazel_path: A path relative to the Bazel execroot, which for artifacts
           must begin with `bazel-out/<config_dir>/bin/`.
    Returns:
        The same path without the `bazel-out/<config_dir>/bin/` prefix.
    Raises:
        ValueError if the path is not formed properly.
    """
    segments = bazel_path.split("/")
    if len(segments) <= 3 or segments[0] != "bazel-out" or segments[2] != "bin":
        raise ValueError(
            f"Expected execroot-relative artifact path, got [{bazel_path}]"
        )
    return "/".join(segments[3:])


class HostTestInfo(object):
    """Record information about a given host test, and produce outputs accordingly."""

    def __init__(
        self,
        test_input: dict[str, T.Any],
        host_info: HostInfo,
        export_subdir: str = "",
    ) -> None:
        """Create new instance.

        Args:
            test_input: A dictionary modeling the host_test_info.cquery output
                for the Bazel test target.

            host_info: A HostInfo instance describing the current host system.

            export_subdir: An optional sub-directory path, relative to the export directory.
                The default is an empty string, meaning no sub-directory will be used.
        Raises:
            AssertionError is test_input is malformed.
        """
        self._host_info = host_info
        self._export_prefix = f"{export_subdir}/" if export_subdir else ""
        # LINT.IfChange
        self._input = test_input
        self._label = self._input["label"]
        assert self._label.startswith(
            "@"
        ), f"Invalid Bazel target label does not begin with @: {self._label}"

        # The path to the tests's entry point, relative to the Ninja build directory
        # (or the top of the test archive). Because the test archive does not support
        # symlinks, this will be a shell script that invokes the real test binary
        # with an "exec" call.
        self._entry_path = self._export_prefix + filter_bazel_bin_path(
            self._input["executable"]
        )

        # The path to the file describing the runtime dependencies for the host test.
        # Its content can only be generated after building the test, as its depends on
        # the runfiles manifest's content, which is not available through Bazel cqueries.
        self._runtime_deps_path = f"{self._entry_path}.runtime_deps.json"
        # LINT.ThenChange(../starlark/host_test_info.cquery)

    def generate_tests_json_entry(self) -> dict[str, T.Any]:
        """Generate a single tests.json entry describing the current host test.

        Returns:
        A new JSON dictionary describing the host test, following the schema at
        https://fuchsia.dev/fuchsia-src/reference/testing/tests-json-format
        """
        # For now, use the test path as the name. `fx test` will use the basename
        # to find it.
        name = self._entry_path

        return {
            "environments": self._host_info.test_json_environments,
            "expects_ssh": False,
            "label": self._label,
            "name": name,
            "path": self._entry_path,
            "runtime_deps": self._runtime_deps_path,
            "os": self._host_info.os,
            "cpu": self._host_info.cpu,
        }

    def export_to(self, build_dir: Path, execroot: Path) -> None:
        """Export this host test to a Ninja {build_dir}.

        Mirror the executable and its runfiles support files into the
        Ninja build directory into {build_dir}/{export_subdir}/

        For example, given the following Bazel test files:

        ```
        {EXECROOT}/bazel-out/k8-fastbuild/bin/
              src/
                  foo
                  foo.repo_manifest
                  foo.runfiles_manifest
                  foo.runfiles/
                      _main/
                          src/
                              foo   ---> {EXECROOT}/bazel-out/k8-fastbuild/bin/src/foo
                          data/
                              file  ---> {WORKSPACE}/data/file
                      _repo_mapping ---> {EXECROOT}/bazel-out/k8-fastbuild/bin/src/foo.repo_mapping
                      MANIFEST      ---> {EXECROOT}/bazel-out/k8-fastbuild/bin/src/foo.runfiles_manifest
        ```

        This will create the following in the Ninja {build_dir}:

        ```
        {export_subdir}/
              src/
                  foo                   --> {EXECROOT}/bazel-out/k8-fastbuild/bin/src/foo
                  foo.repo_manifest     --> {EXECROOT}/bazel-out/k8-fastbuild/
                  foo.runfiles_manifest # ADJUSTED MANIFEST (see below)
                  foo.runtime_deps.json
                  foo.runfiles/
                      _main/
                          src/
                              foo   ---> {EXECROOT}/bazel-out/k8-fastbuild/bin/src/foo
                          data/
                              file  ---> {WORKSPACE}/data/file
                      _repo_mapping ---> {EXECROOT}/bazel-out/k8-fastbuild/bin/src/foo.repo_mapping
                      MANIFEST      ---> {build_dir}/{export_subdir}/src/foo.runfiles_manifest
        ```

        The major differences are:

        - The `{export_subdir}/src/foo.runtime_deps.json` file, which will be referenced from
          `tests.json`, which contains a JSON list of paths to all the runtime-required files,
          relative to the Ninja {build_dir}. In this case, this will look like:

          ```
          {export_subdir}/src/foo
          {export_subdir}/src/foo.repo_mapping
          {export_subdir}/src/foo.runfiles_manifest
          {export_subdir}/src/foo.runfiles/_main/src/foo
          {export_subdir}/src/foo.runfiles/_main/data/file
          {export_subdir}/src/foo.runfiles/MANIFEST
          ```

        - The `{export_subdir}/src/foo.runtime_manifest` is adjusted to list target paths that
          are also relative to the Ninja {build_dir}, so will contain something like:

          ```
          _main/src/foo {export_subdir}/src/foo.runfiles/_main/src/foo
          _main/data/file {export_subdir}/src/foo.runfiles/_main/data/file
          _repo_mapping {export_subdir}/src/foo.runfiles/_repo_mapping
          ```

        IMPORTANT: Running the test requires launching it from the Ninja {build_dir}, though
        there are exceptions for py_test() and sh_test() targets, because they use
        a middle-man script that can probe the environment and file-system to adjust
        paths accordingly.

        Args:
            build_dir: Root export directory, a.k.a. Ninja build directory.
            execroot: Bazel execroot absolute path. Used to resolve relative paths
                in runfiles manifest.
        """
        # First, locate the input runfiles manifest, load it, and clean it up a little.
        input_manifest_path = execroot / self._input["runfiles_manifest"]
        assert (
            input_manifest_path.exists()
        ), f"Missing Bazel runfiles manifest: {input_manifest_path}"
        input_manifest = runfiles_utils.RunfilesManifest.CreateFrom(
            input_manifest_path.read_text()
        )
        input_manifest.remove_legacy_external_runfiles(
            workspace_name=_MAIN_WORKSPACE_NAME
        )

        input_runfiles_dir = execroot / f"{self._input['executable']}.runfiles"
        assert (
            input_runfiles_dir.exists()
        ), f"Missing Bazel runfiles directory: {input_runfiles_dir}"

        # Second, locate the _repo_mapping file from it. This should be an absolute path or
        # an execroot-relative one.
        repo_mapping_path = Path(input_manifest.lookup("_repo_mapping"))
        assert (
            repo_mapping_path
        ), f"Missing _repo_mapping entry from runfiles manifest at: {input_manifest_path}"
        if not repo_mapping_path.is_absolute():
            repo_mapping_path = execroot / repo_mapping_path
        assert (
            repo_mapping_path.exists()
        ), f"Missing Bazel repository mapping file: {repo_mapping_path}"

        # For every entry in the manifest, create a corresponding symlink in the output runfiles directory.
        output_runfiles_dir = build_dir / f"{self._entry_path}.runfiles"
        if output_runfiles_dir.exists():
            shutil.rmtree(output_runfiles_dir)
        output_runfiles_dir.mkdir(parents=True)

        runtime_deps = []
        output_manifest_entries: dict[str, str] = {}
        for source_path, target_path in input_manifest.as_dict().items():
            dest_path = output_runfiles_dir / source_path
            runtime_path = os.path.relpath(dest_path, build_dir)

            if not target_path:
                # This is an empty file in the input runfiles dir, create an empty
                # one in the output runfiles dir too.
                dest_path.parent.mkdir(parents=True, exist_ok=True)
                dest_path.write_text("")
            else:
                if not os.path.isabs(target_path):
                    target_path = f"{execroot}/{target_path}"
                build_utils.force_symlink(dest_path, target_path)

            runtime_deps.append(runtime_path)
            output_manifest_entries[source_path] = runtime_path

        # Create the MANIFEST file in the destination runfiles directory.
        exported_manifest = runfiles_utils.RunfilesManifest(
            output_manifest_entries
        )
        output_manifest_path = output_runfiles_dir / "MANIFEST"
        output_manifest_path.write_text(exported_manifest.generate_content())
        runtime_deps.append(os.path.relpath(output_manifest_path, build_dir))

        # Add foo.runfiles_manifest and foo.repo_mapping symlinks as well.
        output_manifest_link_path = (
            build_dir / f"{self._entry_path}.runfiles_manifest"
        )
        build_utils.force_symlink(
            output_manifest_link_path, output_manifest_path
        )
        runtime_deps.append(
            os.path.relpath(output_manifest_link_path, build_dir)
        )

        output_repo_mapping_link_path = (
            build_dir / f"{self._entry_path}.repo_mapping"
        )
        build_utils.force_symlink(
            output_repo_mapping_link_path, output_runfiles_dir / "_repo_mapping"
        )
        runtime_deps.append(
            os.path.relpath(output_repo_mapping_link_path, build_dir)
        )

        # Write the foo.runtime_deps.json file.
        output_runtime_deps_path = build_dir / self._runtime_deps_path
        output_runtime_deps_path.write_text(json.dumps(runtime_deps, indent=2))

        # Finally, hardlink or copy the executable itself. A symlink will not work for
        # Python and shell tests because the executable is a middleman script that may
        # use expressions like |`readlink $0`.runfiles| to find the runfiles directory
        # and this will not work.
        hardlink_or_copy_file(
            execroot / self._input["executable"],
            build_dir / self._entry_path,
        )


@dataclasses.dataclass
class CommandContext(object):
    """Common context for all script commands."""

    args: argparse.Namespace
    build_dir: Path
    tests_map: dict[str, HostTestInfo]
    bazel_cmd: build_utils.BazelCommand


def list_command(args: argparse.Namespace, context: CommandContext) -> int:
    """Implement the `list` script command."""
    # Create a tests.json file corresponding to the test_info entries.
    tests_json = []
    for label, test_info in sorted(context.tests_map.items()):
        tests_json.append(test_info.generate_tests_json_entry())

    content = json.dumps(tests_json, indent=2)
    if args.output:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(content)
    else:
        print(content)

    return 0


def export_command(args: argparse.Namespace, context: CommandContext) -> int:
    """Implement the `export` command."""
    execroot = context.bazel_cmd.get_execroot()

    # Build them, as we require the runfiles manifest file to parse it.
    context.bazel_cmd.run("build", sorted(args.target_pattern))

    build_dir = context.build_dir
    if args.output_dir:
        build_dir = args.output_dir

    for label, test_info in context.tests_map.items():
        test_info.export_to(build_dir, execroot)

    return 0


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--fuchsia-dir",
        type=Path,
        help="Path to Fuchsia checkout (auto-detected).",
    )
    parser.add_argument(
        "--build-dir",
        type=Path,
        help="Path to Fuchsia build directory (auto-detected).",
    )
    parser.add_argument(
        "--export-subdir",
        type=Path,
        default=_DEFAULT_EXPORT_SUBDIR,
        help=f"Set build_dir sub-directory for export (default {_DEFAULT_EXPORT_SUBDIR}).",
    )
    subparsers = parser.add_subparsers(required=True)

    parser_list = subparsers.add_parser(
        "list",
        help="Generate a tests.json file from one or more Bazel target patterns.",
        formatter_class=argparse.RawTextHelpFormatter,
        description=_LIST_TESTS_COMMAND_HELP,
    )
    parser_list.add_argument(
        "--output",
        type=Path,
        help="Write result to file, default goes to stdout.",
    )
    parser_list.add_argument(
        "target_pattern", nargs="+", help="Bazel host test label patterns"
    )
    parser_list.set_defaults(func=list_command)

    parser_export = subparsers.add_parser(
        "export",
        help="Export Bazel test artifacts to the Ninja build directory.",
        description=_EXPORT_COMMAND_HELP,
        formatter_class=argparse.RawTextHelpFormatter,
    )
    parser_export.add_argument(
        "--output-dir",
        type=Path,
        help="Specify alternative output directory, default is Ninja build directory.",
    )
    parser_export.add_argument(
        "target_pattern", nargs="+", help="Bazel host test label patterns"
    )
    parser_export.set_defaults(func=export_command)

    args = parser.parse_args()

    bazel_paths = build_utils.BazelPaths.new(args.fuchsia_dir, args.build_dir)

    bazel_launcher = bazel_paths.launcher
    if not bazel_launcher.exists():
        parser.error(
            f"Could not find Bazel launcher script! fuchsia_dir={fuchsia_dir} build_dir={build_dir}"
        )

    bazel_cmd = build_utils.BazelCommand(bazel_launcher)

    # Now query each target to get its executable and runfiles_manifest path
    # in the output base, then load the runfiles manifest.
    patterns = " ".join(args.target_pattern)
    output = bazel_cmd.run(
        "cquery",
        [
            "--output=starlark",
            f"--starlark:file={_CQUERY_STARLARK_FILE_PATH.resolve()}",
            f"tests(set({patterns}))",
        ],
    )

    host_info = HostInfo()
    tests_map = {}
    for line in output.splitlines():
        test_entry = json.loads(line)
        test_info = HostTestInfo(test_entry, host_info, args.export_subdir)
        tests_map[test_entry["label"]] = test_info

    context = CommandContext(
        args, bazel_paths.ninja_build_dir, tests_map, bazel_cmd
    )
    return args.func(args, context)


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