/**
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
 * CANN Open Software License Agreement Version 2.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */

#ifndef FLOW_FUNC_THREAD_POOL_H
#define FLOW_FUNC_THREAD_POOL_H

#include <functional>
#include <thread>
#include <vector>
#include <string>
#include "ascend_hal.h"
#include "common/count_down_latch.h"

namespace FlowFunc {
class DeviceCpuInfo {
 public:
  int32_t Init(uint32_t device_id);

  uint32_t GetAicpuPhysIndex(uint32_t logic_idx) const {
    if (logic_idx < aicpu_indexes_.size()) {
      return aicpu_indexes_[logic_idx];
    } else if (aicpu_indexes_.empty()) {
      return 0;
    } else {
      return aicpu_indexes_[0];
    }
  }

  uint32_t GetAicpuNum() const {
    return aicpu_num_;
  }

 private:
  static int32_t GetCpuNum(uint32_t device_id, DEV_MODULE_TYPE model_type, uint32_t &num);

  // aicpu core num
  uint32_t aicpu_num_;
  std::vector<uint32_t> aicpu_indexes_;
};

class FlowFuncThreadPool {
 public:
  explicit FlowFuncThreadPool(std::string thread_name_prefix) : thread_name_prefix_(std::move(thread_name_prefix)) {}

  ~FlowFuncThreadPool();

  int32_t Init(const std::function<void(uint32_t)> &thread_func, uint32_t default_thread_num);

  void WaitForStop() noexcept;

  void WaitAllThreadReady();

  void ThreadReady(uint32_t thread_index);

  void ThreadAbnormal(uint32_t thread_index);

  uint32_t GetThreadNum() const;

  static int32_t ThreadSecureCompute();

 private:
  void SetThreadName(uint32_t thread_idx) const;

  static int32_t WriteTidForAffinity();

  void BindAicpu(uint32_t thread_idx) const;

  std::vector<std::thread> workers_;

  std::string thread_name_prefix_;

  DeviceCpuInfo device_cpu_info_;
  CountDownLatch wait_thread_ready_latch_;
};
}  // namespace FlowFunc
#endif  // FLOW_FUNC_THREAD_POOL_H