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

#include <buildboxworker_worker.h>

#include <buildboxworker_actionutils.h>
#include <buildboxworker_botsessionutils.h>
#include <buildboxworker_config.h>
#include <buildboxworker_expiretime.h>
#include <buildboxworker_logstreamdebugutils.h>
#include <buildboxworker_metricnames.h>

#include <buildboxcommon_connectionoptions_commandline.h>
#include <buildboxcommon_exception.h>
#include <buildboxcommon_fileutils.h>
#include <buildboxcommon_grpcretrier.h>
#include <buildboxcommon_localexecutionclient.h>
#include <buildboxcommon_logging.h>
#include <buildboxcommon_requestmetadata.h>
#include <buildboxcommon_runner.h>
#include <buildboxcommon_stringutils.h>
#include <buildboxcommon_systemutils.h>
#include <buildboxcommon_version.h>
#include <buildboxcommonmetrics_durationmetrictimer.h>
#include <buildboxcommonmetrics_gaugemetricutil.h>
#include <buildboxcommonmetrics_metricguard.h>
#include <buildboxcommonmetrics_metricsconfigurator.h>

#include <google/devtools/remoteworkers/v1test2/worker.grpc.pb.h>
#include <google/rpc/code.pb.h>
#include <grpcpp/channel.h>

#include <algorithm>
#include <csignal>
#include <cstring>
#include <exception>
#include <fcntl.h>
#include <iterator>
#include <mutex>
#include <pthread.h>
#include <signal.h>
#include <stdexcept>
#include <string>
#include <sys/stat.h>
#include <sys/statvfs.h>
#include <sys/types.h>
#include <thread>
#include <unistd.h>

using buildboxcommon::buildboxcommonmetrics::GaugeMetricUtil;

namespace buildboxworker {

const std::chrono::microseconds
    Worker::s_defaultWaitTime(std::chrono::milliseconds(250));

namespace {

std::string generateDefaultBotID()
{
    constexpr int HOSTNAME_MAX = 256;
    char hostname[HOSTNAME_MAX] = {0};
    gethostname(hostname, sizeof(hostname));
    return std::string(hostname) + "-" + std::to_string(getpid());
}

std::atomic_bool d_cancelUpdateBotSession(false);

// Keeps track if re-reading the config was requested
std::atomic_bool d_rereadConfig(false);

// Keeps track if a graceful shutdown was requested
std::atomic_bool d_gracefulShutdown(false);

// Keeps track if a forceful shutdown was requested
std::atomic_bool d_forcefulShutdown(false);

} // namespace

void config_handler(int signal)
{
    d_cancelUpdateBotSession = true;
    d_rereadConfig = true;
}
void graceful_handler(int signal)
{
    d_cancelUpdateBotSession = true;
    d_gracefulShutdown = true;
}
void forceful_handler(int signal)
{
    d_cancelUpdateBotSession = true;
    d_gracefulShutdown = true;
    d_forcefulShutdown = true;
}

bool Worker::cancelUpdateBotSessionRequested()
{
    return d_cancelUpdateBotSession;
}

void Worker::setCancelUpdateBotSessionRequested(bool value)
{
    d_cancelUpdateBotSession = value;
}

Worker::Worker(
    const buildboxcommon::CommandLine &cml, const std::string &botId,
    buildboxcommon::buildboxcommonmetrics::MetricsConfigType *metricsConfig)
    : d_botID(!botId.empty() ? botId : generateDefaultBotID()),
      d_maxConcurrentJobs(cml.getInt("concurrent-jobs")),
      d_stopAfterJobs(cml.getInt("stop-after")),
      d_runnerCommand(cml.getString("buildbox-run")),
      d_extraRunArgs(cml.getVS("runner-arg")),
      d_logLevel(cml.getString("log-level")),
      d_logDirectory(cml.getString("log-directory")),
      d_configFileName(cml.getString("config-file")),
      d_botStatus(Config::getStatusFromConfigFile(d_configFileName)),
      d_consecutiveErrorsThreshold(cml.getInt("consecutive-errors-threshold")),
      d_priority(cml.getInt("priority")),
      d_maxWaitTime(std::chrono::seconds(cml.getInt("max-wait-time"))),
      d_preserveUpdateBotSession(cml.getBool("preserve-update-bot-session"))
{
    if (cml.exists("platform")) {
        d_platform = cml.getVPS("platform");
    }

    if (cml.exists("disk-usage-health-check")) {
        for (const auto &spec : cml.getVS("disk-usage-health-check")) {
            // Parse format: /path:80%
            auto colonPos = spec.rfind(':');
            if (colonPos == std::string::npos) {
                BUILDBOX_LOG_ERROR(
                    "Invalid disk-usage-health-check format: " << spec);
                continue;
            }
            std::string path = spec.substr(0, colonPos);
            std::string percentStr = spec.substr(colonPos + 1);

            // Remove trailing '%' if present
            if (!percentStr.empty() && percentStr.back() == '%') {
                percentStr.pop_back();
            }

            try {
                int threshold = std::stoi(percentStr);
                constexpr int MAX_PERCENTAGE = 100;
                if (threshold < 0 || threshold > MAX_PERCENTAGE) {
                    BUILDBOX_LOG_ERROR(
                        "Invalid percentage (must be 0-100): " << spec);
                    continue;
                }
                d_diskUsageChecks.emplace_back(path, threshold);
            }
            catch (const std::exception &e) {
                BUILDBOX_LOG_ERROR(
                    "Invalid disk-usage-health-check format: " << spec);
            }
        }
    }

    if (buildboxcommon::logging::stringToLogLevelMap().find(d_logLevel) ==
        buildboxcommon::logging::stringToLogLevelMap().end()) {
        BUILDBOX_LOG_ERROR("Invalid log level \"" << d_logLevel
                                                  << "\" setting to INFO");
        d_logLevel = "info";
        BUILDBOX_LOG_SET_LEVEL(buildboxcommon::LogLevel::INFO);
    }
    else if (cml.exists("verbose")) {
        BUILDBOX_LOG_SET_LEVEL(buildboxcommon::LogLevel::DEBUG);
    }
    else {
        BUILDBOX_LOG_SET_LEVEL(
            buildboxcommon::logging::stringToLogLevel.at(d_logLevel));
    }

    if (cml.exists("metrics-mode")) {
        buildboxcommon::buildboxcommonmetrics::MetricsConfigurator::
            metricsParser("metrics-mode", cml.getString("metrics-mode"),
                          metricsConfig);
        metricsConfig->setMetricTag("botID", d_botID);
    }

    if (cml.exists("metrics-publish-interval")) {
        buildboxcommon::buildboxcommonmetrics::MetricsConfigurator::
            metricsParser(
                "metrics-publish-interval",
                std::to_string(cml.getInt("metrics-publish-interval")),
                metricsConfig);
    }

    if (cml.exists("metrics-statsd-format")) {
        buildboxcommon::buildboxcommonmetrics::MetricsConfigurator::
            metricsParser("metrics-statsd-format",
                          cml.getString("metrics-statsd-format"),
                          metricsConfig);
    }

    if (cml.exists("graceful-shutdown-timeout")) {
        const int timeout = cml.getInt("graceful-shutdown-timeout");
        constexpr int MIN_TIMEOUT = 0;
        if (timeout <= MIN_TIMEOUT) {
            BUILDBOX_LOG_ERROR(
                "Invalid graceful-shutdown-timeout value of "
                << timeout << " seconds (timeout must be > 0s)."
                << " Graceful shutdown will not enforce a timeout.");
        }
        else {
            d_gracefulShutdownTimeout = std::chrono::seconds(timeout);
            d_gracefulShutdownTimeoutEnabled = true;
        }
    }

    // Set the defaults for bots, CAS and LogStream
    buildboxcommon::ConnectionOptionsCommandLine::configureChannel(
        cml, "", &d_botsServer);
    buildboxcommon::ConnectionOptionsCommandLine::configureChannel(
        cml, "", &d_casServer);
    buildboxcommon::ConnectionOptionsCommandLine::configureChannel(
        cml, "", &d_logStreamServer);
    buildboxcommon::ConnectionOptionsCommandLine::configureChannel(
        cml, "", &d_assetServer);

    // Configuring the bots, CAS and LogStream (optional), asset (optional)
    // channels:
    buildboxcommon::ConnectionOptionsCommandLine::updateChannelOptions(
        cml, "bots-", &d_botsServer);

    buildboxcommon::ConnectionOptionsCommandLine::updateChannelOptions(
        cml, "cas-", &d_casServer);

    buildboxcommon::ConnectionOptionsCommandLine::updateChannelOptions(
        cml, "logstream-", &d_logStreamServer);

    buildboxcommon::ConnectionOptionsCommandLine::updateChannelOptions(
        cml, "ra-", &d_assetServer);

#ifdef LOGSTREAM_DEBUG
    d_logstreamDebugCommand = cml.getString("launch-logstream-command", "");
#endif
}

bool Worker::validateConfiguration()
{
    if (d_botsServer.d_url.empty()) {
        BUILDBOX_LOG_ERROR("Bots server URL is missing");
        return false;
    }

    if (d_casServer.d_url.empty()) {
        BUILDBOX_LOG_ERROR("CAS server URL is missing");
        return false;
    }

    if (d_stopAfterJobs > 0 && d_maxConcurrentJobs > d_stopAfterJobs) {
        BUILDBOX_LOG_INFO(
            "WARNING: Max concurrent jobs ("
            << d_maxConcurrentJobs << ") > number of jobs to stop after ("
            << d_stopAfterJobs << "). Capping concurrent jobs to "
            << d_stopAfterJobs << ".");
        d_maxConcurrentJobs = d_stopAfterJobs;
    }

    return true;
}

void Worker::setUpStreamingIfNeeded(
    const std::string &stdoutStream, const std::string &stdoutFilePath,
    const std::string &stderrStream, const std::string &stderrFilePath,
    std::unique_ptr<buildboxcommon::StandardOutputStreamer> *stdoutStreamer,
    std::unique_ptr<buildboxcommon::StandardOutputStreamer> *stderrStreamer)
{
    if (this->d_logStreamServer.d_url.empty()) {
        return;
    }

    if (!stdoutStream.empty()) {
        BUILDBOX_LOG_INFO("Reading runner stdout from [ "
                          << stdoutFilePath << "] and streaming it to ["
                          << stdoutStream << "]");

        *stdoutStreamer =
            std::make_unique<buildboxcommon::StandardOutputStreamer>(
                stdoutFilePath, this->d_logStreamServer, stdoutStream);
    }

    if (!stderrStream.empty()) {
        BUILDBOX_LOG_INFO("Reading runner stderr from [ "
                          << stderrFilePath << "] and streaming it to ["
                          << stderrStream << "]");

        *stderrStreamer =
            std::make_unique<buildboxcommon::StandardOutputStreamer>(
                stderrFilePath, this->d_logStreamServer, stderrStream);
    }

#ifdef LOGSTREAM_DEBUG
    LogStreamDebugUtils::launchDebugCommand(d_logstreamDebugCommand,
                                            stdoutStream, stderrStream);
#endif
}

void Worker::stopStreamingIfNeeded(
    const std::unique_ptr<buildboxcommon::StandardOutputStreamer>
        &stdoutStreamer,
    const std::unique_ptr<buildboxcommon::StandardOutputStreamer>
        &stderrStreamer)
{
    if (stdoutStreamer) {
        if (stdoutStreamer->stop()) {
            BUILDBOX_LOG_DEBUG("stdout successfully committed");
        }
        else {
            BUILDBOX_LOG_ERROR(
                "Streaming operation failed, stdout was not committed");
        }
    }

    if (stderrStreamer) {
        if (stderrStreamer->stop()) {
            BUILDBOX_LOG_DEBUG("stderr successfully committed");
        }
        else {
            BUILDBOX_LOG_ERROR(
                "Streaming operation failed, stderr was not committed");
        }
    }
}

google::rpc::Status Worker::executeAction(
    const std::string &leaseId, const buildboxcommon::Action &action,
    const std::string &stdoutStream, const std::string &stderrStream,
    buildboxcommon::ActionResult *actionResult)
{
    buildboxcommon::buildboxcommonmetrics::MetricGuard<
        buildboxcommon::buildboxcommonmetrics::DurationMetricTimer>
        mt(MetricNames::TIMER_NAME_EXECUTE_ACTION);

    buildboxcommon::TemporaryFile stdoutFile;
    stdoutFile.close();
    buildboxcommon::TemporaryFile stderrFile;
    stderrFile.close();

    std::unique_ptr<buildboxcommon::StandardOutputStreamer> stdoutStreamer;
    std::unique_ptr<buildboxcommon::StandardOutputStreamer> stderrStreamer;
    setUpStreamingIfNeeded(stdoutStream, stdoutFile.strname(), stderrStream,
                           stderrFile.strname(), &stdoutStreamer,
                           &stderrStreamer);

    buildboxcommon::TemporaryFile partialExecutionMetadataFile;
    partialExecutionMetadataFile.close();
    {
        std::unique_lock<std::mutex> lock(d_sessionMutex);
        d_activeJobsToPartialExecutionMetadataTmpFile.insert(
            {leaseId, partialExecutionMetadataFile.strname()});
    }

    auto operation = this->spawnSubprocessForLease(
        leaseId, action, stdoutFile.strname(), stderrFile.strname(),
        partialExecutionMetadataFile.strname());

    if (!operation.done()) {
        operation = d_execClient->waitExecution(operation.name());
        this->untrackLeaseSubprocess(leaseId);
        {
            std::unique_lock<std::mutex> lock(d_sessionMutex);
            d_activeJobsToPartialExecutionMetadataTmpFile.erase(leaseId);
        }
    }

    stopStreamingIfNeeded(stdoutStreamer, stderrStreamer);

    buildboxcommon::ExecuteResponse executeResponse;
    if (!operation.response().UnpackTo(&executeResponse)) {
        throw std::runtime_error("Operation response unpacking failed");
    }
    *actionResult = executeResponse.result();
    return executeResponse.status();
}

void Worker::workerThread(const std::string &leaseId)
{
    BUILDBOX_LOG_DEBUG("Starting worker thread for leaseId " << leaseId);

    if (d_gracefulShutdown) {
        decrementThreadCount();
        return;
    }

    buildboxcommon::TemporaryFile actionResultFile;
    buildboxcommon::ActionResult actionResult;

    google::rpc::Status leaseStatus;
    try {
        buildboxcommon::Action action;
        buildboxcommon::Digest actionDigest;

        if (!fetchAction(leaseId, &action, &actionDigest)) {
            return; // Lease not found
        }

        // Fetching the optional stream names that can be set in metadata:
        std::string stdoutStream;
        std::string stderrStream;
        std::tie(stdoutStream, stderrStream) =
            standardOutputsStreamNames(actionDigest);

        // Executing the action:
        leaseStatus = executeAction(leaseId, action, stdoutStream,
                                    stderrStream, &actionResult);
        if (d_forcefulShutdown) {
            // Ensure that if the worker is shutting down the status is
            // non-ok
            leaseStatus.set_code(grpc::StatusCode::UNAVAILABLE);
            leaseStatus.set_message("Worker is shutting down");
        }

        actionResult.mutable_execution_metadata()->set_worker(d_botID);
    }
    catch (const std::exception &e) {
        BUILDBOX_LOG_ERROR("Caught exception in worker: " << e.what());

        leaseStatus.set_code(google::rpc::Code::INTERNAL);
        leaseStatus.set_message(
            "System error when attempting to execute Action: [" +
            std::string(e.what()) + "]");
    }

    BUILDBOX_LOG_DEBUG("Stopping worker thread for leaseId "
                       << leaseId << ". "
                       << (d_preserveUpdateBotSession
                               ? "Preserving UpdateBotSession."
                               : "Cancelling UpdateBotSession."))

    storeActionResultInLease(leaseId, leaseStatus, actionResult);
}

bool Worker::fetchAction(const std::string &leaseId,
                         buildboxcommon::Action *action,
                         buildboxcommon::Digest *actionDigest)
{
    // Lock while fetching the action
    std::lock_guard<std::mutex> lock(d_sessionMutex);

    // Searching for the lease...
    const auto lease = std::find_if(
        d_session.leases().cbegin(), d_session.leases().cend(),
        [&leaseId](const proto::Lease &l) { return l.id() == leaseId; });

    if (lease == d_session.leases().cend()) { // Not found!
        d_detachedThreadCount--;
        d_botSessionUpdate = true;
        d_sessionCondition.notify_all();
        return false;
    }

    // We found it, reading its action:
    std::tie(*action, *actionDigest) =
        ActionUtils::getActionFromLease(*lease, d_casClient);

    return true;
}

void Worker::storeActionResultInLease(
    const std::string &leaseId, const google::rpc::Status &leaseStatus,
    const buildboxcommon::ActionResult &actionResult)
{
    // Store the ActionResult back in the lease.

    std::lock_guard<std::mutex> lock(d_sessionMutex);

    // We need to search for the lease again because it could have been
    // moved/cancelled/deleted/completed by someone else while we were
    // executing it.
    for (auto &lease : *d_session.mutable_leases()) {
        const bool leaseFound = (lease.id() == leaseId &&
                                 lease.state() == proto::LeaseState::ACTIVE);
        if (leaseFound) {
            lease.set_state(proto::LeaseState::COMPLETED);
            lease.mutable_status()->CopyFrom(leaseStatus);
            lease.mutable_result()->PackFrom(actionResult);
        }
    }

    d_activeJobs.erase(leaseId);
    GaugeMetricUtil::setGauge(MetricNames::GAUGE_NUM_JOBS_RUNNING_ON_WORKER,
                              static_cast<int64_t>(this->d_activeJobs.size()));

    d_detachedThreadCount--;

    // Consecutive error tracking for worker health management
    if (d_consecutiveErrorsThreshold > 0) {
        bool isError = leaseStatus.code() != grpc::StatusCode::OK ||
                       actionResult.exit_code() != 0;

        if (isError) {
            d_consecutiveErrorCount++;
            BUILDBOX_LOG_DEBUG("Consecutive error count: "
                               << d_consecutiveErrorCount << "/"
                               << d_consecutiveErrorsThreshold);

            if (d_consecutiveErrorCount >= d_consecutiveErrorsThreshold) {
                d_markedUnhealthyByErrors = true;
                BUILDBOX_LOG_ERROR(
                    "Worker marked UNHEALTHY after "
                    << d_consecutiveErrorCount
                    << " consecutive errors. Status: " << leaseStatus.code()
                    << " (" << leaseStatus.message()
                    << "), Exit code: " << actionResult.exit_code());
            }
        }
        else {
            // Reset on success
            if (d_consecutiveErrorCount > 0) {
                BUILDBOX_LOG_DEBUG("Resetting consecutive error count after "
                                   "successful execution");
                d_consecutiveErrorCount = 0;
            }
        }
    }

    // Only cancel the UpdateBotSession RPC if we're not preserving it
    d_botSessionUpdate = true;
    d_cancelUpdateBotSession = !d_preserveUpdateBotSession;
    d_sessionCondition.notify_all();
}

void Worker::registerSignals()
{
    struct sigaction sa{};
    sa.sa_handler = graceful_handler;
    sigemptyset(&sa.sa_mask);
    sa.sa_flags = SA_SIGINFO;

    if (sigaction(SIGINT, &sa, nullptr) == -1) {
        BUILDBOXCOMMON_THROW_EXCEPTION(
            std::runtime_error,
            "Unable to register signal handler for SIGINT");
    }
    sa.sa_handler = forceful_handler;
    if (sigaction(SIGTERM, &sa, nullptr) == -1) {
        BUILDBOXCOMMON_THROW_EXCEPTION(
            std::runtime_error,
            "Unable to register signal handler for SIGTERM");
    }
    sa.sa_handler = config_handler;
    if (sigaction(SIGHUP, &sa, nullptr) == -1) {
        BUILDBOXCOMMON_THROW_EXCEPTION(
            std::runtime_error,
            "Unable to register signal handler for SIGHUP");
    }

    // Ignore SIGPIPE in case of using sockets + grpc without MSG_NOSIGNAL
    // support configured
    sa.sa_handler = SIG_IGN;
    sa.sa_flags = 0;
    if (sigaction(SIGPIPE, &sa, nullptr) == -1) {
        BUILDBOXCOMMON_THROW_EXCEPTION(std::runtime_error,
                                       "Unable to ignore SIGPIPE");
    }
}

void Worker::setPlatformProperties()
{
    auto workerProto = this->d_session.mutable_worker();

    auto device = workerProto->add_devices();
    device->set_handle(this->d_botID);

    for (const auto &platformPair : this->d_platform) {
        // TODO Differentiate worker properties and device properties?
        auto workerProperty = workerProto->add_properties();
        workerProperty->set_key(platformPair.first);
        workerProperty->set_value(platformPair.second);

        auto deviceProperty = device->add_properties();
        deviceProperty->set_key(platformPair.first);
        deviceProperty->set_value(platformPair.second);
    }
}

void Worker::setWorkerConfig()
{
    auto workerProto = this->d_session.mutable_worker();
    workerProto->clear_configs();
    auto config = workerProto->add_configs();
    config->set_key("capacity");
    config->set_value(std::to_string(this->d_maxConcurrentJobs));

    auto priorityConfig = workerProto->add_configs();
    priorityConfig->set_key("priority");
    priorityConfig->set_value(std::to_string(this->d_priority));
}

bool Worker::hasJobsToProcess() const
{
    return (this->d_stopAfterJobs != 0 || this->d_activeJobs.size() > 0 ||
            this->d_jobsPendingAck.size() > 0);
}

void Worker::initLocalExecutionClient()
{
    this->d_execClient =
        std::make_unique<buildboxcommon::LocalExecutionClient>(d_casServer,
                                                               d_assetServer);
    this->d_execClient->setRunner(d_runnerCommand, d_extraRunArgs);
    this->d_execClient->init();
}

void Worker::processLeases(bool *skipPollDelay)
{
    std::set<std::string> droppedLeases;
    std::set_union(this->d_activeJobs.begin(), this->d_activeJobs.end(),
                   this->d_jobsPendingAck.begin(),
                   this->d_jobsPendingAck.end(),
                   std::inserter(droppedLeases, droppedLeases.begin()));
    for (const auto &lease : this->d_session.leases()) {
        droppedLeases.erase(lease.id());
    }
    // If there are any leases we're keeping track of that were dropped by the
    // server, stop tracking them.
    for (auto &leaseId : droppedLeases) {
        BUILDBOX_LOG_DEBUG(
            "Lease ["
            << leaseId
            << "] was tracked locally but not by the server, removing.");
        if (this->d_jobsPendingAck.count(leaseId) != 0) {
            this->d_jobsPendingAck.erase(leaseId);
            if (this->d_stopAfterJobs >= 0) {
                this->d_stopAfterJobs++;
            }
        }
        else if (this->d_activeJobs.count(leaseId) != 0) {
            this->d_activeJobs.erase(leaseId);
            GaugeMetricUtil::setGauge(
                MetricNames::GAUGE_NUM_JOBS_RUNNING_ON_WORKER,
                static_cast<int64_t>(this->d_activeJobs.size()));
            terminateJobSubprocess(leaseId);
            if (this->d_stopAfterJobs >= 0) {
                this->d_stopAfterJobs++;
            }
        }
    }

    // Handle cancellation first
    for (auto &lease : *this->d_session.mutable_leases()) {
        if (lease.state() == proto::LeaseState::CANCELLED) {
            lease.mutable_status()->set_code(google::rpc::Code::CANCELLED);
            lease.mutable_status()->set_message("Lease cancelled");
            processCancelledLease(lease);
        }
    }

    // Handle pending/active session leases
    for (auto &lease : *this->d_session.mutable_leases()) {
        if (lease.state() == proto::LeaseState::PENDING) {
            processPendingLease(&lease, skipPollDelay);
        }
        else if (lease.state() == proto::LeaseState::ACTIVE) {
            processActiveLease(lease);
        }
    }
}

void Worker::processPendingLease(proto::Lease *lease, bool *skipPollDelay)
{
    BUILDBOX_LOG_DEBUG("Processing pending lease: " << lease->DebugString());
    if (lease->payload().Is<buildboxcommon::Action>() ||
        lease->payload().Is<buildboxcommon::Digest>()) {

        if (static_cast<int>(this->d_activeJobs.size() +
                             this->d_jobsPendingAck.size()) <
                this->d_maxConcurrentJobs &&
            this->d_stopAfterJobs != 0) {
            // Accept the lease, but wait for the server's ack
            // before actually starting work on it.
            lease->set_state(proto::LeaseState::ACTIVE);
            *skipPollDelay = true;
            this->d_jobsPendingAck.insert(lease->id());

            if (this->d_stopAfterJobs > 0) {
                this->d_stopAfterJobs--;
            }
        }
    }
    else {
        BUILDBOX_LOG_ERROR("Lease payload is of an unexpected type (not "
                           "`Action` or `Digest`): "
                           << lease->DebugString());

        auto status = lease->mutable_status();
        status->set_message("Invalid lease");
        status->set_code(google::rpc::Code::INVALID_ARGUMENT);

        lease->set_state(proto::LeaseState::COMPLETED);
    }
}

void Worker::processActiveLease(const proto::Lease &lease)
{
    BUILDBOX_LOG_DEBUG("Processing active lease: " << lease.DebugString());

    this->d_jobsPendingAck.erase(lease.id());

    if (this->d_activeJobs.count(lease.id()) == 0) {
        const std::string &leaseId = lease.id();
        auto thread =
            std::thread([this, leaseId] { this->workerThread(leaseId); });
        ++this->d_detachedThreadCount;
        thread.detach();
        // (the thread will remove its job from activeJobs when
        // it's done)
        this->d_activeJobs.insert(lease.id());
        GaugeMetricUtil::setGauge(
            MetricNames::GAUGE_NUM_JOBS_RUNNING_ON_WORKER,
            static_cast<int64_t>(this->d_activeJobs.size()));
    }
}

void Worker::processCancelledLease(const proto::Lease &lease)
{
    BUILDBOX_LOG_DEBUG("Processing cancelled lease: " << lease.DebugString());

    if (this->d_jobsPendingAck.count(lease.id()) != 0) {
        // We accepted the job, but the server decided that we
        // shouldn't run it after all.
        this->d_jobsPendingAck.erase(lease.id());
        if (this->d_stopAfterJobs >= 0) {
            this->d_stopAfterJobs++;
        }
    }
    else if (this->d_activeJobs.count(lease.id()) != 0) {
        // TODO: Refactor how we spin off threads to execute work
        this->d_activeJobs.erase(lease.id());
        GaugeMetricUtil::setGauge(
            MetricNames::GAUGE_NUM_JOBS_RUNNING_ON_WORKER,
            static_cast<int64_t>(this->d_activeJobs.size()));

        // Terminate the runner subprocess
        terminateJobSubprocess(lease.id());
        if (this->d_stopAfterJobs >= 0) {
            this->d_stopAfterJobs++;
        }
    }
}

std::chrono::system_clock::time_point Worker::calculateWaitTime(
    const std::chrono::system_clock::time_point currentTime) const
{
    if (!this->d_session.has_expire_time() ||
        (static_cast<int>(this->d_activeJobs.size() +
                          this->d_jobsPendingAck.size()) <
             this->d_maxConcurrentJobs &&
         this->d_botStatus == proto::BotStatus::OK)) {
        // `expire_time` may be way too long if this bot is healthy and
        // we don't have any pending work, and that would mean workers
        // won't pick up work for a while. So in that case we use the default
        // value.
        return currentTime + s_defaultWaitTime;
    }

    // convert google.protobuf.Timestamp to std::chrono::time_point
    const auto expireTime =
        ExpireTime::convertToTimePoint(this->d_session.expire_time());

    if (expireTime <= currentTime) {
        using namespace buildboxcommon;
        BUILDBOX_LOG_WARNING(
            "BotSession::expire_time is either "
            "now or in the past: expireTime = "
            << TimeUtils::timePointToStr(expireTime)
            << ", now = " << TimeUtils::timePointToStr(currentTime)
            << ", using default wait time of "
            << std::chrono::duration_cast<std::chrono::seconds>(
                   s_defaultWaitTime)
                   .count()
            << " seconds");

        return currentTime + s_defaultWaitTime;
    }

    // calc the difference
    auto waitTime = std::chrono::duration_cast<std::chrono::microseconds>(
        expireTime - currentTime);

    // minus some percentage of time allow us to respond
    const auto factor =
        static_cast<int64_t>(static_cast<float>(waitTime.count()) *
                             ExpireTime::updateTimeoutPaddingFactor());
    waitTime -= std::chrono::microseconds(factor);

    // check for sane upper bounds
    if (waitTime > d_maxWaitTime) {
        BUILDBOX_LOG_DEBUG(
            "detected excessive BotSession::expire_time set to "
            << waitTime.count() << "us, using max wait time of "
            << std::chrono::duration_cast<std::chrono::seconds>(d_maxWaitTime)
                   .count()
            << " seconds");

        waitTime = d_maxWaitTime;
    }
    return currentTime + waitTime;
}

void Worker::decrementThreadCount()
{
    std::lock_guard<std::mutex> lock(this->d_sessionMutex);
    --this->d_detachedThreadCount;
    this->d_sessionCondition.notify_all();
}

void Worker::shutdownTrackedSubprocessPgids()
{
    this->d_execClient->cancelAllOperations();
}

std::pair<std::string, std::string> Worker::standardOutputsStreamNames(
    const buildboxcommon::Digest &actionDigest) const
{
    std::string stdoutStreamName;
    std::string stderrStreamName;

    const auto it = d_sessionExecuteOperationMetadata.find(actionDigest);
    if (it != d_sessionExecuteOperationMetadata.cend()) {
        stdoutStreamName = it->second.stdout_stream_name();
        stderrStreamName = it->second.stderr_stream_name();
    }

    return std::make_pair(stdoutStreamName, stderrStreamName);
}

void Worker::logSuppliedParameters() const
{
    std::ostringstream oss;
    oss << "Starting buildbox-worker with BotID = \"" << d_botID
        << "\", instanceName = \"" << d_botsServer.d_instanceName
        << "\", CAS-remote = [" << d_casServer << "], BOTS-remote = ["
        << d_botsServer << "], log-level = \"" << d_logLevel
        << ", buildbox-run = \"" << d_runnerCommand << "\"";

    oss << ", runner-arg = [";
    for (size_t i = 0; i < d_extraRunArgs.size(); ++i) {
        if (i > 0) {
            oss << ", ";
        }
        oss << "\"" << d_extraRunArgs[i] << "\"";
    }

    oss << "], platform = [";
    for (size_t i = 0; i < d_platform.size(); ++i) {
        if (i > 0) {
            oss << ", ";
        }
        oss << "[\"" << d_platform[i].first << "\" = \""
            << d_platform[i].second << "\"]";
    }

    oss << "], disk-usage-health-check = [";
    for (size_t i = 0; i < d_diskUsageChecks.size(); ++i) {
        if (i > 0) {
            oss << ", ";
        }
        oss << "[\"" << d_diskUsageChecks[i].first
            << "\" = " << d_diskUsageChecks[i].second << "%]";
    }

    oss << "], config-file = \"" << d_configFileName
        << "\", botStatus = " << d_botStatus;

    BUILDBOX_LOG_INFO(oss.str());
}

void Worker::runWorker()
{
    registerSignals(); // Handle SIGINT, SIGTERM and SIGHUP and ignore SIGPIPE

    std::unique_lock<std::mutex> lock(this->d_sessionMutex);

    logSuppliedParameters();

    this->d_session.set_bot_id(this->d_botID);
    setPlatformProperties();
    setWorkerConfig();

    this->d_botsStub =
        proto::Bots::NewStub(this->d_botsServer.createChannel());

    auto grpcClient = std::make_shared<buildboxcommon::GrpcClient>();
    grpcClient->init(d_casServer);
    grpcClient->setToolDetails("buildbox-worker", buildboxcommon::VERSION);
    const std::string toolInvocationId =
        buildboxcommon::StringUtils::getUUIDString();
    const std::string correlatedInvocationsId = this->d_botID;

    grpcClient->setRequestMetadata("", toolInvocationId,
                                   correlatedInvocationsId);
    this->d_casClient =
        std::make_shared<buildboxcommon::CASClient>(grpcClient);
    this->d_casClient->init();

    const buildboxcommon::GrpcRetrier::GrpcStatusCodes retryableStatusCodes = {
        grpc::StatusCode::UNAVAILABLE};

    buildboxcommon::RequestMetadataGenerator requestMetadata =
        buildboxcommon::RequestMetadataGenerator();
    requestMetadata.set_tool_details("buildbox-worker",
                                     buildboxcommon::VERSION);
    requestMetadata.set_tool_invocation_id(toolInvocationId);
    requestMetadata.set_correlated_invocations_id(correlatedInvocationsId);

    const auto createSessionStatus = BotSessionUtils::createBotSession(
        d_botsStub, d_botsServer, &d_session, retryableStatusCodes,
        requestMetadata, d_gracefulShutdown);
    if (!createSessionStatus.ok()) {
        BUILDBOXCOMMON_THROW_EXCEPTION(
            std::runtime_error,
            "Failed to create bot session: gRPC error " +
                std::to_string(createSessionStatus.error_code()) + ": " +
                createSessionStatus.error_message());
    }

    BUILDBOX_LOG_DEBUG("Bot Session created. Now waiting for jobs.");

    GaugeMetricUtil::setGauge(MetricNames::GAUGE_MAXIMUM_WORKER_CAPACITY,
                              static_cast<int64_t>(d_maxConcurrentJobs));

    bool skipPollDelay = true;
    std::chrono::time_point<std::chrono::system_clock> waitUntil =
        std::chrono::system_clock::now();
    while (this->hasJobsToProcess() && !d_gracefulShutdown) {
        if (skipPollDelay) {
            skipPollDelay = false;
            d_botSessionUpdate = true;
        }
        else {
            this->d_sessionCondition.wait_for(lock, this->s_defaultWaitTime);
        }

        if (d_rereadConfig) {
            d_botStatus =
                Config::getStatusFromConfigFile(this->d_configFileName);

            // Note: consecutive errors threshold is from command line only,
            // so it cannot be changed via config reload. The UNHEALTHY state
            // from consecutive errors persists until worker restart.
            if (d_markedUnhealthyByErrors) {
                BUILDBOX_LOG_INFO(
                    "Config reloaded, but worker remains UNHEALTHY due to "
                    "consecutive errors (restart required to clear)");
            }

            d_rereadConfig = false;
            d_botSessionUpdate = true;
        }
        if (d_botSessionUpdate ||
            waitUntil < std::chrono::system_clock::now()) {
            processLeases(&skipPollDelay);
            setWorkerConfig();

            // Compute effective status for this update
            proto::BotStatus effectiveStatus = d_botStatus;

            // Override config-based status when marked unhealthy by
            // consecutive errors (persistent until restart)
            if (d_markedUnhealthyByErrors) {
                effectiveStatus = proto::BotStatus::UNHEALTHY;
                BUILDBOX_LOG_DEBUG("Status override: consecutive errors "
                                   "UNHEALTHY takes precedence over config");
            }

            // Check disk usage health (re-evaluated each iteration)
            if (!d_diskUsageChecks.empty() && !checkDiskUsageHealth()) {
                effectiveStatus = proto::BotStatus::UNHEALTHY;
                BUILDBOX_LOG_DEBUG("Status override: disk usage UNHEALTHY "
                                   "takes precedence over config");
            }

            grpc::Status updateStatus = BotSessionUtils::updateBotSession(
                effectiveStatus, d_botsStub, d_botsServer, &d_session,
                retryableStatusCodes, d_cancelUpdateBotSession,
                requestMetadata, getPartialExecutionMetadataFromRunner(),
                &lock, &d_sessionExecuteOperationMetadata);
            d_botSessionUpdate = false;

            if (updateStatus.ok()) {
                const bool hasPendingLeases = std::any_of(
                    this->d_session.leases().begin(),
                    this->d_session.leases().end(), [](const auto &lease) {
                        return lease.state() == proto::LeaseState::PENDING;
                    });
                if (hasPendingLeases) {
                    // If we have pending leases, we do not need to wait to
                    // mark them active
                    skipPollDelay = true;
                }

                waitUntil = this->calculateWaitTime();
            }
            else if (updateStatus.error_code() ==
                     grpc::StatusCode::CANCELLED) {
                skipPollDelay = true;
                d_cancelUpdateBotSession = false;
            }
            else {
                // If communication with the server fails, shutdown
                // immediately.
                d_gracefulShutdown = true;
                d_forcefulShutdown = true;
                BUILDBOX_LOG_ERROR(
                    "Failed to update bot session: gRPC error " +
                    std::to_string(updateStatus.error_code()) + ": " +
                    updateStatus.error_message());
            }
        }
    }

    if (d_gracefulShutdown) {
        // Reset flag now that we're processing shutdown
        d_gracefulShutdown = false;
        BUILDBOX_LOG_DEBUG("Received signal, exiting.");
        // Wait for all detached threads to finish.
        d_botSessionUpdate = true;

        std::chrono::time_point<std::chrono::system_clock> shutdownTimeout =
            std::chrono::system_clock::now() + d_gracefulShutdownTimeout;

        if (d_gracefulShutdownTimeoutEnabled && !d_forcefulShutdown) {
            BUILDBOX_LOG_DEBUG(
                "Graceful shutdown timeout is enabled, waiting up to "
                << std::chrono::duration_cast<std::chrono::seconds>(
                       d_gracefulShutdownTimeout)
                       .count()
                << " seconds for active jobs to finish, before forcefully "
                   "shutting down.");
        }

        // TODO: Instead of detaching threads and checking
        // d_detachedThreadCount > 0 during shutdown, explicitly managing the
        // threads and calling thread.join() might be more robust
        bool signalSentToSubprocesses = false;
        while (this->d_detachedThreadCount > 0) {
            // Timeout reached for graceful shutdown, forcefully shutdown
            if (d_gracefulShutdownTimeoutEnabled &&
                shutdownTimeout < std::chrono::system_clock::now() &&
                !d_forcefulShutdown) {
                BUILDBOX_LOG_DEBUG("Graceful shutdown timeout reached, "
                                   "forcefully shutting down "
                                   "now.");
                d_forcefulShutdown = true;
            }

            if (d_forcefulShutdown && !signalSentToSubprocesses) {
                BUILDBOX_LOG_DEBUG(
                    "Immediate shutdown requested, sending SIGTERM "
                    "to runner processes.");
                shutdownTrackedSubprocessPgids();
                signalSentToSubprocesses = true;
            }
            if (d_botSessionUpdate ||
                waitUntil < std::chrono::system_clock::now()) {
                // Attempt to gracefully shutdown by notifying the server of
                // our status and waiting for work to finish
                grpc::Status updateStatus = BotSessionUtils::updateBotSession(
                    proto::BotStatus::BOT_TERMINATING, d_botsStub,
                    d_botsServer, &d_session, retryableStatusCodes,
                    d_forcefulShutdown, requestMetadata,
                    getPartialExecutionMetadataFromRunner(), &lock);
                d_botSessionUpdate = false;

                // But if trying to talk to the server fails, enable forceful
                // shutdown to shutdown any spun off subprocesses doing work.
                if (!updateStatus.ok() &&
                    updateStatus.error_code() != grpc::StatusCode::CANCELLED) {
                    d_forcefulShutdown = true;
                    BUILDBOX_LOG_ERROR(
                        "Failed to update bot session for graceful shutdown: "
                        "gRPC error " +
                        std::to_string(updateStatus.error_code()) + ": " +
                        updateStatus.error_message());
                }
                waitUntil = this->calculateWaitTime();
            }

            this->d_sessionCondition.wait_for(lock, this->s_defaultWaitTime);
        }
        // Send a final update before exiting.
        grpc::Status finalStatus = BotSessionUtils::updateBotSession(
            proto::BotStatus::BOT_TERMINATING, d_botsStub, d_botsServer,
            &d_session, retryableStatusCodes, requestMetadata,
            getPartialExecutionMetadataFromRunner(), &lock);
        if (!finalStatus.ok()) {
            BUILDBOXCOMMON_THROW_EXCEPTION(
                std::runtime_error,
                "Final bot session update failed: gRPC error " +
                    std::to_string(finalStatus.error_code()) + ": " +
                    finalStatus.error_message());
        }
    }
}

void Worker::runWorkerWithoutBotSession()
{
    bool skipPollDelay = true;
    std::unique_lock<std::mutex> lock(this->d_sessionMutex);
    while (this->hasJobsToProcess()) {
        if (skipPollDelay) {
            skipPollDelay = false;
        }
        else {
            this->d_sessionCondition.wait(lock);
        }

        processLeases(&skipPollDelay);
    }
}

bool Worker::testRunnerCommand()
{
    return this->d_execClient->validateRunner();
}

const std::chrono::microseconds &Worker::maxWaitTime() const
{
    return d_maxWaitTime;
}

bool Worker::terminateJobSubprocess(const std::string &leaseId)
{
    const auto iter = d_activeJobsToOperations.find(leaseId);
    if (iter != d_activeJobsToOperations.end()) {
        if (!d_execClient->cancelOperation(iter->second)) {
            return false;
        }
        d_activeJobsToOperations.erase(iter);
        return true;
    }
    return false;
}

buildboxcommon::Operation Worker::spawnSubprocessForLease(
    const std::string &leaseId, const buildboxcommon::Action &action,
    const std::string &stdoutFile, const std::string &stderrFile,
    const std::string &partialExecutionMetadataFile)
{
    std::lock_guard<std::mutex> lock(d_sessionMutex);
    if (d_activeJobs.count(leaseId) == 0) {
        BUILDBOXCOMMON_THROW_EXCEPTION(
            std::runtime_error, "Lease is not active when trying to fork a "
                                "subprocess to execute it, leaseId ["
                                    << leaseId << "]");
    }
    if (d_activeJobsToOperations.count(leaseId) != 0) {
        BUILDBOXCOMMON_THROW_EXCEPTION(
            std::runtime_error, "A subprocess is already created for lease : ["
                                    << leaseId << ", operation : ["
                                    << d_activeJobsToOperations.at(leaseId));
    }

    auto operation = d_execClient->asyncExecuteAction(
        action, stdoutFile, stderrFile, partialExecutionMetadataFile);
    if (!operation.done()) {
        d_activeJobsToOperations.insert({leaseId, operation.name()});
    }
    return operation;
}

bool Worker::untrackLeaseSubprocess(const std::string &leaseId)
{
    std::lock_guard<std::mutex> lock(d_sessionMutex);
    size_t num_erased = d_activeJobsToOperations.erase(leaseId);
    return num_erased != 0;
}

RunnerPartialExecutionMetadata Worker::getPartialExecutionMetadataFromRunner()
{
    RunnerPartialExecutionMetadata partialExecutionMetadata;
    for (const auto &[leaseID, path] :
         d_activeJobsToPartialExecutionMetadataTmpFile) {
        try {
            ExecutedActionMetadata metadata;
            std::string data =
                buildboxcommon::FileUtils::getFileContents(path.c_str());
            metadata.ParseFromString(data);
            partialExecutionMetadata.insert({leaseID, metadata});
        }
        catch (std::exception &e) {
            BUILDBOX_LOG_ERROR("Failed to get metadata for leaseId: "
                               << leaseID << " with exception: " << e.what());
        }
    }

    return partialExecutionMetadata;
}

bool Worker::checkDiskUsageHealth()
{
    for (const auto &[path, threshold] : d_diskUsageChecks) {
        struct statvfs buf{};
        if (statvfs(path.c_str(), &buf) < 0) {
            BUILDBOX_LOG_ERROR("Failed to get disk stats for "
                               << path << ": " << strerror(errno));
            // Treat statvfs failure as unhealthy
            return false;
        }

        if (buf.f_blocks == 0) {
            BUILDBOX_LOG_ERROR("Invalid filesystem at " << path
                                                        << ": zero blocks");
            return false;
        }

        // Use integer arithmetic to avoid floating point precision issues when
        // narrowing to double
        constexpr decltype(buf.f_blocks) hundred = 100;
        auto usedBlocks = buf.f_blocks - buf.f_bavail;
        auto usedTimes100 = usedBlocks * hundred;
        auto thresholdTimesTotal =
            static_cast<decltype(usedBlocks)>(threshold) * buf.f_blocks;

        if (usedTimes100 >= thresholdTimesTotal) {
            auto usagePercent = usedBlocks * hundred / buf.f_blocks;
            BUILDBOX_LOG_WARNING("Disk usage at "
                                 << path << " is " << usagePercent
                                 << "% (threshold: " << threshold << "%)");
            return false;
        }

        BUILDBOX_LOG_DEBUG("Disk usage at "
                           << path << ": "
                           << usedBlocks * hundred / buf.f_blocks
                           << "% (threshold: " << threshold << "%)");
    }
    return true;
}

} // namespace buildboxworker
