
#include <buildboxcasd_requestcontextmanager.h>
#include <buildboxcommon_requestmetadata.h>

using namespace buildboxcasd;
using namespace build::bazel::remote::execution::v2;

thread_local RequestMetadata RequestContextManager::s_requestMetadata =
    RequestMetadata::default_instance();
thread_local bool RequestContextManager::s_metadataSet = false;
thread_local const grpc::ServerContext
    *RequestContextManager::s_serverContext = nullptr;

void RequestContextManager::setRequestMetadata(const RequestMetadata &metadata)
{
    s_requestMetadata.CopyFrom(metadata);
    s_requestMetadata.mutable_tool_details()->set_tool_name(
        s_requestMetadata.tool_details().tool_name() + " (via buildbox-casd)");
    s_metadataSet = true;
}

RequestMetadata RequestContextManager::getRequestMetadata()
{
    return s_requestMetadata;
}

void RequestContextManager::attachMetadata(grpc::ClientContext *ctx)
{
    if (s_metadataSet) {
        ctx->AddMetadata(
            buildboxcommon::RequestMetadataGenerator::HEADER_NAME,
            RequestContextManager::getRequestMetadata().SerializeAsString());
    }
}

void RequestContextManager::setServerContext(const grpc::ServerContext *ctx)
{
    s_serverContext = ctx;
}

const grpc::ServerContext *RequestContextManager::getServerContext()
{
    return s_serverContext;
}

void RequestContextManager::reset()
{
    s_requestMetadata = RequestMetadata::default_instance();
    s_metadataSet = false;
    s_serverContext = nullptr;
}

void RequestContextManager::configureGrpcClient(
    buildboxcommon::GrpcClient *client)
{
    client->setMetadataAttacher(RequestContextManager::attachMetadata);
    client->setServerContextGetter(RequestContextManager::getServerContext);
}
