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

#pragma once

#include "device/pal/palkernel.hpp"
#include "amd_hsa_loader.hpp"

namespace amd {
namespace option {
class Options;
}  // namespace option
namespace hsa {
namespace loader {
class Loader;
class Executable;
class Context;
}  // namespace loader
}  // namespace hsa
}  // namespace amd

//! \namespace amd::pal PAL Device Implementation
namespace amd::pal {

/*! \addtogroup pal PAL Device Implementation
 *  @{
 */

using namespace amd::hsa::loader;
class Program;

class Segment {
 public:
  Segment();
  ~Segment();

  //! Allocates a segment
  bool alloc(Program& prog, amdgpu_hsa_elf_segment_t segment, size_t size, size_t align,
             bool zero);

  //! Copies data from host to the segment
  void copy(size_t offset, const void* src, size_t size);

  //! Segment freeze
  bool freeze(bool destroySysmem);

  //! Returns address for GPU access in the segment
  uint64_t gpuAddress(size_t offset) const { return gpuAccess_->vmAddress() + offset; }

  bool gpuAddressOffset(uint64_t offAddr, size_t* offset);

  //! Returns address for CPU access in the segment
  void* cpuAddress(size_t offset) const {
    return ((cpuAccess_ != nullptr) ? cpuAccess_->data() : cpuMem_) + offset;
  }

  void DestroyCpuAccess();

 private:
  Memory* gpuAccess_;  //!< GPU memory for segment access
  Memory* cpuAccess_;  //!< CPU memory for segment (backing store)
  address cpuMem_;     //!< CPU memory for segment without GPU direct access (backing store)
};

class PALHSALoaderContext final : public hsa::loader::Context {
 public:
  PALHSALoaderContext(pal::Program* program) : program_(program) {}

  virtual ~PALHSALoaderContext() {}

  hsa_isa_t IsaFromName(const char* name) override;

  bool IsaSupportedByAgent(hsa_agent_t agent, hsa_isa_t isa) override;

  void* SegmentAlloc(amdgpu_hsa_elf_segment_t segment, hsa_agent_t agent, size_t size, size_t align,
                     bool zero) override;

  bool SegmentCopy(amdgpu_hsa_elf_segment_t segment, hsa_agent_t agent, void* dst, size_t offset,
                   const void* src, size_t size) override;

  void SegmentFree(amdgpu_hsa_elf_segment_t segment, hsa_agent_t agent, void* seg,
                   size_t size = 0) override;

  void* SegmentAddress(amdgpu_hsa_elf_segment_t segment, hsa_agent_t agent, void* seg,
                       size_t offset) override;

  void* SegmentHostAddress(amdgpu_hsa_elf_segment_t segment, hsa_agent_t agent, void* seg,
                           size_t offset) override;

  bool SegmentFreeze(amdgpu_hsa_elf_segment_t segment, hsa_agent_t agent, void* seg,
                     size_t size) override;

  bool ImageExtensionSupported() override { return false; }

  hsa_status_t ImageCreate(hsa_agent_t agent, hsa_access_permission_t image_permission,
                           const hsa_ext_image_descriptor_t* image_descriptor,
                           const void* image_data, hsa_ext_image_t* image_handle) override {
    // not supported
    assert(false);
    return HSA_STATUS_ERROR;
  }

  hsa_status_t ImageDestroy(hsa_agent_t agent, hsa_ext_image_t image_handle) override {
    // not supported
    assert(false);
    return HSA_STATUS_ERROR;
  }

  hsa_status_t SamplerCreate(hsa_agent_t agent,
                             const hsa_ext_sampler_descriptor_t* sampler_descriptor,
                             hsa_ext_sampler_t* sampler_handle) override;

  //! All samplers are owned by pal program and are deleted in its destructor.
  hsa_status_t SamplerDestroy(hsa_agent_t agent, hsa_ext_sampler_t sampler_handle) override;

 private:
  PALHSALoaderContext(const PALHSALoaderContext& c);
  PALHSALoaderContext& operator=(const PALHSALoaderContext& c);

  pal::Program* program_;
};

//! \class pal program
class Program : public device::Program {
  friend class ClBinary;

 public:
  //! Default constructor
  Program(Device& device, amd::Program& owner);
  Program(NullDevice& device, amd::Program& owner);
  //! Default destructor
  virtual ~Program();

  void addGlobalStore(Memory* mem) { globalStores_.push_back(mem); }

  void setCodeObjects(Segment* seg, Memory* codeGpu, address codeCpu) {
    codeSegGpu_ = codeGpu;
    codeSegment_ = seg;
  }

  const std::vector<Memory*>& globalStores() const { return globalStores_; }

  //! Return a typecasted PAL null device.
  pal::NullDevice& palNullDevice() {
    return const_cast<pal::NullDevice&>(static_cast<const pal::NullDevice&>(device()));
  }

  //! Return a typecasted PAL device. The device must not be the NullDevice.
  pal::Device& palDevice() {
    return const_cast<pal::Device&>(static_cast<const pal::Device&>(device()));
  }

  //! Returns GPU kernel table
  const Memory* kernelTable() const { return kernels_; }

  //! Adds all kernels to the mem handle lists
  void fillResListWithKernels(VirtualGPU& gpu) const;

  //! Returns the maximum number of scratch regs used in the program
  uint maxScratchRegs() const { return maxScratchRegs_; }

  //! Returns the maximum number of VGPR(s) used in the program
  uint maxVgprs() const { return maxVgprs_; }

  //! Add internal static sampler
  void addSampler(Sampler* sampler) { staticSamplers_.push_back(sampler); }

  //! Returns TRUE if the program contains static samplers
  bool isStaticSampler() const { return (staticSamplers_.size() != 0); }

  //! Returns code segement on GPU
  const Memory& codeSegGpu() const { return *codeSegGpu_; }

  //! Returns CPU address for a kernel
  uint64_t findHostKernelAddress(uint64_t devAddr) const {
    return loader_->FindHostAddress(devAddr);
  }

  //! Get symbol by name
  amd::hsa::loader::Symbol* getSymbol(const char* symbol_name, const hsa_agent_t* agent) const {
    return executable_->GetSymbol(symbol_name, agent);
  }

  //! Returns API hash value of the program for RGP thread trace
  uint64_t ApiHash() const { return apiHash_; }

  //! Returns the load address of the trap handler
  uint64_t GetTrapHandlerAddress() const;

  virtual bool createKernels(void* binary, size_t binSize, bool useUniformWorkGroupSize,
                             bool internalKernel) override;

  virtual bool setKernels(void* binary, size_t binSize,
                          amd::Os::FileDesc fdesc = amd::Os::FDescInit(), size_t foffset = 0,
                          std::string uri = std::string()) override;

  virtual bool createBinary(amd::option::Options* options) override;

 protected:
  //! Destroys CPU allocations in the code segment
  void DestroySegmentCpuAccess() const {
    if (codeSegment_ != nullptr) {
      codeSegment_->DestroyCpuAccess();
    }
  }

  virtual bool defineGlobalVar(const char* name, void* dptr);
  virtual bool createGlobalVarObj(amd::Memory** amd_mem_obj, void** dptr, size_t* bytes,
                                  const char* globalName) const;

 private:
  //! Disable default copy constructor
  Program(const Program&);

  //! Disable operator=
  Program& operator=(const Program&);

 protected:
  //! Allocate kernel table
  bool allocKernelTable();

  void* rawBinary_;                    //!< Pointer to the raw binary
  std::vector<Memory*> globalStores_;  //!< Global memory for the program
  Memory* kernels_;                    //!< Table with kernel object pointers
  Memory* codeSegGpu_;                 //!< GPU memory with code objects
  Segment* codeSegment_;               //!< Pointer to the code segment for this program
  uint maxScratchRegs_;                //!< Maximum number of scratch regs used
                                       //!< in the program by individual kernel
  uint maxVgprs_;                      //!< Maximum number of VGPR(s) used
                                       //!< in the program by individual kernel
  uint64_t apiHash_ = 0;               //!< API hash value for RGP thread trace

  std::list<Sampler*> staticSamplers_;  //!< List od internal static samplers

  amd::hsa::loader::Loader* loader_;          //!< Loader object
  amd::hsa::loader::Executable* executable_;  //!< Executable for HSA Loader
  PALHSALoaderContext loaderContext_;         //!< Context for HSA Loader
};

/*@}*/  // namespace amd::pal
}  // namespace amd::pal