/*
 * Copyright 2018 Bloomberg Finance LP
 *
 * 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.
 */

#ifndef INCLUDED_BUILDBOXCOMMON_RUNNER
#define INCLUDED_BUILDBOXCOMMON_RUNNER

#include <optional>
#include <signal.h>
#include <unordered_map>

#include <buildboxcommon_assetclient.h>
#include <buildboxcommon_casclient.h>
#include <buildboxcommon_connectionoptions.h>
#include <buildboxcommon_exception.h>
#include <buildboxcommon_logging.h>
#include <buildboxcommon_stageddirectory.h>
#include <buildboxcommon_systemutils.h>
#include <buildboxcommon_temporaryfile.h>
#include <buildboxcommon_timeutils.h>

namespace buildboxcommon {
// To avoid repetition, we define these macros that will augment the messages
// that are logged with the action digest that a `Runner` instance is
// processing.
// NOLINTBEGIN (cppcoreguidelines-macro-usage)
#define BUILDBOX_RUNNER_LOG_MESSAGE_FORMAT(message)                           \
    "[actionDigest=" << this->d_action_digest << "] " << message

#define BUILDBOX_RUNNER_LOG(level, message)                                   \
    {                                                                         \
        BUILDBOX_LOG_##level(BUILDBOX_RUNNER_LOG_MESSAGE_FORMAT(message));    \
    }

// NOLINTEND (cppcoreguidelines-macro-usage)

class Runner {
  public:
    Runner() {}

    /**
     * Execute the given Command in the given input root and return an
     * ActionResult. Subclasses should override this to implement
     * sandboxing behaviors.
     */
    virtual ActionResult execute(const Command &command,
                                 const Digest &inputRootDigest,
                                 const Platform &platform) = 0;

    /**
     * Subclasses can override this to add support for special arguments.
     * Return true if an argument was handled successfully.
     */
    virtual bool parseArg(std::string_view) { return false; }

    /**
     * Subclasses can override this to print a message after Runner prints
     * its usage message.
     */
    virtual void printSpecialUsage() {}

    /**
     * Subclasses can override this to print Runner-specific capabilities.
     * The format is one capability name per line. In the common case where
     * the capability is associated with a CLI option, the printed
     * capability name should match the name of the option.
     */
    virtual void printSpecialCapabilities() {}

    static void handleSignal(int signal);
    static sig_atomic_t getSignalStatus();
    /**
     * Chmod a directory and all subdirectories recursively.
     */
    static void recursively_chmod_directories(const char *path, mode_t mode);

    int main(int argc, char *argv[]);
    virtual ~Runner() {};

    // Delete the copy constructor
    Runner(const Runner &) = delete;

    // Delete the copy assignment operator
    Runner &operator=(const Runner &) = delete;

    // Define the move constructor
    Runner(Runner &&) noexcept = default;

    // Define the move assignment operator
    Runner &operator=(Runner &&) noexcept = default;

    // When a Runner exits early (non-zero exit code), it will attempt to write
    // a `StatusCode` proto that describes the cause of the error next to where
    // the `ActionResult` is expected. This function returns the path where
    // that proto might be found.
    static std::string
    errorStatusCodeFilePath(const std::string &actionResultPath);

    static bool readStatusFile(const std::string &path,
                               google::rpc::Status *status);

  protected:
    /**
     * Execute the given command (without attempting to sandbox it) and
     * store its stdout, stderr, and exit code in the given ActionResult.
     *
     * If `skip_logs_capture` is true, do not capture the contents of stdout
     * and stderr (`stdout_digest` and `stderr_digest` will be left unset in
     * the `ActionResult`.)
     */
    void executeAndStore(const std::vector<std::string> &command,
                         ActionResult *result,
                         const std::string &workdir = ".");

    typedef std::function<grpc::Status(
        const std::string &stdout_file, const std::string &stderr_file,
        Digest *stdout_digest, Digest *stderr_digest)>
        UploadOutputsCallback;
    /**
     * Helper method to unit test the runner-facing implementation.
     * It invokes the `upload_output_function` callback for `stdout` and
     * `stderr` unless the `skip_capture` option is set in this instance's
     * `StandardOutputsCaptureConfig`, in which case the callback is ignored.
     */
    void executeAndStore(const std::vector<std::string> &command,
                         const UploadOutputsCallback &uploadOutputsFunction,
                         ActionResult *result,
                         const std::string &workdir = ".");
    /**
     * Stage the directory with the given digest, to stage_path, and
     * return a StagedDirectory object representing it.
     * If stage_path is empty, will stage directory to tmpdir.
     *
     * If `use_localcas_protocol` is `true` uses `LocalCasStagedDirectory`
     * instead of `FallBackStagedDirectory`.
     */
    std::unique_ptr<StagedDirectory>
    stageDirectory(const Digest &digest,
                   const ProcessCredentialsGetter &processCredentialsGetter =
                       SystemUtils::getProcessCredentials);

    /**
     * Create parent output directories, in staged directory, as specified
     * by command
     *
     * Given an output file or directory, creates all the parent
     * directories leading up to the directory or file. But not including
     * it. The output files and directories should be relative to
     * workingDir. They should also not contain any trailing or leading
     * slashes.
     */
    void createOutputDirectories(const Command &command,
                                 const std::string &workingDir) const;

    // NOLINTBEGIN (cppcoreguidelines-non-private-member-variables-in-classes)
    std::shared_ptr<CASClient> d_casClient;
    std::shared_ptr<AssetClient> d_assetClient;

    bool d_use_localcas_protocol = true; // Use LocalCAS by default.
    std::string d_stage_path = "";
    StageTreeRequest_StagingMode d_stagingMode =
        StageTreeRequest_StagingMode_DEFAULT;

    Digest d_action_digest;
    std::chrono::seconds d_timeout = std::chrono::seconds::zero();

    // If requested, fill an `execution_stats.proto` message with metrics
    // from the command's execution, upload it to CAS, and attach its digest to
    // `auxiliary_metadata` in the `ActionResult`.
    bool d_collect_execution_stats = false;
    // NOLINTEND (cppcoreguidelines-non-private-member-variables-in-classes)

    // Helpers to set the timestamps of the different stages that are
    // carried out in the runner. Those will be returned in the
    // `ActionResult`, as part of its `ExecutedActionMetadata` field. The
    // helpers also update the partial execution metadata file so the
    // timestamps can be reported back by the worker to buildgrid.
    inline void
    metadata_mark_input_download_start(ExecutedActionMetadata *metadata)
    {
        set_timestamp_to_now(metadata->mutable_input_fetch_start_timestamp());
        updatePartialExecutionMetadataFile(metadata);
    }

    inline void
    metadata_mark_input_download_end(ExecutedActionMetadata *metadata)
    {
        set_timestamp_to_now(
            metadata->mutable_input_fetch_completed_timestamp());
        updatePartialExecutionMetadataFile(metadata);
    }

    inline void
    metadata_mark_output_upload_start(ExecutedActionMetadata *metadata)
    {
        set_timestamp_to_now(
            metadata->mutable_output_upload_start_timestamp());
        updatePartialExecutionMetadataFile(metadata);
    }

    inline void
    metadata_mark_output_upload_end(ExecutedActionMetadata *metadata)
    {
        set_timestamp_to_now(
            metadata->mutable_output_upload_completed_timestamp());
        updatePartialExecutionMetadataFile(metadata);
    }

    inline void
    updatePartialExecutionMetadataFile(ExecutedActionMetadata *metadata)
    {
        if (!d_partialExecutionMetadataFilePath.has_value() ||
            d_partialExecutionMetadataFilePath == "") {
            BUILDBOX_LOG_DEBUG(
                "No file path specified to write partial execution metadata.");
            return;
        }

        BUILDBOX_LOG_DEBUG("Writing partial execution metadata to path: "
                           << d_partialExecutionMetadataFilePath.value());

        buildboxcommon::FileUtils::writeFileAtomically(
            d_partialExecutionMetadataFilePath.value().c_str(),
            metadata->SerializeAsString());
    }

    void writeErrorStatusFile(const google::protobuf::int32 errorCode,
                              const std::string &errorMessage) const;

    void writeStatusFile(const google::rpc::Status &status,
                         const std::string &path) const;

    // Allow derived classes the ability to set output path for testing
    void setOutputPath(const std::string &path)
    {
        if (!this->d_outputPath.empty()) {
            BUILDBOXCOMMON_THROW_EXCEPTION(std::runtime_error,
                                           "output path already set to \"" +
                                               d_outputPath + "\"");
        }
        this->d_outputPath = path;
    }

    struct StandardOutputsCaptureConfig {
        // If not empty, redirect the command's standard output to that file.
        std::string stdout_file_path;
        std::string stderr_file_path;

        // If set, skips capturing and uploading the outputs written by the
        // command to stdout and stderr.
        bool skip_capture = false;
    };

    StandardOutputsCaptureConfig
        d_standardOutputsCaptureConfig; // NOLINT
                                        // (cppcoreguidelines-non-private-member-variables-in-classes)

    /**
     * These helpers `exit(1)` on failures.
     */
    void registerSignals() const;
    void resetSignals() const;

  private:
    /**
     * Attempt to parse all of the given arguments and update this object
     * to reflect them. If an argument is invalid or missing, return false.
     * Otherwise, return true.
     */
    bool parseArguments(int argc, char *argv[]);

    // If set, `main()` will exit immediately after calling  `parseArguments()`
    // with an exit code that represents its result.
    bool d_validateParametersAndExit = false;

    // Commands to run after action is done
    std::vector<std::string> d_postActionCommands;

    /**
     * Upload the contents of `stdout` and `stderr` and return a grpc::Status.
     * The digests of the uploaded files are stored in the provided Digest
     * objects.
     */
    grpc::Status uploadOutputs(const std::string &stdout_file,
                               const std::string &stderr_file,
                               Digest *stdout_digest,
                               Digest *stderr_digest) const;

    ConnectionOptions d_casRemote;
    std::unique_ptr<ConnectionOptions> d_assetRemote;

    std::string d_inputPath;
    std::string d_outputPath;

    // File path to write partial execution metadata to during execution.
    std::optional<std::string> d_partialExecutionMetadataFilePath =
        std::nullopt;

    // When capturing outputs, skip caching them in local cache
    bool d_outputBypassLocalCache = false;

    static volatile sig_atomic_t d_signal_status;

    LogLevel d_logLevel = LogLevel::INFO;

    std::optional<std::string> d_nestedInstanceName;
    std::string d_remoteApisSocketPath;
    bool d_nestedActionCacheUpdateEnabled = false;

    std::optional<std::chrono::seconds> d_graceful_shutdown_seconds =
        std::nullopt;

    /**
     * These helpers `exit(1)` on failures.
     */
    Action readAction(const std::string &path) const;
    void initializeCasClient(const std::string &name,
                             const std::string &actionId,
                             const std::string &toolInvocationId,
                             const std::string &correlatedInvocationId);
    void initializeAssetClient(const std::string &name,
                               const std::string &actionId,
                               const std::string &toolInvocationId,
                               const std::string &correlatedInvocationId);

    void writeActionResult(const ActionResult &action_result,
                           const std::string &path) const;

    // Fetch a `Command` message from the remote CAS. If that fails, log
    // the error and `exit(1)`.
    Command fetchCommand(const Digest &command_digest) const;

    // Populates a `protobuf::Timestamp` with the current time.
    inline static void set_timestamp_to_now(google::protobuf::Timestamp *t)
    {
        t->CopyFrom(buildboxcommon::TimeUtils::now());
    }

    // Collect `ExecutionStatistics` from the child process and return its
    // digest. If a CAS connection is configured (not the case in unit tests),
    // upload the proto message (wrapped inside an `Any` message) to CAS and
    // return its digest. On errors returns an empty `Digest`.
    Digest collectExecutionStatsAndUploadIfNeeded();

    void maskSignals(int how) const;

    // The environment passed to the executed command. If it is
    // `std::nullopt_t`, the whole `environ` from the parent process is passed.
    virtual inline std::optional<std::unordered_map<std::string, std::string>>
    commandEnvironment()
    {
        return std::nullopt;
    }
};
// NOLINTBEGIN (cppcoreguidelines-macro-usage)
#define BUILDBOX_RUNNER_MAIN(x)                                               \
    int main(int argc, char *argv[]) { return x().main(argc, argv); }
// NOLINTEND (cppcoreguidelines-macro-usage)

} // namespace buildboxcommon

#endif
