/*******************************************************************************
* Copyright 上海赜睿信息科技有限公司(Zilliz) - All Rights Reserved
* Unauthorized copying of this file, via any medium is strictly prohibited.
* Proprietary and confidential.
******************************************************************************/
#pragma once

#include "utils/BlockingQueue.h"
#include "status.grpc.pb.h"
#include "status.pb.h"

#include <map>
#include <vector>
#include <thread>

namespace zilliz {
namespace milvus {
namespace server {
namespace grpc {

class GrpcBaseTask {
protected:
    GrpcBaseTask(const std::string &task_group, bool async = false);

    virtual ~GrpcBaseTask();

public:
    ErrorCode Execute();

    void Done();

    ErrorCode WaitToFinish();

    std::string TaskGroup() const { return task_group_; }

    ErrorCode ErrorID() const { return error_code_; }

    std::string ErrorMsg() const { return error_msg_; }

    bool IsAsync() const { return async_; }

protected:
    virtual ErrorCode OnExecute() = 0;

    ErrorCode SetError(ErrorCode error_code, const std::string &msg);

protected:
    mutable std::mutex finish_mtx_;
    std::condition_variable finish_cond_;

    std::string task_group_;
    bool async_;
    bool done_;
    ErrorCode error_code_;
    std::string error_msg_;
};

using BaseTaskPtr = std::shared_ptr<GrpcBaseTask>;
using TaskQueue = BlockingQueue<BaseTaskPtr>;
using TaskQueuePtr = std::shared_ptr<TaskQueue>;
using ThreadPtr = std::shared_ptr<std::thread>;

class GrpcRequestScheduler {
public:
    static GrpcRequestScheduler &GetInstance() {
        static GrpcRequestScheduler scheduler;
        return scheduler;
    }

    void Start();

    void Stop();

    ErrorCode ExecuteTask(const BaseTaskPtr &task_ptr);

    static void ExecTask(BaseTaskPtr &task_ptr, ::milvus::grpc::Status *grpc_status);

protected:
    GrpcRequestScheduler();

    virtual ~GrpcRequestScheduler();

    void TakeTaskToExecute(TaskQueuePtr task_queue);

    ErrorCode PutTaskToQueue(const BaseTaskPtr &task_ptr);

private:
    mutable std::mutex queue_mtx_;

    std::map<std::string, TaskQueuePtr> task_groups_;

    std::vector<ThreadPtr> execute_threads_;

    bool stopped_;
};

}
}
}
}
