/*
 * Copyright 2019 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_BUILDBOXCASD_SERVER_H
#define INCLUDED_BUILDBOXCASD_SERVER_H

#include <grpc/grpc.h>
#include <grpcpp/server.h>
#include <grpcpp/server_builder.h>
#include <grpcpp/server_context.h>
#include <memory>
#include <optional>
#include <ThreadPool.h>
#include <unordered_map>

#include <buildboxcasd_actionstorage.h>
#include <buildboxcasd_casinstance.h>
#include <buildboxcasd_fslocalassetstorage.h>
#include <buildboxcasd_inodecache.h>
#include <buildboxcasd_instancemanager.h>
#include <buildboxcasd_localacinstance.h>
#include <buildboxcasd_localexecutionscheduler.h>
#include <buildboxcommon_connectionoptions.h>
#include <buildboxcommon_filestager.h>
#include <buildboxcommon_localcas.h>
#include <buildboxcommon_protos.h>

namespace buildboxcasd {

using namespace build::bazel::remote::asset::v1;
using namespace build::bazel::remote::execution::v2;
using namespace build::buildgrid;
using namespace google::protobuf::util;
using namespace google::bytestream;
using namespace google::longrunning;

using grpc::ServerContext;
using grpc::ServerReader;
using grpc::ServerWriter;
using grpc::Status;

class CapabilitiesServicer final : public Capabilities::Service {
  public:
    explicit CapabilitiesServicer(
        const std::shared_ptr<InstanceManager> &instance_manager);

    Status GetCapabilities(grpc::ServerContext *ctx,
                           const GetCapabilitiesRequest *request,
                           ServerCapabilities *response) override;

  private:
    std::shared_ptr<InstanceManager> d_instance_manager;
    ServerCapabilities d_server_capabilities;
};

class CasRemoteExecutionServicer final
    : public ContentAddressableStorage::Service {
    /* Implements the CAS RPC methods. */

  public:
    explicit CasRemoteExecutionServicer(
        const std::shared_ptr<InstanceManager> &instance_manager);

    Status FindMissingBlobs(ServerContext *ctx,
                            const FindMissingBlobsRequest *request,
                            FindMissingBlobsResponse *response) override;

    Status BatchUpdateBlobs(ServerContext *ctx,
                            const BatchUpdateBlobsRequest *request,
                            BatchUpdateBlobsResponse *response) override;

    Status BatchReadBlobs(ServerContext *ctx,
                          const BatchReadBlobsRequest *request,
                          BatchReadBlobsResponse *response) override;

    Status GetTree(ServerContext *ctx, const GetTreeRequest *request,
                   ServerWriter<GetTreeResponse> *writer) override;

  private:
    std::shared_ptr<InstanceManager> d_instance_manager;
};

class CasBytestreamServicer final : public ByteStream::Service {
    /* Implements the Bytestream API Read() and Write() methods for the CAS. */
  public:
    explicit CasBytestreamServicer(
        const std::shared_ptr<InstanceManager> &instance_manager);

    Status Read(grpc::ServerContext *ctx, const ReadRequest *request,
                ServerWriter<ReadResponse> *writer) override;

    Status Write(grpc::ServerContext *ctx, ServerReader<WriteRequest> *request,
                 WriteResponse *response) override;

  private:
    std::string readInstanceName(const std::string &resource_name) const;
    std::string writeInstanceName(const std::string &resource_name) const;

    std::string
    instanceNameFromResourceName(const std::string &resource_name,
                                 const std::string &root_path) const;

    std::shared_ptr<InstanceManager> d_instance_manager;
};

class LocalCasServicer final : public LocalContentAddressableStorage::Service {
    /* Implements the LocalContentAddressableStorage RPC methods */

  public:
    LocalCasServicer(
        Server *server,
        const std::shared_ptr<InstanceManager> &instance_manager,
        const std::shared_ptr<buildboxcommon::LocalCas> &storage,
        const std::shared_ptr<FsLocalAssetStorage> &asset_storage,
        const std::shared_ptr<ActionStorage> &actioncache_storage,
        std::unordered_map<buildboxcommon::StageTreeRequest::StagingMode,
                           std::shared_ptr<buildboxcommon::FileStager>>
            fileStagers,
        const std::shared_ptr<LocalExecutionScheduler> &executionScheduler,
        const std::shared_ptr<ThreadPool> &digestThreadPool,
        const std::shared_ptr<ThreadPool> &ioThreadPool,
        const std::shared_ptr<std::unordered_set<std::string>>
            &allowedPreUnstageCommands);

    Status GetLocalServerDetails(ServerContext *context,
                                 const GetLocalServerDetailsRequest *request,
                                 LocalServerDetails *response) override;

    Status FetchMissingBlobs(ServerContext *ctx,
                             const FetchMissingBlobsRequest *request,
                             FetchMissingBlobsResponse *response) override;

    Status UploadMissingBlobs(ServerContext *ctx,
                              const UploadMissingBlobsRequest *request,
                              UploadMissingBlobsResponse *response) override;

    Status FetchTree(ServerContext *ctx, const FetchTreeRequest *request,
                     FetchTreeResponse *response) override;

    Status UploadTree(ServerContext *ctx, const UploadTreeRequest *request,
                      UploadTreeResponse *response) override;

    Status StageTree(ServerContext *ctx,
                     ServerReaderWriter<StageTreeResponse, StageTreeRequest>
                         *stream) override;

    Status
    NestedServer(ServerContext *ctx,
                 ServerReaderWriter<NestedServerResponse, NestedServerRequest>
                     *stream) override;

    Status ExportFiles(ServerContext *ctx, const ExportFilesRequest *request,
                       ExportFilesResponse *response) override;

    Status CaptureTree(ServerContext *ctx, const CaptureTreeRequest *request,
                       CaptureTreeResponse *response) override;

    Status CaptureFiles(ServerContext *ctx, const CaptureFilesRequest *request,
                        CaptureFilesResponse *response) override;

    Status HashFiles(ServerContext *ctx, const HashFilesRequest *request,
                     HashFilesResponse *response) override;

    Status GetInstanceNameForRemote(
        ServerContext *ctx, const GetInstanceNameForRemoteRequest *request,
        GetInstanceNameForRemoteResponse *response) override;

    Status GetInstanceNameForRemotes(
        ServerContext *ctx, const GetInstanceNameForRemotesRequest *request,
        GetInstanceNameForRemotesResponse *response) override;

    Status GetLocalDiskUsage(ServerContext *ctx,
                             const GetLocalDiskUsageRequest *request,
                             GetLocalDiskUsageResponse *response) override;

  private:
    Server *d_server;
    std::shared_ptr<InstanceManager> d_instance_manager;

    std::shared_ptr<buildboxcommon::LocalCas> d_storage;
    std::shared_ptr<FsLocalAssetStorage> d_asset_storage;
    std::shared_ptr<ActionStorage> d_actioncache_storage;
    std::unordered_map<buildboxcommon::StageTreeRequest::StagingMode,
                       std::shared_ptr<buildboxcommon::FileStager>>
        d_fileStagers;
    std::shared_ptr<LocalExecutionScheduler> d_executionScheduler;
    std::shared_ptr<ThreadPool> d_digestThreadPool;
    std::shared_ptr<ThreadPool> d_ioThreadPool;
    std::shared_ptr<std::unordered_set<std::string>>
        d_allowedPreUnstageCommands;
};

class AssetFetchServicer final : public Fetch::Service {
  public:
    explicit AssetFetchServicer(
        const std::shared_ptr<InstanceManager> &instance_manager);

    Status FetchBlob(ServerContext *ctx, const FetchBlobRequest *request,
                     FetchBlobResponse *response) override;

    Status FetchDirectory(ServerContext *ctx,
                          const FetchDirectoryRequest *request,
                          FetchDirectoryResponse *response) override;

  private:
    std::shared_ptr<InstanceManager> d_instance_manager;
};

class AssetPushServicer final : public Push::Service {
  public:
    explicit AssetPushServicer(
        const std::shared_ptr<InstanceManager> &instance_manager);

    Status PushBlob(ServerContext *ctx, const PushBlobRequest *request,
                    PushBlobResponse *response) override;

    Status PushDirectory(ServerContext *ctx,
                         const PushDirectoryRequest *request,
                         PushDirectoryResponse *response) override;

  private:
    std::shared_ptr<InstanceManager> d_instance_manager;
};

class ActionCacheServicer final : public ActionCache::Service {
    /* Implements the ActionCache RPC methods */
  public:
    explicit ActionCacheServicer(
        const std::shared_ptr<InstanceManager> &instance_manager);
    ActionCacheServicer() = delete;

    Status GetActionResult(ServerContext *ctx,
                           const GetActionResultRequest *request,
                           ActionResult *result) override;

    Status UpdateActionResult(ServerContext *ctx,
                              const UpdateActionResultRequest *request,
                              ActionResult *result) override;

  private:
    std::shared_ptr<InstanceManager> d_instance_manager;
};

class ExecutionServicer final : public Execution::Service {
  public:
    explicit ExecutionServicer(
        const std::shared_ptr<InstanceManager> &instance_manager);

    Status Execute(ServerContext *ctx, const ExecuteRequest *request,
                   ServerWriter<Operation> *writer) override;

    Status WaitExecution(ServerContext *ctx,
                         const WaitExecutionRequest *request,
                         ServerWriter<Operation> *writer) override;

  private:
    std::shared_ptr<InstanceManager> d_instance_manager;
};

class OperationsServicer final : public Operations::Service {
  public:
    explicit OperationsServicer(
        const std::shared_ptr<InstanceManager> &instance_manager);

    Status GetOperation(ServerContext *ctx, const GetOperationRequest *request,
                        Operation *response) override;

    Status ListOperations(ServerContext *ctx,
                          const ListOperationsRequest *request,
                          ListOperationsResponse *response) override;

    Status CancelOperation(ServerContext *ctx,
                           const CancelOperationRequest *request,
                           google::protobuf::Empty *response) override;

  private:
    std::shared_ptr<InstanceManager> d_instance_manager;
};

class Server {
    /* gRPC server that provide the buildbox-casd services:
     * `ContentAddressableStorage`, `ByteStream`, `Capabilities`,
     * `Fetch`, `Push` and `LocalContentAddressableStorage`.
     */
  public:
    Server(const std::shared_ptr<buildboxcommon::LocalCas> &cas_storage,
           const std::shared_ptr<FsLocalAssetStorage> &asset_storage,
           const std::shared_ptr<ActionStorage> &actioncache_storage,
           const std::shared_ptr<LocalExecutionScheduler> &executionScheduler =
               nullptr,
           const std::shared_ptr<ThreadPool> &digestThreadPool = nullptr,
           const std::shared_ptr<ThreadPool> &ioThreadPool = nullptr,
           const std::shared_ptr<std::unordered_set<std::string>>
               &allowedPreUnstageCommands = nullptr);

    Server(
        const std::shared_ptr<buildboxcommon::LocalCas> &cas_storage,
        const std::shared_ptr<FsLocalAssetStorage> &asset_storage,
        const std::shared_ptr<ActionStorage> &actioncache_storage,
        const std::shared_ptr<ThreadPool> &digestThreadPool,
        const std::shared_ptr<ThreadPool> &ioThreadPool,
        const std::shared_ptr<std::unordered_set<std::string>>
            &allowedPreUnstageCommands,
        const buildboxcommon::ConnectionOptions &cas_endpoint,
        const std::optional<buildboxcommon::ConnectionOptions> &ra_endpoint,
        const std::optional<buildboxcommon::ConnectionOptions> &ac_endpoint,
        const std::optional<buildboxcommon::ConnectionOptions>
            &execution_endpoint,
        const std::vector<std::string> &instance_names,
        const std::unordered_map<std::string, std::string> &instanceMappings,
        const bool read_only_remote = false,
        const int proxy_findmissingblobs_cache_ttl = 0,
        const int proxyAcValidationCacheTtl = 0);

    /* Create a new, restricted server for sandboxing with multiple instances,
     * sharing underlying storage, scheduler and thread pools.
     * instanceCloneMap maps cloned instance name -> base instance name.
     */
    Server(
        Server *base,
        const std::unordered_map<std::string, std::string> &instanceCloneMap,
        const std::string &rootPath, bool actionCacheUpdateEnabled,
        const std::shared_ptr<InodeCache> &inodeCache);

    /* Add port or UNIX socket to bind the server to. Must be called before
     * starting the server.
     */
    void addListeningPort(const std::string &addr);

    /* Set the minimum allowed time between a server receiving successive ping
     * frames without sending any data/header frame. Must be called before
     * starting the server.
     */
    void permitKeepaliveTime(int seconds);

    /* Start the server. */
    void start();

    /* Shutdown the server. */
    void shutdown();
    /* Shutdown the server with the given timeout (chiefly for tests) */
    void shutdown(gpr_timespec timeout);

    /* Block until the server shuts down. */
    void wait();

    inline std::shared_ptr<grpc::Service> remoteExecutionCasServicer()
    {
        return d_remote_execution_cas_servicer;
    }
    inline std::shared_ptr<grpc::Service> bytestreamServicer()
    {
        return d_cas_bytestream_servicer;
    }
    inline std::shared_ptr<grpc::Service> localCasServicer()
    {
        return d_local_cas_servicer;
    }
    inline std::shared_ptr<grpc::Service> capabilitiesServicer()
    {
        return d_capabilities_servicer;
    }
    inline std::shared_ptr<grpc::Service> assetFetchServicer()
    {
        return d_asset_fetch_servicer;
    }
    inline std::shared_ptr<grpc::Service> assetPushServicer()
    {
        return d_asset_push_servicer;
    }
    inline std::shared_ptr<grpc::Service> localAcServicer()
    {
        return d_actioncache_servicer;
    }
    inline std::shared_ptr<grpc::Service> executionServicer()
    {
        return d_execution_servicer;
    }
    inline std::shared_ptr<grpc::Service> operationsServicer()
    {
        return d_operations_servicer;
    }

    /* Create a proxy instance. */
    void addProxyInstance(
        const std::string &instanceName,
        const buildboxcommon::ConnectionOptions &casEndpoint,
        const std::optional<buildboxcommon::ConnectionOptions> &raEndpoint,
        const std::optional<buildboxcommon::ConnectionOptions> &acEndpoint,
        const std::optional<buildboxcommon::ConnectionOptions>
            &executionEndpoint,
        const bool readOnlyRemote = false,
        const int proxyFindmissingblobsCacheTtl = 0,
        const int proxyAcValidationCacheTtl = 0,
        std::optional<int> execHybridQueueLimit = std::nullopt,
        const bool enableFallbackToLocalExecution = false);

    /* Add a vector of proxy instances. */
    void addProxyInstances(
        const std::vector<std::string> &instanceNames,
        const std::unordered_map<std::string, std::string> &instanceMappings,
        const buildboxcommon::ConnectionOptions &casEndpoint,
        const std::optional<buildboxcommon::ConnectionOptions> &raEndpoint,
        const std::optional<buildboxcommon::ConnectionOptions> &acEndpoint,
        const std::optional<buildboxcommon::ConnectionOptions>
            &executionEndpoint,
        const bool readOnlyRemote = false,
        const int proxyFindmissingblobsCacheTtl = 0,
        const int proxyAcValidationCacheTtl = 0,
        std::optional<int> execHybridQueueLimit = std::nullopt,
        const bool enableFallbackToLocalExecution = false);

    /* Create a local server instance without proxy functionality. */
    void addLocalServerInstance(const std::string &instanceName);

    std::shared_ptr<InstanceManager> getInstanceManager();

  private:
    Server *d_baseServer = nullptr;

    grpc::ServerBuilder d_server_builder;
    std::shared_ptr<grpc::Server> d_server;

    std::shared_ptr<buildboxcommon::LocalCas> d_cas_storage;
    std::shared_ptr<FsLocalAssetStorage> d_asset_storage;
    std::shared_ptr<ActionStorage> d_actioncache_storage;
    std::unordered_map<buildboxcommon::StageTreeRequest::StagingMode,
                       std::shared_ptr<buildboxcommon::FileStager>>
        d_fileStagers;
    std::shared_ptr<LocalExecutionScheduler> d_executionScheduler;
    std::shared_ptr<InstanceManager> d_instance_manager;
    std::shared_ptr<ThreadPool> d_digestThreadPool;
    std::shared_ptr<ThreadPool> d_ioThreadPool;
    std::shared_ptr<std::unordered_set<std::string>>
        d_allowedPreUnstageCommands;

    const std::shared_ptr<CasRemoteExecutionServicer>
        d_remote_execution_cas_servicer;
    const std::shared_ptr<CasBytestreamServicer> d_cas_bytestream_servicer;
    const std::shared_ptr<LocalCasServicer> d_local_cas_servicer;
    const std::shared_ptr<CapabilitiesServicer> d_capabilities_servicer;

    const std::shared_ptr<AssetFetchServicer> d_asset_fetch_servicer;
    const std::shared_ptr<AssetPushServicer> d_asset_push_servicer;

    const std::shared_ptr<ActionCacheServicer> d_actioncache_servicer;

    const std::shared_ptr<ExecutionServicer> d_execution_servicer;
    const std::shared_ptr<OperationsServicer> d_operations_servicer;
};

} // namespace buildboxcasd

#endif // INCLUDED_BUILDBOXCASD_SERVER_H
