/*
 * Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
 *
 * SPDX-License-Identifier: MIT
 */

#pragma once

#include <functional>
#include "top.hpp"
#include "device/device.hpp"
#include "object.hpp"
#include "commandqueue.hpp"

namespace amd {

class HeapBlock;
class VmHeap;
class VmHeapArray;

class HeapBlock {
 public:
  friend VmHeap;
  //! Constructor
  HeapBlock(VmHeap* owner = nullptr,  //!< VmHeap object that owns this heap block
            size_t size = 0,          //!< Heap block size for allocation
            size_t offset = 0)        //!< Heap block offset
      : owner_(owner), size_(size), offset_(offset), next_(nullptr), prev_(nullptr), busy_(false) {}

  //! Destructor does some sanity checks
  ~HeapBlock() { assert(!busy_ && "The blocked must be destroyed explicitly!"); }

  //! Gets the offset
  size_t Offset() const { return offset_; }

 private:
  HeapBlock() = delete;
  HeapBlock(const HeapBlock&) = delete;
  HeapBlock& operator=(const HeapBlock&) = delete;

  VmHeap* owner_;    //!< Heap that owns this block
  size_t size_;      //!< Size of the block in bytes
  size_t offset_;    //!< Offset of this block in the heap
  HeapBlock* next_;  //!< Next block on the list, or nullptr
  HeapBlock* prev_;  //!< Previous block on the list, or nullptr
  bool busy_;        //!< True if the block is in use
};

class VmHeap {
 public:
  friend VmHeapArray;
  static const size_t kChunkSize = 128 * Mi;  //!< Chunk size, must be power of 2
  static const size_t kMinBlockAlignment = 256;
  typedef std::function<amd::HostQueue&()> GetQueueFunc;

  VmHeap(Device* device,         //!< GPU device object
         GetQueueFunc get_queue  //!< Function to retrieve a map queue
         )
      : VmHeap(device, device->info().globalMemSize_ / 8, kChunkSize, get_queue) {}

  VmHeap(Device* device,         //!< GPU device object
         size_t va_size,         //!< The size of the allocated heap (bytes).Virtual address space
         size_t chunk_size,      //!< The size of single chunk for physical memory growth
         GetQueueFunc get_queue  //!< Function to retrieve a map/unmap queue
  );

  //! Heap destructor
  virtual ~VmHeap();

  //! Returns a pointer to the allocated device memory from a heap
  address Alloc(size_t size  //! The allocation size
  );

  //! Release memory back to the VM heap
  void Free(amd::Memory* memory);

  //! Unmaps freed memory based on the threshold
  void TrimPhysMemory(size_t unmap_threshold);

  //! Enable memory unmap threashold (default 0 unmap always)
  void SetUnmapThreshold(uint64_t threshold) { unmap_threshold_ = threshold; }

  //! Returns mapped memory size (total allocated physical memory)
  uint64_t MappedSize() const { return mapped_size_; }

  //! Returns mapped memory size (allocated physical memory) without actual allocations
  uint64_t FreeMappedSize() const { return mapped_size_ - (va_size_ - free_size_); }

  //! Returns true if the address is in the range of this heap
  bool InRange(void* addr) {
    return ((addr >= base_address_) && (addr <= (base_address_ + va_size_))) ? true : false;
  }

 private:
  VmHeap() = delete;
  VmHeap(const VmHeap&) = delete;
  VmHeap& operator=(const VmHeap&) = delete;

  //! Ceates heap object. Reserves virtual address range for the heap operation
  bool Create();

  //! Reseves address range for memory allocations
  address ReserveAddressRange(address start, size_t size, size_t alignment);

  //! Releases address range specified by the address
  bool ReleaseAddressRange(void* addr);

  //! Commits actual physical memory on the specified address
  bool CommitMemory(void* addr, size_t size);

  //! Uncommits physical memory from the spcified address
  bool UncommitMemory(void* addr, size_t size);

  HeapBlock* AllocBlock(size_t size  //! The allocation size
  );

  //! Release memory back to a heap
  void FreeBlock(HeapBlock* blk);

  //! Insert a block into a list
  void InsertBlock(HeapBlock** list, HeapBlock* node);

  //! Merge a block into a list
  void MergeBlock(HeapBlock** list, HeapBlock* node);

  //! Remove a block from a list
  void DetachBlock(HeapBlock** list, HeapBlock* node);

  //! Splits a block into two pieces
  HeapBlock* SplitBlock(HeapBlock* node, size_t size);

  //! Maps physical memory into specified address space
  bool MapPhysMemory(size_t offset, size_t size);

  //! Unmaps physical memory from the specified address
  void UnmapPhysMemory(size_t offset, size_t size);

  //! Join two blocks, transferring the size of the second into the first and deleting the second
  void Join2Blocks(HeapBlock* first, HeapBlock* second) const;

  //! Returns a queue for VM map/unmap operations
  amd::HostQueue& GetVmQueue() const { return get_vm_queue_(); }

  address base_address_ = nullptr;      //!< GPU virtual address base of the heap
  amd::Memory* base_memory_ = nullptr;  //!< VA space base object, used in the view creation
  HeapBlock* free_list_ = nullptr;      //!< Head block for free list
  HeapBlock* busy_list_ = nullptr;      //!< Head block for busy list
  size_t free_size_ = 0;                //!< Total free size of the heap (both mapped and unmapped)
  size_t va_size_ = 0;                  //!< Heap virtual address space size
  size_t block_alignment_ = 1;          //!< Size of an allocation page
  size_t chunk_size_ = 0;               //!< Chunk size (min physical allocation for the growth)
  uint64_t unmap_threshold_ = 0;        //!< Unmap threshold in bytes,used to release phys memory
  uint64_t mapped_size_ = 0;            //!< Size of mapped memory
  uint64_t max_mapped_size_ = 0;        //!< Max size of mapped memory in this heap
  bool created_ = false;                //!< Used for deferred VM heap allocation
  std::recursive_mutex lock_;           //!< Lock to serialise heap accesses
  Device* device_;                      //!< Device that owns this heap
  GetQueueFunc get_vm_queue_;           //!< Queue for VM operations

  std::vector<bool> mapped_mem_;  //!< A map of mapped memory, the size is total_size/chunk_size
};

//! Implements an array of vm heaps of different sizes for more efficient management
class VmHeapArray {
 public:
  VmHeapArray(Device* device,                 //!< GPU device object
              VmHeap::GetQueueFunc get_queue  //!< Function to retrieve a map queue
              )
      : heap0_(device, device->info().globalMemSize_ / 8, VmHeap::kChunkSize, get_queue),
        heap1_(device, device->info().globalMemSize_ / 4, VmHeap::kChunkSize, get_queue),
        heap2_(device, device->info().globalMemSize_ * 5 / 8, VmHeap::kChunkSize, get_queue),
        heap3_(device, device->info().globalMemSize_, VmHeap::kChunkSize, get_queue),
        device_(device) {}

  //! Returns a pointer to the allocated device memory from a heap
  address Alloc(size_t size  //! The allocation size
  );

  //! Release memory back to the VM heap
  void Free(amd::Memory* memory);

  //! Unmaps freed memory based on the threshold
  void TrimPhysMemory(size_t unmap_threshold);

  //! Enable memory unmap threashold (default 0 unmap always)
  void SetUnmapThreshold(uint64_t threshold);

  //! Returns mapped memory size (total allocated physical memory)
  uint64_t MappedSize() const;

  //! Returns mapped memory size (allocated physical memory) without actual allocations
  uint64_t FreeMappedSize() const;

  //! Returns the maximum mapped memory size
  uint64_t MaxMappedSize() const;

  //! Returns the maximum mapped memory size
  void ResetMaxMappedSize();

 private:
  VmHeapArray() = delete;
  VmHeapArray(const VmHeapArray&) = delete;
  VmHeapArray& operator=(const VmHeapArray&) = delete;

  static const uint32_t kMaxArraySize = 4;  //!< The number of vm heap in the array
  // @note: gcc10.2 or lower wrongly uses copy constructor in the initialization
  // of VmHeap array of objects. Hence, use an array of VmHeap pointers for now
  VmHeap* vm_heaps_[kMaxArraySize] = {&heap0_, &heap1_, &heap2_, &heap3_};  //!< The array of heaps
  VmHeap heap0_;
  VmHeap heap1_;
  VmHeap heap2_;
  VmHeap heap3_;

  uint64_t unmap_threshold_ = 0;  //!< Unmap threshold in bytes,used to release phys memory
  Device* device_;                //!< Device that owns this heap
};

}  // namespace amd