/**
 * Copyright 2020 Huawei Technologies Co., Ltd
 *
 * 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 GE_GRAPH_LOAD_NEW_MODEL_MANAGER_CPU_QUEUE_SCHEDULE_H_
#define GE_GRAPH_LOAD_NEW_MODEL_MANAGER_CPU_QUEUE_SCHEDULE_H_

#include <cstdint>
#include <vector>

#include "framework/common/ge_inner_error_codes.h"
#include "graph/load/model_manager/task_info/task_info.h"
#include "graph/load/model_manager/zero_copy_offset.h"
#include "runtime/kernel.h"

namespace ge {
// For AICPU task "modelPrepare" / "modelPostpare"
struct AicpuPareInfo {
  uint32_t aicpu_info_size;
  uint32_t model_id;
  uint32_t input_addr_num;
  uint64_t input_addr_list;
  uint64_t input_index_list;
  uint32_t output_addr_num;
  uint64_t output_addr_list;
  uint64_t output_index_list;
  uint32_t output_num;
  uint64_t output_size_list;
  uint32_t in_queue_num;
  uint64_t in_queueid_list;
  uint32_t out_queue_num;
  uint64_t out_queueid_list;
  uint64_t mbufptr_list;
};

///
/// @ingroup ge
/// @brief CpuTask base, inherit from TaskInfo used for manage.
///
class CpuTaskInfo : public TaskInfo {
 public:
  explicit CpuTaskInfo(rtStream_t stream);
  ~CpuTaskInfo() override;

 protected:
  void *args_;
  uint32_t args_size_;
};

class CpuTaskModelPrepare : public CpuTaskInfo {
 public:
  explicit CpuTaskModelPrepare(rtStream_t stream) : CpuTaskInfo(stream) {}
  ~CpuTaskModelPrepare() override;

  Status Init(const domi::TaskDef &task_def, DavinciModel *davinci_model) override { return SUCCESS; }
  Status Init(const vector<uint32_t> &input_queue_ids, const vector<uint32_t> &output_queue_ids,
              const map<uint32_t, ZeroCopyOffset> &inside_addrs, const map<uint32_t, ZeroCopyOffset> &outside_addrs,
              uintptr_t &out_mbuf);

  Status Distribute() override;

 private:
  Status GenerateOutSizeAddr(const map<uint32_t, ZeroCopyOffset> &outside_addrs, void *&output_size_list_addr);
  Status GenerateCpuAddr(const map<uint32_t, ZeroCopyOffset> &node_addrs, void *&data_list_addr, void *&index_list_addr,
                         uint32_t &num);
  void *input_list_addr_ = nullptr;
  void *input_index_list_addr_ = nullptr;
  void *output_list_addr_ = nullptr;
  void *output_index_list_addr_ = nullptr;
  void *output_size_list_addr_ = nullptr;
  void *queue_id_list_addr_ = nullptr;
  void *mbufptr_list_ = nullptr;
};

class CpuTaskModelPostpare : public CpuTaskInfo {
 public:
  explicit CpuTaskModelPostpare(rtStream_t stream) : CpuTaskInfo(stream) {}
  ~CpuTaskModelPostpare() override;

  Status Init(const domi::TaskDef &task_def, DavinciModel *davinci_model) override { return SUCCESS; }
  Status Init(uint32_t model_id, const vector<uint32_t> &output_queue_ids, uintptr_t out_mbuf);

  Status Distribute() override;

 private:
  void *queue_id_list_addr_ = nullptr;
};

///
/// @ingroup ge
/// @brief definiteness queue schedule, active entry stream.
///
class CpuTaskActiveEntry : public CpuTaskInfo {
 public:
  explicit CpuTaskActiveEntry(rtStream_t stream) : CpuTaskInfo(stream), active_stream_(nullptr) {}
  ~CpuTaskActiveEntry() override {}

  Status Init(const domi::TaskDef &task_def, DavinciModel *davinci_model) override { return SUCCESS; }
  Status Init(rtStream_t stream);

  Status Distribute() override;

 private:
  rtStream_t active_stream_;
};

///
/// @ingroup ge
/// @brief definiteness queue schedule, wait for end graph.
///
class CpuTaskWaitEndGraph : public CpuTaskInfo {
 public:
  explicit CpuTaskWaitEndGraph(rtStream_t stream) : CpuTaskInfo(stream) {}
  ~CpuTaskWaitEndGraph() override {}

  Status Init(const domi::TaskDef &task_def, DavinciModel *davinci_model) override { return SUCCESS; }
  Status Init(uint32_t model_id);

  Status Distribute() override;
};
}  // namespace ge
#endif  // GE_GRAPH_LOAD_NEW_MODEL_MANAGER_CPU_QUEUE_SCHEDULE_H_