# Copyright 2022 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.
"""CI recipe for Bazel-SDK-based projects in the Fuchsia ecosystem.

TODO(fxbug.dev/112403): Actually make this recipe Bazel-specific. We will move
incrementally to make diffs with build_test_upload.py more apparent in code
review.

This recipe supports projects that need the following steps in their CI jobs:

1. Checkout code.
2. Invoke well-known Bazel bootstrapping script included in the Bazel SDK.
3. Build with Bazel.
4. Upload build artifacts (images, blobs, drivers, etc.).
5. Run tests in parallel shards on separate machines.
6. Trigger tests in downstream projects, passing through the uploaded build
   artifacts.

## Entrypoint commands

The recipe interacts with each project via a set of high-level entrypoint
commands (often shell scripts), configured via properties. Interacting with each
project only via these entrypoints has several benefits:

- Prevents project-specific implementation details from leaking into the recipe.
  This keeps the recipe simple, and minimizes the need for soft transitions by
  allowing project owners to make changes within the project without changing
  the infrastructure.
- Ensures that continuous integration workflows can be easily run locally just
  by running the same commands as the recipe.
- Enforces that each project's continuous integration workflow follows a
  reasonable process that's consistent with Fuchsia's CI guidelines:
  https://fuchsia.dev/fuchsia-src/contribute/governance/rfcs/0148_ci_guidelines

Some entrypoint commands run actions themselves. Other commands generate output
that instructs the recipe what actions to take, for infrastructure-specific
steps that are needed across many projects (e.g. `gcs_manifest_command`
generates a list of files to upload to GCS).

### Variable substitution

Some of the project entrypoints commands need to take inputs from the recipe for
values that can only be known at runtime. For example, a project's build command
generally requires the absolute path to an infrastructure-created build
directory as an argument so it knows where to place build outputs.

For any input parameters that must come from the recipe, the recipe supports
substitution of a small predefined set of variables used in the entrypoint
configuration properties. For example, a builder for a given project might set
its `gcs_manifest_command` property to something like
["scripts/emit_gcs_manifest.sh", "--build-dir", "{build_dir}"], where
"{build_dir}" is a magic string that will be replaced with the actual path to
the build directory when the recipe runs the command.

An alternative to variable substitution would be to require each entrypoint
executable to implement a certain set of flags. However, this would make each
project more coupled to the infrastructure, since the flag names would need to
match what the recipe uses, and each entrypoint's command-line argument parser
must support the flag format that the recipe expects. So the recipe should not
assume that any entrypoint commands support any specific flags or flag formats.
"""

from google.protobuf import json_format as jsonpb

from PB.recipes.fuchsia.bazel_build_test_upload import BuilderManifest, InputProperties
from PB.infra.build_test_upload.upload_manifest import CIPDUploadManifest
from PB.infra.build_test_upload.test_manifest import TestManifest

from RECIPE_MODULES.fuchsia.swarming_retry import api as swarming_retry_api

DEPS = [
    "fuchsia/artifacts",
    "fuchsia/buildbucket_util",
    "fuchsia/cas_util",
    "fuchsia/checkout",
    "fuchsia/cipd_util",
    "fuchsia/dpi",
    "fuchsia/fxt",
    "fuchsia/git",
    "fuchsia/git_checkout",
    "fuchsia/swarming_retry",
    "recipe_engine/context",
    "recipe_engine/file",
    "recipe_engine/json",
    "recipe_engine/path",
    "recipe_engine/properties",
    "recipe_engine/proto",
    "recipe_engine/raw_io",
    "recipe_engine/step",
    "recipe_engine/swarming",
]

PROPERTIES = InputProperties


def RunSteps(api, props):
    # Projects should not assume they'll always be checked out at the same
    # absolute path, so do the checkout in a random temporary directory.
    checkout_dir = api.path.mkdtemp("checkout")

    if props.jiri_manifest:
        api.checkout.with_options(
            path=checkout_dir,
            manifest=props.jiri_manifest,
            remote=props.remote,
            project=props.jiri_project,
        )
        with api.context(cwd=checkout_dir):
            git_revision = api.git.rev_parse(
                "HEAD",
                step_name="resolve HEAD",
                step_test_data=lambda: api.raw_io.test_api.stream_output_text("abc123"),
            )
    else:
        _, git_revision = api.git_checkout(props.remote, path=checkout_dir)

    upload_namespace = api.buildbucket_util.id

    # Any of these variables can be used as a format string in a command's
    # configured arguments and it will be substituted for the actual value. For
    # example, "--build-dir={build_dir}" in a command's args will be
    # replaced with "--build-dir=/actual/path/to/build/dir".
    build_dir = api.path.mkdtemp("build")
    output_user_root_dir = build_dir.join("user_root")
    variables = {
        "build_dir": build_dir,
        "upload_namespace": upload_namespace,
    }

    def run_command(step_name, command_pb, **kwargs):
        cmd = [checkout_dir.join(*command_pb.path.split("/"))]
        for arg in command_pb.args:
            # Perform variable substitutions.
            cmd.append(arg.format(**variables))
        return api.step(step_name, cmd, **kwargs)

    # TODO(fxbug.dev/101594): These should be configured higher in the
    # stack in one of {swarming_bot, bbagent, recipe_bootstrap}. Set them
    # here for now to unblock isolation work (b/234060366).
    xdg_env_vars = [
        "HOME",
        "XDG_CACHE_HOME",
        "XDG_CONFIG_HOME",
        "XDG_DATA_HOME",
        "XDG_HOME",
        "XDG_STATE_HOME",
    ]
    env = api.context.env
    for env_var in xdg_env_vars:
        env[env_var] = env.get("TMPDIR", str(api.path["cleanup"]))
    with api.context(cwd=checkout_dir, env=env):
        # Invoke well-known bootstrapping script which should install a pinned
        # version of Bazel from CIPD at tools/bazel.
        api.step(
            "bootstrap bazel",
            [checkout_dir.join("scripts", "bootstrap.sh")],
        )
        bazel_path = checkout_dir.join("tools", "bazel")
        # TODO(fxbug.dev/112403): Explicitly fetch dependencies first rather
        # than implicitly depending on `bazel build` to fetch dependencies.
        # Then, run `bazel build` in a sandbox.
        api.step(
            "bazel build",
            [
                bazel_path,
                "--output_base",
                build_dir,
                "--output_user_root",
                output_user_root_dir,
                "build",
                *props.build_args,
            ],
        )

        builder_manifest = BuilderManifest()
        # TODO(fxbug.dev/112403): Make this unconditionally true once every repo
        # has migrated to fuchsia_builder_group().
        if props.query_builder_manifest:
            # The builder manifest is a file generated by the build that
            # specifies test and upload requests. The following query is a
            # stable query that will discover the manifest path.
            query_step = api.step(
                "query builder manifest path",
                [
                    bazel_path,
                    "cquery",
                    *props.build_args,
                    "--starlark:expr",
                    "providers(target).get('@fuchsia_infra//infra/private:providers.bzl%FuchsiaBuilderGroupInfo').manifest.path",
                    "--output",
                    "starlark",
                ],
                stdout=api.raw_io.output_text(),
                step_test_data=lambda: api.raw_io.test_api.stream_output_text(
                    "builder_manifest.json",
                ),
            )
            query_step.presentation.logs["stdout"] = query_step.stdout
            builder_manifest_path = checkout_dir.join(query_step.stdout.strip())
            api.path.mock_add_paths(builder_manifest_path)
            builder_manifest = api.file.read_proto(
                "read builder manifest",
                builder_manifest_path,
                BuilderManifest,
                codec="JSONPB",
                # TODO(fxbug.dev/112403): Supply test data.
                test_proto=BuilderManifest(),
            )

        test_manifests = [
            api.file.read_proto(
                "read test manifest %s" % test_manifest_path,
                checkout_dir.join(test_manifest_path),
                TestManifest,
                codec="JSONPB",
            )
            for test_manifest_path in builder_manifest.test_manifests
        ]
        # TODO(fxbug.dev/112403): Delete in favor of `test_manifests`.
        if props.test_manifest_command.path:
            assert (
                props.testing_pool
            ), "testing_pool must be set if test_manifest_command is set"
            test_manifest = run_command(
                "generate test manifest",
                props.test_manifest_command,
                stdout=api.proto.output(TestManifest, codec="JSONPB"),
            ).stdout
            test_manifests.append(test_manifest)
        # TODO(fxbug.dev/112403): Probably should not use a loop here.
        for test_manifest in test_manifests:
            with api.step.nest("upload test dependencies"):
                task_requests = create_task_requests(
                    api, test_manifest, testing_pool=props.testing_pool
                )
            run_tests(api, task_requests)

        gcs_manifests = [
            api.file.read_json(
                "read gcs manifest %s" % gcs_manifest_path,
                checkout_dir.join(gcs_manifest_path),
            )
            for gcs_manifest_path in builder_manifest.gcs_manifests
        ]
        # TODO(fxbug.dev/112403): Delete in favor of `gcs_manifests`.
        if props.gcs_manifest_command.path:
            with api.step.nest("gcs upload"):
                gcs_manifest = run_command(
                    "generate gcs manifest",
                    props.gcs_manifest_command,
                    stdout=api.json.output(),
                ).stdout
                gcs_manifests.append(gcs_manifest)
        api.artifacts.gcs_bucket = props.gcs_bucket
        api.artifacts.namespace = upload_namespace
        # TODO(fxbug.dev/112403): There should probably only be one GCS manifest
        # allowed.
        for gcs_manifest in gcs_manifests:
            api.artifacts.upload_from_manifest(
                "upload from manifest",
                api.json.input(gcs_manifest),
                sign_artifacts=props.sign_artifacts,
            )

        cipd_manifests = [
            api.file.read_proto(
                "read cipd manifest %s" % cipd_manifest_path,
                checkout_dir.join(cipd_manifest_path),
                CIPDUploadManifest,
                codec="JSONPB",
            )
            for cipd_manifest_path in builder_manifest.cipd_manifests
        ]
        # TODO(fxbug.dev/112403): Delete in favor of `cipd_manifests`. This can
        # happen once the CIPD manifests include the `pkg_name` field.
        for cipd_cmd in props.cipd_manifest_commands:
            cipd_manifest = run_command(
                "generate cipd manifest %s" % cipd_cmd.package,
                cipd_cmd.command,
                stdout=api.proto.output(CIPDUploadManifest, codec="JSONPB"),
            ).stdout
            # TODO(fxbug.dev/112403): Delete this line. Necessary for
            # soft-transition.
            cipd_manifest.pkg_name = cipd_cmd.package
            cipd_manifests.append(cipd_manifest)

        for cipd_manifest in cipd_manifests:
            with api.step.nest("cipd upload %s" % cipd_manifest.pkg_name):
                cipd_upload_from_manifest(
                    api,
                    cipd_manifest.pkg_name,
                    cipd_manifest,
                    build_dir=build_dir,
                    repository=props.remote,
                    git_revision=git_revision,
                )

    if props.mos_upload_options.repo_hostname:
        api.dpi.upload(
            "upload to MOS-TUF repos",
            build_dir=build_dir,
            options=props.mos_upload_options,
        )

    if props.gcs_bucket and props.HasField("fxt_options"):
        api.fxt.orchestrate_fxt_tests(
            bucket=props.gcs_bucket,
            namespace=upload_namespace,
            options=props.fxt_options,
        )


def dirnames(api, path):
    """Lists all parent directories of a path."""
    dirname = api.path.dirname(path)
    # When dirname() becomes a no-op it means we've reached the root.
    if path == dirname:
        return
    yield dirname
    yield from dirnames(api, dirname)


def create_task_requests(api, test_manifest, testing_pool):
    root_dir = api.path.abs_to_path(test_manifest.root_dir)

    task_requests = []
    for test_group in test_manifest.test_groups:
        original_runfiles = set(test_group.runfiles)
        # TODO(crbug.com/1216363): Once the `cas` tool is capable of following
        # symlinks we won't need to reconstruct the runfiles in a temporary
        # directory and can instead upload the runfiles directly from
        # `root_dir`.
        tree = api.cas_util.hardlink_tree(api.path.mkdtemp("cas-inputs"))
        for f in test_group.runfiles:
            # Ignore a file if we're also linking in the full contents of one of
            # its parent directories.
            if any(d in original_runfiles for d in dirnames(api, f)):
                continue
            tree.register_link(root_dir.join(f), linkname=tree.root.join(f))
        tree.create_links("create links")

        cas_digest = api.cas_util.upload(
            tree.root,
            step_name="upload files for %s" % test_group.name,
        )

        request = api.swarming.task_request().with_name(test_group.name)
        request = request.with_slice(
            0,
            request[0]
            .with_command(list(test_group.command))
            .with_relative_cwd(test_group.exec_dir)
            .with_cas_input_root(cas_digest)
            .with_dimensions(pool=testing_pool, **test_group.scheduling_dimensions)
            .with_execution_timeout_secs(test_group.timeout_secs)
            .with_io_timeout_secs(5 * 60)
            .with_expiration_secs(60 * 60),
        )
        task_requests.append(request)

    return task_requests


def run_tests(api, task_requests):
    tasks = [Task(request, api) for request in task_requests]
    api.swarming_retry.run_and_present_tasks(tasks, max_attempts=1)


class Task(swarming_retry_api.TriggeredTask):
    pass


def cipd_upload_from_manifest(
    api, cipd_package, cipd_manifest, build_dir, repository, git_revision
):
    staging_dir = api.path.mkdtemp("cipd")
    tree = api.cas_util.hardlink_tree(staging_dir)

    for f in cipd_manifest.files:
        # We should generally not be uploading artifacts from outside the build
        # directory, in order to ensure everything flows through the
        # checkout->build->upload pipeline. So disallow uploading files from
        # outside the build directory.
        #
        # Links to files outside the build directory are still allowed.
        if ".." in f.source.split("/"):
            raise api.step.StepFailure(
                "CIPD upload file source must within the build directory: %s" % f.source
            )
        abs_source = build_dir.join(*f.source.split("/"))

        # For convenience, projects can specify dest="." to upload an entire
        # directory as the contents of the package.
        if f.dest == ".":
            if len(cipd_manifest.files) > 1:  # pragma: no cover
                raise api.step.StepFailure(
                    "Only one CIPD manifest entry is allowed if any entry's destination is '.'"
                )
            # No need for a tree of symlinks anymore, we can just treat the
            # source directory as the staging directory.
            staging_dir = abs_source
            tree = None
            break

        abs_dest = tree.root.join(*f.dest.split("/"))
        tree.register_link(abs_source, linkname=abs_dest)

    if tree:
        tree.create_links("create hardlinks")

    api.cipd_util.upload_package(
        cipd_package,
        staging_dir,
        search_tag={"git_revision": git_revision},
        repository=repository,
    )


def GenTests(api):
    default_remote = "https://fuchsia.googlesource.com/foo"

    def properties(**kwargs):
        props = {
            "remote": default_remote,
            "build_args": ["//src:buildbot"],
            "query_builder_manifest": True,
        }
        props.update(kwargs)
        return api.properties(jsonpb.ParseDict(props, InputProperties()))

    def cipd_manifest_data(files=None):
        if not files:
            files = {"local/foo/bar": "package/foobar"}
        return api.proto.output(
            CIPDUploadManifest(
                files=[
                    CIPDUploadManifest.FileToUpload(source=k, dest=v)
                    for k, v in files.items()
                ]
            )
        )

    yield (
        api.buildbucket_util.test("fxt_tests", git_repo=default_remote)
        + properties(
            gcs_manifest_command={"path": "emit_gcs_manifest.sh"},
            gcs_bucket="foo-artifacts",
            sign_artifacts=True,
            fxt_options={"image_name": "image-name"},
        )
        + api.step_data(
            "gcs upload.generate gcs manifest",
            stdout=api.json.output([{"source": "foo.txt", "destination": "foo.txt"}]),
        )
        + api.fxt.orchestrate_fxt_tests()
    )

    yield (
        api.buildbucket_util.test("jiri__mos__cipd", git_repo=default_remote)
        + properties(
            jiri_manifest="path/to/manifest",
            mos_upload_options={
                "repo_hostname": "test.fuchsia-update.googleusercontent.com",
                "gcs_bucket": "discover-cloud.appspot.com",
                "manifest_path": "path/to/mos/manifest",
            },
            cipd_manifest_commands=[
                {
                    "command": {
                        "path": "scripts/emit_foo_cipd_manifest.sh",
                        "args": ["-o", "{build_dir}"],
                    },
                    "package": "fuchsia/tools/foo",
                }
            ],
        )
        + api.step_data(
            "generate cipd manifest fuchsia/tools/foo",
            stdout=cipd_manifest_data(
                # Cover the case where we upload an entire subdirectory into the
                # root of the CIPD package.
                {"dir/pkg": "."},
            ),
        )
    )

    yield (
        api.buildbucket_util.test(
            "invalid_cipd_upload",
            git_repo=default_remote,
            # The manifest should be validated even in presubmit.
            tryjob=True,
            status="failure",
        )
        + properties(
            cipd_manifest_commands=[
                {
                    "command": {"path": "scripts/emit_cipd_manifest.sh"},
                    "package": "fuchsia/tools/foo",
                }
            ],
        )
        + api.step_data(
            "generate cipd manifest fuchsia/tools/foo",
            stdout=cipd_manifest_data(
                # Paths outside the build directory are disallowed.
                {"foo/../../checkout-path": "package_path/foo/bar"}
            ),
        )
    )

    yield (
        api.buildbucket_util.test(
            "cipd_upload_tryjob", git_repo=default_remote, tryjob=True
        )
        + properties(
            cipd_manifest_commands=[
                {
                    "command": {"path": "scripts/emit_cipd_manifest.sh"},
                    "package": "fuchsia/tools/foo",
                }
            ],
        )
        + api.step_data(
            "generate cipd manifest fuchsia/tools/foo",
            stdout=cipd_manifest_data(),
        )
    )

    yield (
        api.buildbucket_util.test(
            "failed_tests", git_repo=default_remote, status="failure"
        )
        + properties(
            testing_pool="testpool",
            test_manifest_command={"path": "scripts/emit_test_manifest.sh"},
        )
        + api.step_data(
            "generate test manifest",
            stdout=api.proto.output(
                TestManifest(
                    root_dir="[START_DIR]/foo",
                    test_groups=[
                        TestManifest.TestGroup(
                            name="group1",
                            command=["run-tests", "--verbose"],
                            exec_dir="path/to/exec/dir",
                            scheduling_dimensions=[("foo", "bar")],
                            # The first file should be ignored since its entire
                            # directory is included.
                            runfiles=["some/runtime/dep", "some/runtime"],
                        )
                    ],
                )
            ),
        )
        + api.swarming_retry.collect_data(
            [api.swarming_retry.failed_task(name="group1", task_id=0)]
        )
    )