/*
 * 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_MEMORY_BUFFER_HPP_
#define NVIDIA_GXF_STD_MEMORY_BUFFER_HPP_

#include <functional>
#include <memory>
#include <utility>

#include "common/byte.hpp"
#include "gxf/core/expected.hpp"
#include "gxf/core/handle.hpp"
#include "gxf/std/allocator.hpp"

namespace nvidia {
namespace gxf {

class MemoryBuffer {
 public:
  MemoryBuffer() = default;
  MemoryBuffer(const MemoryBuffer&) = delete;
  MemoryBuffer& operator=(const MemoryBuffer&) = delete;

  MemoryBuffer(MemoryBuffer&& other) { *this = std::move(other); }

  MemoryBuffer& operator=(MemoryBuffer&& other) {
    if (this != &other) {
      freeBuffer();  // Release existing resources first (ignore return value, like destructor)
      size_ = other.size_;
      storage_type_ = other.storage_type_;
      pointer_ = other.pointer_;
      release_func_ = std::move(other.release_func_);
      stream_state_ = std::move(other.stream_state_);

      other.pointer_ = nullptr;
      other.release_func_ = nullptr;
      other.stream_state_ = nullptr;
    }

    return *this;
  }

  // Type of the callback function to release memory passed to the MemoryBuffer
  // using the wrapMemory method
  using release_function_t = std::function<Expected<void>(void* pointer)>;

  Expected<void> freeBuffer() {
    if (pointer_) {
      if (release_func_) {
        const Expected<void> result = release_func_(pointer_);
        if (!result) {
          return ForwardError(result);
        }
        release_func_ = nullptr;
      }

      pointer_ = nullptr;
      size_ = 0;
      stream_state_ = nullptr;
    }

    return Success;
  }

  virtual ~MemoryBuffer() { freeBuffer(); }

  // Resizes the buffer, allocating memory from the given allocator.
  // The optional stream parameter specifies the CUDA stream for stream-aware deallocation.
  // If stream is provided, it will be used when freeing the memory to enable deferred
  // reuse in allocators like BlockMemoryPool.
  // Note: The stream can be updated later via setStream() to reflect the last-used stream.
  Expected<void> resize(Handle<Allocator> allocator, uint64_t size, MemoryStorageType storage_type,
                        void* stream = nullptr) {
    const auto result = freeBuffer();
    if (!result) {
      GXF_LOG_ERROR("Failed to free memory. Error code: %s", GxfResultStr(result.error()));
      return ForwardError(result);
    }

    const auto maybe = allocator->allocate(size, storage_type);
    if (!maybe) {
      GXF_LOG_ERROR("%s Failed to allocate %ld size of memory of type %d. Error code: %s",
                    allocator->name(),
                    size,
                    static_cast<int32_t>(storage_type),
                    GxfResultStr(maybe.error()));
      return ForwardError(maybe);
    }

    storage_type_ = storage_type;
    pointer_ = maybe.value();
    size_ = size;

    // Create shared stream state that can be updated via setStream() and accessed by the
    // release function. This allows stream-aware deallocation even after the MemoryBuffer
    // is moved, since the release function captures the shared_ptr.
    stream_state_ = std::make_shared<StreamState>();
    stream_state_->stream = stream;

    // Create release function that captures the allocator and stream state.
    // The allocator handle is captured by value, and the stream state is captured
    // by shared_ptr so setStream() updates are visible to the release function.
    release_func_ = [allocator, stream_state = stream_state_](void* ptr) -> Expected<void> {
      void* stream = stream_state ? stream_state->stream : nullptr;
      return allocator->free(reinterpret_cast<byte*>(ptr), stream);
    };

    return Success;
  }

  // Wrap existing memory inside the MemoryBuffer. A callback function of type
  // release_function_t may be passed that will be called when the MemoryBuffer
  // wants to release the memory. If release_func is null/empty, the memory is
  // assumed to be externally managed and will not be released.
  Expected<void> wrapMemory(void* pointer, uint64_t size, MemoryStorageType storage_type,
                            release_function_t release_func) {
    if (pointer != this->pointer()) {
      const auto result = freeBuffer();
      if (!result) {
        return ForwardError(result);
      }

      pointer_ = reinterpret_cast<byte*>(pointer);
    }
    storage_type_ = storage_type;
    size_ = size;
    release_func_ = std::move(release_func);

    // Create stream state so setStream() can be used after wrapMemory().
    // Only create if not already present (allows re-wrapping same pointer with new release_func).
    if (!stream_state_) {
      stream_state_ = std::make_shared<StreamState>();
    }

    return Success;
  }

  // The type of memory where the data is stored.
  MemoryStorageType storage_type() const { return storage_type_; }

  // Raw pointer to the first byte of elements stored in the buffer.
  byte* pointer() const { return pointer_; }

  // Size of buffer contents in bytes
  uint64_t size() const { return size_; }

  // Sets the CUDA stream that last operated on this buffer's data.
  // This stream will be passed to the allocator during deallocation to enable
  // stream-aware memory reuse (e.g., deferred block reuse in BlockMemoryPool).
  // The stream parameter is void* to avoid CUDA header dependencies; pass a
  // cudaStream_t cast to void*.
  // Note: This updates shared state that is captured by the release function,
  // so the updated stream will be used even if this MemoryBuffer is moved.
  // For buffers created via wrapMemory(), the release function may not use the
  // stream unless it explicitly captures the stream state.
  void setStream(void* stream) {
    if (stream_state_) {
      stream_state_->stream = stream;
    }
  }

  // Gets the CUDA stream that last operated on this buffer's data.
  // Returns nullptr if no stream has been set.
  // The return value should be cast to cudaStream_t for use with CUDA APIs.
  void* stream() const {
    return stream_state_ ? stream_state_->stream : nullptr;
  }

  // Shared state for the CUDA stream, allowing setStream() to update the stream
  // that will be used by the release function even after the buffer is moved.
  struct StreamState {
    void* stream = nullptr;
  };

  // Returns the shared stream state. This allows callers (like Tensor::initializeDLContext)
  // to capture the stream state in a release function so that stream updates via setStream()
  // can be propagated to inner buffers during deallocation.
  std::shared_ptr<StreamState> streamState() const { return stream_state_; }

 protected:

  uint64_t size_ = 0;
  byte* pointer_ = nullptr;
  MemoryStorageType storage_type_ = MemoryStorageType::kHost;
  release_function_t release_func_ = nullptr;
  // Shared stream state that can be updated via setStream() and is captured by
  // the release function created in resize(). This enables stream-aware deallocation
  // with dynamically updatable streams.
  std::shared_ptr<StreamState> stream_state_;
};

}  // namespace gxf
}  // namespace nvidia

#endif  // NVIDIA_GXF_STD_MEMORY_BUFFER_HPP_
