/*
 * SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: Apache-2.0
 *
 * 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 NVIDIA_GXF_STD_RESOURCES_HPP_
#define NVIDIA_GXF_STD_RESOURCES_HPP_

#include <map>
#include <string>
#include <vector>

#include "gxf/core/component.hpp"
#include "gxf/core/entity.hpp"

namespace nvidia {
namespace gxf {

// A base type for all kinds of resources.
// This is used to filter out resource Component from an entity
class ResourceBase : public Component {};

class ThreadPool : public ResourceBase {
 public:
  struct Thread {
    gxf_uid_t uid;
    gxf_uid_t cpu_thread_cid;
    // Scalability: std::thread
  };
  gxf_result_t registerInterface(Registrar* registrar) override;
  gxf_result_t initialize() override;
  Expected<gxf_uid_t> addThread(gxf_uid_t key, gxf_uid_t cpu_thread_cid = kUnspecifiedUid);
  Expected<gxf_uid_t> addThread(gxf_uid_t key, const std::vector<uint32_t>& pin_cores);
  const Expected<Thread> getThread(gxf_uid_t key) const;
  const std::map<gxf_uid_t, Thread>& get() const;
  int64_t size() const;
  int64_t priority() const;

  std::vector<uint32_t> coresToPin() const {
    const auto maybe_pin_cores = pin_cores_.try_get();
    if (maybe_pin_cores) {
      return maybe_pin_cores.value();
    }
    return {};
  }

 private:
  Parameter<int64_t> initial_size_;
  Parameter<int64_t> priority_;
  // CPU core IDs to pin the worker threads to (empty means no core pinning)
  Parameter<std::vector<uint32_t>> pin_cores_;
  std::map<gxf_uid_t, Thread> thread_pool_;
  static inline gxf_uid_t next_uid_ = 0;
};

class GPUDevice : public ResourceBase {
 public:
  gxf_result_t registerInterface(Registrar* registrar) override;
  int32_t device_id() const { return dev_id_; }

 private:
  Parameter<int32_t> dev_id_;
};

}  // namespace gxf
}  // namespace nvidia

#endif  // NVIDIA_GXF_STD_RESOURCES_HPP_
