/*
 * Copyright 2022 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.
 */

#include <exception>
#include <fstream>
#include <google/protobuf/util/json_util.h>
#include <google/rpc/code.pb.h>
#include <iostream>
#include <memory>
#include <optional>
#include <ostream>
#include <stdexcept>
#include <ThreadPool.h>
#include <utility>

#include <buildboxcommon_fileutils.h>
#include <buildboxcommon_grpcclient.h>
#include <buildboxcommon_grpcerror.h>
#include <buildboxcommon_localexecutionclient.h>
#include <buildboxcommon_logging.h>
#include <buildboxcommon_protojsonutils.h>
#include <buildboxcommon_protos.h>
#include <buildboxcommon_remoteexecutionclient.h>
#include <buildboxcommon_stringutils.h>
#include <buildboxcommon_timeutils.h>
#include <trexe_executioncontext.h>

namespace trexe {

ExecutionContext::ExecutionContext(
    const std::shared_ptr<const ExecutionOptions> &options,
    const ConnectionOptions &casConnectionOptions,
    const ConnectionOptions &executionConnectionOptions,
    const ConnectionOptions &actioncacheConnectionOptions,
    const std::shared_ptr<const ConnectionOptions> &logstreamConnectionOptions)
    : d_executionOptions(options)
{
    std::string toolInvocationId = "";
    if (options->d_metadata.count("tool-invocation-id")) {
        toolInvocationId = options->d_metadata.at("tool-invocation-id");
    }

    std::string correlatedInvocationsId = "";
    if (options->d_metadata.count("correlated-invocations-id")) {
        correlatedInvocationsId =
            options->d_metadata.at("correlated-invocations-id");
    }

    std::string actionMnemonic = "";
    if (options->d_metadata.count("action-mnemonic")) {
        actionMnemonic = options->d_metadata.at("action-mnemonic");
    }

    std::string targetId = "";
    if (options->d_metadata.count("target-id")) {
        targetId = options->d_metadata.at("target-id");
    }

    std::string configurationId = "";
    if (options->d_metadata.count("configuration-id")) {
        configurationId = options->d_metadata.at("configuration-id");
    }

    // Initialize CAS Client. This client does not yet have action_id,
    // as action_id needs a CAS lookup to calculate.
    auto grpcClient = std::make_shared<GrpcClient>();
    {
        grpcClient->init(casConnectionOptions);
        if (options->d_metadata.count("tool-name") &&
            options->d_metadata.count("tool-version")) {
            grpcClient->setToolDetails(options->d_metadata.at("tool-name"),
                                       options->d_metadata.at("tool-version"));
        }

        grpcClient->setRequestMetadata("", toolInvocationId,
                                       correlatedInvocationsId, actionMnemonic,
                                       targetId, configurationId);

        d_casClient = std::make_shared<CASClient>(grpcClient);
        d_casClient->init();
    }
    std::shared_ptr<Digest> inputRootDigest;
    // Validate input-root-digest if necessary
    {
        if (!options->d_inputRootDigest.empty()) {
            inputRootDigest = std::make_shared<Digest>();
            if (!StringUtils::parseDigest(options->d_inputRootDigest,
                                          inputRootDigest.get())) {
                BUILDBOXCOMMON_THROW_EXCEPTION(
                    std::invalid_argument, "Invalid input root digest: "
                                               << options->d_inputRootDigest);
            }
            try {
                d_casClient->fetchMessage<Directory>(*inputRootDigest);
            }
            catch (const std::exception &) {
                BUILDBOX_LOG_ERROR("Cannot validate that input-root-digest "
                                   "points to a Directory");
                throw;
            }
        }
    }
    // Generate Action
    {
        if (options->d_operation.empty()) {
            setActionData(buildAction(
                d_casClient, options->d_argv, options->d_workingDir,
                options->d_inputPaths, inputRootDigest, options->d_outputPaths,
                options->d_platform, options->d_environment,
                options->d_execTimeout, options->d_doNotCache,
                options->d_followSymlinks, options->d_salt,
                options->d_outputNodeProperties, options->d_numDigestThreads));
        }
    }
    // Update CAS client now that action_id is available
    {
        if (options->d_operation.empty() &&
            options->d_metadata.count("tool-invocation-id") &&
            options->d_metadata.count("correlated-invocations-id")) {
            grpcClient->setRequestMetadata(
                buildboxcommon::toString(d_actionData.d_actionDigest),
                toolInvocationId, correlatedInvocationsId, actionMnemonic,
                targetId, configurationId);
        }
    }
    // Initialize Execution Client
    {
        std::shared_ptr<GrpcClient> actioncacheGrpcClient;
        if (!actioncacheConnectionOptions.d_url.empty()) {
            actioncacheGrpcClient = std::make_shared<GrpcClient>();
            actioncacheGrpcClient->init(actioncacheConnectionOptions);
            if (options->d_metadata.count("tool-name") &&
                options->d_metadata.count("tool-version")) {
                actioncacheGrpcClient->setToolDetails(
                    options->d_metadata.at("tool-name"),
                    options->d_metadata.at("tool-version"));
            }
            if (options->d_metadata.count("tool-invocation-id") &&
                options->d_metadata.count("correlated-invocations-id")) {
                actioncacheGrpcClient->setRequestMetadata(
                    !options->d_operation.empty()
                        ? ""
                        : buildboxcommon::toString(
                              d_actionData.d_actionDigest),
                    toolInvocationId, correlatedInvocationsId, actionMnemonic,
                    targetId, configurationId);
            }
        }

        if (options->d_cacheOnly) {
            auto localExecutionClient = std::make_shared<LocalExecutionClient>(
                casConnectionOptions, actioncacheGrpcClient);
            localExecutionClient->setRunner(options->d_runnerCommand,
                                            options->d_extraRunArgs);
            d_executionClient = localExecutionClient;
        }
        else {
            auto executionGrpcClient = std::make_shared<GrpcClient>();
            executionGrpcClient->init(executionConnectionOptions);
            if (options->d_metadata.count("tool-name") &&
                options->d_metadata.count("tool-version")) {
                executionGrpcClient->setToolDetails(
                    options->d_metadata.at("tool-name"),
                    options->d_metadata.at("tool-version"));
            }
            if (options->d_metadata.count("tool-invocation-id") &&
                options->d_metadata.count("correlated-invocations-id")) {
                executionGrpcClient->setRequestMetadata(
                    !options->d_operation.empty()
                        ? ""
                        : buildboxcommon::toString(
                              d_actionData.d_actionDigest),
                    toolInvocationId, correlatedInvocationsId, actionMnemonic,
                    targetId, configurationId);
            }
            if (options->d_schedulingMetadata.ByteSizeLong() > 0) {
                executionGrpcClient->setSchedulingMetadata(
                    options->d_schedulingMetadata);
            }

            d_executionClient = std::make_shared<RemoteExecutionClient>(
                executionGrpcClient, actioncacheGrpcClient,
                logstreamConnectionOptions);

            if (options->d_streamLogs && options->d_operation.empty()) {
                this->enableLogStream(d_actionData.d_actionDigest);
            }
            if (options->d_logProgress) {
                std::dynamic_pointer_cast<RemoteExecutionClient>(
                    d_executionClient)
                    ->enableProgressLog();
            }
        }

        d_executionClient->init();
    }

    if (options->d_numIOThreads > 0) {
        d_ioThreadPool =
            std::make_unique<ThreadPool>(options->d_numIOThreads, "trexe.io");
    }
}

ExecutionContext::ExecutionContext(
    const std::shared_ptr<const ExecutionOptions> &options,
    const std::shared_ptr<CASClient> &casClient,
    const std::shared_ptr<ExecutionClient> &execClient)
    : d_executionOptions(options), d_casClient(casClient),
      d_executionClient(execClient)
{
    d_casClient->init();
    d_executionClient->init();
}

ExecutionContext::~ExecutionContext()
{
    try {
        // When the ExecutionContext is out of the scope,
        // it writes the metadata file if user specifies the path
        if (!d_executionOptions->d_resultMetadataFile.empty()) {
            BUILDBOX_LOG_DEBUG("Writing result metadata to: "
                               << d_executionOptions->d_resultMetadataFile);
            writeProtoInJSONToFile(d_executionOptions->d_resultMetadataFile,
                                   d_resultMetadata, false);
        }
    }
    catch (const std::exception &exception) {
        BUILDBOX_LOG_ERROR("Caught exception while writing result metadata: "
                           << exception.what());
    }
}

void ExecutionContext::enableLogStream(const Digest &actionDigest)
{
    if (!d_executionOptions->d_cacheOnly) {
        // Downcasting here is safe because `d_executionClient` is
        // guaranteed to be a `RemoteExecutionClient` in this remote
        // build code branch.
        std::dynamic_pointer_cast<RemoteExecutionClient>(d_executionClient)
            ->enableLogStream(actionDigest, d_executionOptions->d_stdoutFile,
                              d_executionOptions->d_stderrFile);
        std::dynamic_pointer_cast<RemoteExecutionClient>(d_executionClient)
            ->enableLogStream(actionDigest, &std::cout, &std::cerr);
    }
}

void ExecutionContext::setActionData(const ActionData &actionData)
{
    d_actionData = actionData;
    *d_resultMetadata.mutable_action_digest() = actionData.d_actionDigest;
}

bool ExecutionContext::skipsCacheLookup() const
{
    return d_executionOptions->d_skipCacheLookup;
}

bool ExecutionContext::isResultCached(bool allowRemoteQuery)
{
    return (getActionResult(allowRemoteQuery) != nullptr);
}

bool ExecutionContext::downloadResults(const ActionResult &ar,
                                       const std::string &where)
{
    if (where.size() == 0) {
        BUILDBOXCOMMON_THROW_EXCEPTION(std::runtime_error,
                                       "Download location is empty.");
    }

    buildboxcommon::FileUtils::createDirectory(where.c_str());
    buildboxcommon::FileDescriptor dirfd(
        open(where.c_str(), O_RDONLY | O_DIRECTORY | O_CLOEXEC));
    if (dirfd.get() < 0) {
        BUILDBOXCOMMON_THROW_EXCEPTION(std::runtime_error,
                                       "Download location "
                                           << where << " can't be opened.");
    }

    BUILDBOX_LOG_DEBUG("Output directory ready: " << where);

    try {

        d_executionClient->downloadOutputs(d_casClient.get(), ar, dirfd.get(),
                                           d_ioThreadPool.get());
    }
    catch (buildboxcommon::GrpcError &grpcError) {
        updateGrpcError(grpcError.status);
        throw;
    }

    BUILDBOX_LOG_DEBUG("Results downloaded!");
    return true;
}

/**
 * Return the ActionResult for the given Operation. Throws an exception
 * if the Operation finished with an error, or if the Operation hasn't
 * finished yet.
 */
void ExecutionContext::getActionResult(const Operation &operation,
                                       ActionResult *ar)
{
    if (!operation.done()) {
        throw std::logic_error(
            "Called getActionResult on an unfinished Operation");
    }
    else if (operation.has_error()) {
        throw GrpcError("Operation failed: " + operation.error().message(),
                        grpc::Status(static_cast<grpc::StatusCode>(
                                         operation.error().code()),
                                     operation.error().message()));
    }
    else if (!operation.response().Is<ExecuteResponse>()) {
        throw std::runtime_error("Server returned invalid Operation result");
    }

    ExecuteResponse executeResponse;
    if (!operation.response().UnpackTo(&executeResponse)) {
        throw std::runtime_error("Operation response unpacking failed");
    }

    const auto executeStatus = executeResponse.status();
    if (executeStatus.code() != google::rpc::Code::OK) {
        throw GrpcError(
            "Execution failed: " + executeStatus.message(),
            grpc::Status(static_cast<grpc::StatusCode>(executeStatus.code()),
                         executeStatus.message()));
    }

    const ActionResult actionResult = executeResponse.result();
    if (actionResult.exit_code() == 0) {
        BUILDBOX_LOG_DEBUG("Execute response message: " +
                           executeResponse.message());
    }
    else if (!executeResponse.message().empty()) {
        BUILDBOX_LOG_INFO("Remote execution message: " +
                          executeResponse.message());
    }

    *ar = actionResult;
}

bool ExecutionContext::downloadCompletedOperation(const std::string &where,
                                                  bool wait)
{
    try {
        Operation op;
        if (!wait || d_executionOptions->d_streamLogs) {
            // Get one single operation if we don't need to wait,
            // or we need the metadata for logstream
            op = d_executionClient->getOperation(
                d_executionOptions->d_operation);
            updateOperationMetadata(op);
            if (op.done()) {
                // we are certain we don't need to wait
                wait = false;
            }
        }
        if (wait) {
            if (d_executionOptions->d_streamLogs) {
                enableLogStream(d_resultMetadata.action_digest());
            }
            // Use WaitExecution to block until the operation is done
            op = d_executionClient->waitExecution(
                d_executionOptions->d_operation);
            updateOperationMetadata(op);
        }

        if (op.done()) {
            if (op.has_error()) {
                BUILDBOX_LOG_ERROR("Operation FAILED with error: " +
                                   op.error().DebugString());
                return false;
            }
            else {
                ActionResult ar;
                getActionResult(op, &ar);
                d_actionResult = std::make_shared<ActionResult>(ar);
                if (!d_executionOptions->d_actionResultJsonFile.empty()) {
                    writeProtoInJSONToFile(
                        d_executionOptions->d_actionResultJsonFile, ar);
                }
                updateActionResultMetadata(ar, false);
                return where.empty() ? true : downloadResults(ar, where);
            }
        }
        else {
            BUILDBOX_LOG_DEBUG(std::string("Operation ") +
                               d_executionOptions->d_operation +
                               std::string(" not yet completed"));
            return false;
        }
    }
    catch (buildboxcommon::GrpcError &grpcError) {
        updateGrpcError(grpcError.status);
        throw;
    }
}

bool ExecutionContext::cancelOperation()
{
    BUILDBOX_LOG_DEBUG("Cancelling operation ["
                       << d_executionOptions->d_operation << "]");
    try {
        d_executionClient->cancelOperation(d_executionOptions->d_operation);
    }
    catch (buildboxcommon::GrpcError &grpcError) {
        updateGrpcError(grpcError.status);
        throw;
    }
    return true;
}

bool ExecutionContext::execute(bool async)
{
    BUILDBOX_LOG_DEBUG("Execute called with async=" << async);
    // Upload missing blobs as necessary
    {
        BUILDBOX_LOG_DEBUG("Figuring out missing blobs to upload.");

        // Action and Command protos
        digest_string_map additionalProtosNeeded;
        additionalProtosNeeded.emplace(
            d_actionData.d_actionDigest,
            d_actionData.d_actionProto.SerializeAsString());
        additionalProtosNeeded.emplace(
            d_actionData.d_commandDigest,
            d_actionData.d_commandProto.SerializeAsString());

        std::vector<Digest> inputDigests;
        for (const auto &blob : d_actionData.d_inputDigestsToPaths) {
            inputDigests.emplace_back(blob.first);
        }
        for (const auto &blob :
             d_actionData.d_inputDigestsToSerializedProtos) {
            inputDigests.emplace_back(blob.first);
        }
        for (const auto &blob : additionalProtosNeeded) {
            inputDigests.emplace_back(blob.first);
        }
        BUILDBOX_LOG_DEBUG("Count of blobs needed: " << inputDigests.size());

        std::vector<Digest> missingDigests =
            d_casClient->findMissingBlobs(inputDigests);

        BUILDBOX_LOG_DEBUG(
            "Count of missing blobs: " << missingDigests.size());

        std::vector<buildboxcommon::CASClient::UploadRequest> upload_requests;
        std::vector<std::string> captureRequests;
        upload_requests.reserve(missingDigests.size());
        for (const auto &digest : missingDigests) {
            if (d_actionData.d_inputDigestsToPaths.count(digest) == 1) {
                if (d_executionOptions->d_useLocalCas) {
                    captureRequests.push_back(
                        d_actionData.d_inputDigestsToPaths.at(digest));
                }
                else {
                    upload_requests.push_back(
                        buildboxcommon::CASClient::UploadRequest::from_path(
                            digest,
                            d_actionData.d_inputDigestsToPaths.at(digest)));
                }
            }
            else if (d_actionData.d_inputDigestsToSerializedProtos.count(
                         digest)) {
                auto blob =
                    d_actionData.d_inputDigestsToSerializedProtos.at(digest);
                upload_requests.emplace_back(digest, blob);
            }
            else if (additionalProtosNeeded.count(digest)) {
                const auto &blob = additionalProtosNeeded.at(digest);
                upload_requests.emplace_back(digest, blob);
            }
            else {
                BUILDBOXCOMMON_THROW_EXCEPTION(
                    std::runtime_error, "FMB reported unknown blobs digest=["
                                            << toString(digest)
                                            << "] missing!");
            }
        }

        if (d_executionOptions->d_useLocalCas) {
            // If preferring local cas, CaptureFiles with the contents of
            // inputDigestsToPaths
            std::vector<std::string> captureProperties;
            BUILDBOX_LOG_DEBUG("Capturing " << captureRequests.size()
                                            << " blobs.");
            d_casClient->captureFiles(
                captureRequests, captureProperties, false,
                d_executionOptions->d_localCasSkipUpload);
        }

        BUILDBOX_LOG_DEBUG("Need to upload " << upload_requests.size()
                                             << " blobs");
        auto failedUploads =
            d_casClient->uploadBlobs(upload_requests, d_ioThreadPool.get());
        if (failedUploads.size() != 0) {
            const auto &firstError = failedUploads.at(0).status;
            updateGrpcError(firstError);
            BUILDBOXCOMMON_THROW_EXCEPTION(
                std::runtime_error,
                "Failed to upload "
                    << failedUploads.size() << " blobs. "
                    << "First error: " << firstError.error_code() << ": "
                    << firstError.error_message());
        }
    }

    bool stop_requested = false;
    ExecutionPolicy *policyPtr = nullptr;
    ExecutionPolicy policy;
    if (d_executionOptions->d_priority != nullptr) {
        policy.set_priority(*d_executionOptions->d_priority);
        policyPtr = &policy;
    }
    // Execute vs WaitExecute
    try {
        if (async) {
            google::longrunning::Operation op =
                d_executionClient->asyncExecuteAction(
                    d_actionData.d_actionDigest, stop_requested,
                    d_executionOptions->d_skipCacheLookup, policyPtr);
            updateOperationMetadata(op);
            std::cout << op.name() << '\n';
        }
        else {
            d_actionResult = std::make_shared<ActionResult>(
                d_executionClient->executeAction(
                    d_actionData.d_actionDigest, stop_requested,
                    d_executionOptions->d_skipCacheLookup, policyPtr));
            if (!d_executionOptions->d_actionResultJsonFile.empty()) {
                writeProtoInJSONToFile(
                    d_executionOptions->d_actionResultJsonFile,
                    *d_actionResult);
            }
            updateActionResultMetadata(*d_actionResult, false);
            d_getActionResultRemoteQueryCalled =
                true; // implicitly got ActionResult
            BUILDBOX_LOG_DEBUG("ExecuteAction returned!");
        }
    }
    catch (const buildboxcommon::GrpcError &e) {
        updateGrpcError(e.status);
        BUILDBOX_LOG_ERROR("Error in ExecuteAction: "
                           << e.status.error_code() << ": "
                           << e.status.error_message());
        return false;
    }

    if (d_actionResult != nullptr ||
        async) { // if async return true, as the actionResult will not
                 // be finished
        return true;
    }
    return false;
}

std::shared_ptr<ActionResult>
ExecutionContext::getActionResult(bool allowRemoteQuery)
{
    if (!d_actionResult && !d_getActionResultRemoteQueryCalled) {
        if (!allowRemoteQuery) {
            BUILDBOXCOMMON_THROW_EXCEPTION(
                std::runtime_error, "Called getActionResult for the first "
                                    "time but not allowing remote query!");
        }
        ActionResult actionResult;
        std::set<std::string>
            outputsToInline; // we don't want anything to be inlined
        try {
            if (d_executionClient->fetchFromActionCache(
                    d_actionData.d_actionDigest, outputsToInline,
                    &actionResult)) {
                d_actionResult = std::make_shared<ActionResult>(actionResult);
                if (!d_executionOptions->d_actionResultJsonFile.empty()) {
                    writeProtoInJSONToFile(
                        d_executionOptions->d_actionResultJsonFile,
                        actionResult);
                }
                updateActionResultMetadata(actionResult, true);
            }
            d_getActionResultRemoteQueryCalled = true;
        }
        catch (const GrpcError &e) {
            updateGrpcError(e.status);
            throw;
        }
    }

    return d_actionResult;
}

template <typename T>
void ExecutionContext::writeProtoInJSONToFile(const std::string &filepath,
                                              const T &proto, bool throwExc)
{
    std::ofstream outFile(filepath, std::fstream::out | std::fstream::trunc);
    if (!outFile.good()) {
        BUILDBOX_LOG_ERROR("cannot write to file " << filepath);
        return;
    }
    try {
        google::protobuf::util::JsonPrintOptions printOptions;
        buildboxcommon::jsonOptionsAlwaysPrintFieldsWithNoPresence(
            printOptions);

        std::string messageJson;
        auto status = google::protobuf::util::MessageToJsonString(
            proto, &messageJson, printOptions);
        if (!status.ok()) {
            BUILDBOX_LOG_ERROR(
                "cannot serialize proto message to JSON outputfile. filepath="
                << filepath << " proto=" << proto.SerializeAsString());
            if (throwExc) {
                throw std::invalid_argument(
                    "cannot serialize proto message to JSON");
            }
        }

        outFile << messageJson;
    }
    catch (const std::exception &e) {
        BUILDBOX_LOG_ERROR(
            "cannot deserialize proto message to output file. filepath="
            << filepath << " proto=" << proto.SerializeAsString());
        if (throwExc) {
            throw std::invalid_argument(
                "cannot deserialize proto message to output file");
        }
    }
}

void ExecutionContext::updateActionResultMetadata(
    const ActionResult &actionResult, bool cached_result)
{
    d_resultMetadata.mutable_action_result_metadata()->set_exit_code(
        actionResult.exit_code());
    d_resultMetadata.mutable_action_result_metadata()->set_cached_result(
        cached_result);

    *d_resultMetadata.mutable_execution_metadata() =
        actionResult.execution_metadata();
}

void ExecutionContext::updateOperationMetadata(const Operation &operation)
{
    d_resultMetadata.mutable_operation_metadata()->set_name(operation.name());
    if (operation.has_error()) {
        *d_resultMetadata.mutable_operation_metadata()
             ->mutable_operation_error() = operation.error();
    }
    ExecuteOperationMetadata executeOperationMetadata;
    if (operation.has_metadata() &&
        operation.metadata().UnpackTo(&executeOperationMetadata)) {
        *d_resultMetadata.mutable_action_digest() =
            executeOperationMetadata.action_digest();

        *d_resultMetadata.mutable_execution_metadata() =
            executeOperationMetadata.partial_execution_metadata();

        logExecutionMetadata(
            executeOperationMetadata.partial_execution_metadata());
    }
}

void ExecutionContext::logExecutionMetadata(
    const ::build::bazel::remote::execution::v2::ExecutedActionMetadata
        &metadata)
{
    if (!d_executionOptions->d_logProgress) {
        return;
    }

    using namespace buildboxcommon;

    if (metadata.has_queued_timestamp()) {
        BUILDBOX_LOG_INFO(
            "Reached scheduler queue at: "
            << TimeUtils::timeStampToStr(metadata.queued_timestamp()));
    }

    if (metadata.worker() != "") {
        BUILDBOX_LOG_INFO("Assigned to worker: " << metadata.worker());
    }

    if (metadata.has_worker_start_timestamp()) {
        BUILDBOX_LOG_INFO(
            "Worker started job at: "
            << TimeUtils::timeStampToStr(metadata.worker_start_timestamp()));
    }

    if (metadata.has_input_fetch_start_timestamp()) {
        BUILDBOX_LOG_INFO("Worker started fetching action inputs at: "
                          << TimeUtils::timeStampToStr(
                                 metadata.input_fetch_start_timestamp()));
    }

    if (metadata.has_input_fetch_completed_timestamp()) {
        BUILDBOX_LOG_INFO("Worker finished fetching action inputs at: "
                          << TimeUtils::timeStampToStr(
                                 metadata.input_fetch_completed_timestamp()));
    }

    if (metadata.has_execution_start_timestamp()) {
        BUILDBOX_LOG_INFO(
            "Worker started execution at: " << TimeUtils::timeStampToStr(
                metadata.execution_start_timestamp()));
    }

    if (metadata.has_execution_completed_timestamp()) {
        BUILDBOX_LOG_INFO(
            "Worker completed execution at: " << TimeUtils::timeStampToStr(
                metadata.execution_completed_timestamp()));
    }

    if (metadata.has_output_upload_start_timestamp()) {
        BUILDBOX_LOG_INFO("Worker started uploading action outputs at: "
                          << TimeUtils::timeStampToStr(
                                 metadata.output_upload_start_timestamp()));
    }

    if (metadata.has_output_upload_completed_timestamp()) {
        BUILDBOX_LOG_INFO(
            "Worker finished uploading action outputs at: "
            << TimeUtils::timeStampToStr(
                   metadata.output_upload_completed_timestamp()));
    }

    if (metadata.has_worker_completed_timestamp()) {
        BUILDBOX_LOG_INFO(
            "Worker completed job at: " << TimeUtils::timeStampToStr(
                metadata.worker_completed_timestamp()));
    }
}

std::shared_ptr<const ActionResult> ExecutionContext::actionResult() const
{
    return d_actionResult;
}

void ExecutionContext::updateGrpcError(const grpc::Status &error)
{
    auto metadataError = d_resultMetadata.mutable_error();
    metadataError->set_code(error.error_code());
    *metadataError->mutable_message() = error.error_message();
}

const TrexeResultMetadata &ExecutionContext::resultMetadata() const
{
    return d_resultMetadata;
}

bool ExecutionContext::outputStreamContent(std::ostream &stream,
                                           const std::string &inlineStream,
                                           const Digest &streamDigest)
{
    if (!inlineStream.empty()) {
        stream << inlineStream;
        return stream.good();
    }
    if (streamDigest.size_bytes() != 0) {
        StreamConsumer streamConsumer = [&](const ReadResponse &response,
                                            size_t /*offset*/) {
            stream << response.data();
            return stream.good()
                       ? grpc::Status::OK
                       : grpc::Status(
                             grpc::StatusCode::INTERNAL,
                             "Failed to write stream content to file");
        };
        try {
            d_casClient->fetchStream(streamDigest, streamConsumer, 0,
                                     std::nullopt);
        }
        catch (GrpcError &e) {
            updateGrpcError(e.status);
            return false;
        }
    }
    return true;
}

bool ExecutionContext::outputStdoutStderr()
{

    if (d_actionResult == nullptr) {
        BUILDBOX_LOG_ERROR(
            "Cannot output stdout/stderr before fetching ActionResult");
        return false;
    }

    bool success = true;
    {
        std::optional<std::ofstream> stdoutFile;
        if (!d_executionOptions->d_stdoutFile.empty()) {
            stdoutFile =
                std::ofstream(d_executionOptions->d_stdoutFile,
                              std::fstream::out | std::fstream::trunc);
        }
        if (!outputStreamContent(stdoutFile.has_value() ? *stdoutFile
                                                        : std::cout,
                                 d_actionResult->stdout_raw(),
                                 d_actionResult->stdout_digest())) {
            success = false;
        }
    }

    {
        std::optional<std::ofstream> stderrFile;
        if (!d_executionOptions->d_stderrFile.empty()) {
            stderrFile =
                std::ofstream(d_executionOptions->d_stderrFile,
                              std::fstream::out | std::fstream::trunc);
        }
        if (!outputStreamContent(stderrFile.has_value() ? *stderrFile
                                                        : std::cerr,
                                 d_actionResult->stderr_raw(),
                                 d_actionResult->stderr_digest())) {
            success = false;
        }
    }

    return success;
}

bool ExecutionContext::checkLogStreamStatus()
{
    // In cache-only mode, we never even try to stream logs
    if (d_executionOptions->d_cacheOnly) {
        return true;
    }

    // This downcast is safe because the `d_cacheOnly` execution option
    // determines which `ExecutionClient` subclass we have here. That
    // option must be `false` to reach this code, which means we are
    // using a `RemoteExecutionClient`.
    std::optional<buildboxcommon::LogStreamInfo> logStreamInfo =
        std::dynamic_pointer_cast<RemoteExecutionClient>(d_executionClient)
            ->getLogStreamInfo(d_actionData.d_actionDigest);
    if (logStreamInfo.has_value()) {
        if (logStreamInfo.value().stdoutStreamName.has_value() &&
            logStreamInfo.value().stderrStreamName.has_value()) {
            try {
                if (logStreamInfo.value().streamResults.at(
                        logStreamInfo.value().stdoutStreamName.value()) &&
                    logStreamInfo.value().streamResults.at(
                        logStreamInfo.value().stderrStreamName.value())) {
                    return true;
                }
            }
            catch (const std::out_of_range &) {
                return false;
            }
        }
    }

    return false;
}

} // namespace trexe
