#ifndef HSA_RUNTIME_EXT_IMAGE_BLIT_KERNEL_H
#define HSA_RUNTIME_EXT_IMAGE_BLIT_KERNEL_H
#include <assert.h>
#include <atomic>
#include <mutex>
#include <unordered_map>
#include <vector>

#include "inc/hsa.h"
#include "resource.h"

namespace amd {
typedef struct BlitQueue {
  hsa_queue_t* queue_;
  volatile std::atomic<uint64_t> cached_index_;
} BlitQueue;

typedef struct BlitCodeInfo {
  uint64_t code_handle_;
  uint32_t group_segment_size_;
  uint32_t private_segment_size_;
} BlitCodeInfo;

class BlitKernel {
 public:
  typedef enum KernelOp {
    KERNEL_OP_COPY_IMAGE_TO_BUFFER = 0,
    KERNEL_OP_COPY_BUFFER_TO_IMAGE = 1,
    KERNEL_OP_COPY_IMAGE_DEFAULT = 2,
    KERNEL_OP_COPY_IMAGE_LINEAR_TO_STANDARD = 3,
    KERNEL_OP_COPY_IMAGE_STANDARD_TO_LINEAR = 4,
    KERNEL_OP_COPY_IMAGE_1DB = 5,
    KERNEL_OP_COPY_IMAGE_1DB_TO_REG = 6,
    KERNEL_OP_COPY_IMAGE_REG_TO_1DB = 7,
    KERNEL_OP_CLEAR_IMAGE = 8,
    KERNEL_OP_CLEAR_IMAGE_1DB = 9,
    KERNEL_OP_COUNT = 10
  } KernelOp;

  explicit BlitKernel();
  ~BlitKernel();

  hsa_status_t Initialize();

  hsa_status_t Cleanup();

  hsa_status_t BuildBlitCode(hsa_agent_t agent,
                             std::vector<BlitCodeInfo>& blit_code_catalog);

  hsa_status_t CopyBufferToImage(
      BlitQueue& blit_queue,
      const std::vector<BlitCodeInfo>& blit_code_catalog,
      const void* src_memory, size_t src_row_pitch, size_t src_slice_pitch,
      const Image& dst_image, const hsa_ext_image_region_t& image_region);

  hsa_status_t CopyImageToBuffer(
      BlitQueue& blit_queue,
      const std::vector<BlitCodeInfo>& blit_code_catalog,
      const Image& src_image, void* dst_memory, size_t dst_row_pitch,
      size_t dst_slice_pitch, const hsa_ext_image_region_t& image_region);

  hsa_status_t CopyImage(BlitQueue& blit_queue,
                         const std::vector<BlitCodeInfo>& blit_code_catalog,
                         const Image& dst_image, const Image& src_image,
                         const hsa_dim3_t& dst_origin,
                         const hsa_dim3_t& src_origin, const hsa_dim3_t size,
                         KernelOp copy_type);

  hsa_status_t FillImage(BlitQueue& blit_queue,
                         const std::vector<BlitCodeInfo>& blit_code_catalog,
                         const Image& image, const void* pattern,
                         const hsa_ext_image_region_t& region);

 private:

  hsa_status_t PopulateKernelCode(
      hsa_agent_t agent, hsa_executable_t executable,
      std::vector<BlitCodeInfo>& blit_code_catalog);

  inline void CalcBufferRowSlicePitchesInPixel(
      hsa_ext_image_geometry_t geometry, uint32_t element_size,
      const hsa_dim3_t& copy_size, size_t in_row_pitch_byte,
      size_t in_slice_pitch_byte, unsigned long* out_pitch_pixel);

  inline uint32_t GetDimSize(const Image& image);

  inline uint32_t GetNumChannel(const Image& image);

  inline uint32_t GetImageAccessType(const Image& image);

  void CalcWorkingSize(const Image& image, const hsa_dim3_t& range,
                       hsa_kernel_dispatch_packet_t& packet);

  void CalcWorkingSize(const Image& src_image, const Image& dst_image,
                       const hsa_dim3_t& range,
                       hsa_kernel_dispatch_packet_t& packet);

  hsa_status_t ConvertImage(const Image& original_image,
                            const Image** new_image);

  hsa_status_t LaunchKernel(BlitQueue& queue,
                            hsa_kernel_dispatch_packet_t& packet);

  // The kernels' name.
  static const char* kernel_name_[KERNEL_OP_COUNT];
  static const char* ocl_kernel_name_[KERNEL_OP_COUNT];

  // Mapping of ISA and kernel object.
  std::unordered_map<uint64_t, hsa_code_object_t> code_object_map_;

  // Mapping of ISA and kernel executable.
  std::unordered_map<uint64_t, hsa_executable_t> code_executable_map_;

  std::mutex lock_;

  DISALLOW_COPY_AND_ASSIGN(BlitKernel);

  // Get the patched code object
  hsa_status_t GetPatchedBlitObject(const char* agent_name, uint8_t** code_object_handle);
};
}  // namespace amd

#endif  // HSA_RUNTIME_EXT_IMAGE_BLIT_KERNEL_H