/*
 * SPDX-FileCopyrightText: Copyright (c) 2020-2026 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_BLOCK_MEMORY_POOL_HPP
#define NVIDIA_GXF_STD_BLOCK_MEMORY_POOL_HPP

#include <atomic>
#include <cstdint>
#include <memory>
#include <mutex>
#include <vector>

#include "gxf/std/allocator.hpp"
#include "gxf/std/resources.hpp"

namespace nvidia {
namespace gxf {

class FixedPoolUint64;

// A memory pools which provides a maximum number of equally sized blocks of
// memory.
// Supports stream-aware deallocation: when free_abi is called with a CUDA stream,
// the block is not immediately returned to the pool. Instead, a CUDA event is
// recorded on the stream, and the block is deferred until the event completes.
// This prevents data corruption when GPU operations are still reading from the buffer.
class BlockMemoryPool : public Allocator {
 public:
  BlockMemoryPool();
  ~BlockMemoryPool();

  gxf_result_t registerInterface(Registrar* registrar) override;
  gxf_result_t initialize() override;
  gxf_result_t is_available_abi(uint64_t size) override;
  gxf_result_t allocate_abi(uint64_t size, int32_t type, void** pointer) override;
  gxf_result_t free_abi(void* pointer) override;
  // Stream-aware free: defers block reuse until operations on the stream complete.
  // If stream is not null, records a CUDA event and defers the block.
  // If stream is null, falls back to immediate reuse (same as free_abi(pointer)).
  gxf_result_t free_abi(void* pointer, void* stream) override;
  gxf_result_t deinitialize() override;
  uint64_t block_size_abi() const override;

  // Returns the storage type of the memory blocks
  MemoryStorageType storage_type() const {
    return static_cast<MemoryStorageType>(storage_type_.get());
  }

  // Returns the total number of blocks
  uint64_t num_blocks() const { return num_blocks_.get(); }

 private:
  // Converts a pointer to a block index. Returns -1 if the pointer is invalid.
  int64_t pointerToIndex(void* pointer) const;

  // Reclaims deferred blocks whose CUDA events have completed.
  // Called internally before allocation to make completed blocks available.
  void reclaimCompletedDeferredBlocks();

  Parameter<int32_t> storage_type_;
  Parameter<uint64_t> block_size_;
  Parameter<uint64_t> num_blocks_;
  Resource<Handle<GPUDevice>> gpu_device_;

  void* pointer_;
  std::unique_ptr<FixedPoolUint64> stack_;
  std::mutex stack_mutex_;
  // Holds lifecycle stage of the allocator
  std::atomic<AllocatorStage> stage_{AllocatorStage::kUninitialized};
  int32_t dev_id_ = -1;

  // Tracks blocks that are deferred until their CUDA events complete.
  // Used for stream-aware deallocation to prevent data corruption.
  // Protected by stack_mutex_ (same as stack_ since they're accessed together).
  // RAII wrapper ensures cudaEventDestroy is called when the block is removed.
  struct DeferredBlock {
    uint64_t index = 0;     // Block index in the pool
    void* event = nullptr;  // CUDA event (stored as void* to avoid CUDA header in .hpp)

    DeferredBlock() = default;
    DeferredBlock(uint64_t idx, void* evt) : index(idx), event(evt) {}

    // Move constructor - takes ownership, nulls source to prevent double-destroy
    DeferredBlock(DeferredBlock&& other) noexcept : index(other.index), event(other.event) {
      other.event = nullptr;
    }

    // Move assignment - destroy old event, take ownership of new
    DeferredBlock& operator=(DeferredBlock&& other) noexcept;

    // Delete copy operations to prevent double-destroy (enforced at compile time)
    DeferredBlock(const DeferredBlock&) = delete;
    DeferredBlock& operator=(const DeferredBlock&) = delete;

    // Destructor - automatically destroys the CUDA event
    ~DeferredBlock();
  };
  std::vector<DeferredBlock> deferred_blocks_;
};

}  // namespace gxf
}  // namespace nvidia

#endif
