#pragma once

#include "llvm-project/clang/include/clang/Basic/LLVM.h"
#include "llvm/Support/JSON.h"

#include "onnx-mlir/Compiler/OMCompilerTypes.h"
#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"

namespace onnx_mlir {

struct MemEntry {
  size_t address;
  size_t size;
};

class PimMemory {
  llvm::SmallVector<std::pair<MemEntry, mlir::Value>, 32> memEntries;
  llvm::SmallDenseMap<mlir::Value, MemEntry, 32>& globalMemEntriesMap;

  size_t maxSize = 0; // 0 for unbounded memory
  size_t startAddress = 0;
  size_t minAlignment = 4;
  size_t firstAvailableAddress = 0;

  MemEntry* gatherMemEntry(mlir::Value value);
  void allocateMemoryForValue(mlir::Value value, MemEntry& memEntry);

public:
  PimMemory(llvm::SmallDenseMap<mlir::Value, MemEntry, 32>& globalMemEntriesMap)
  : globalMemEntriesMap(globalMemEntriesMap) {}

  void allocateHost(mlir::ModuleOp moduleOp, mlir::func::FuncOp funcOp);
  void allocateCore(mlir::Operation* op);

  size_t getFirstAvailableAddress() const { return firstAvailableAddress; }
  MemEntry getMemEntry(mlir::Value value) const;
};

class PimAcceleratorMemory {
public:
  llvm::SmallDenseMap<mlir::Value, MemEntry, 32> memEntriesMap;
  PimMemory hostMem;

private:
  llvm::SmallDenseMap<size_t, PimMemory> deviceMem;

public:
  PimAcceleratorMemory()
  : hostMem(memEntriesMap) {}

  PimMemory& getOrCreateDeviceMem(size_t id);

  size_t getValueAddress(mlir::Value value) const;
};

class PimCodeGen {
  PimAcceleratorMemory& memory;
  llvm::raw_fd_ostream& coreFileStream;

  static llvm::json::Object createEmptyOffset();
  void emitInstruction(llvm::json::Object instruction) const;

  void genSetRegisterImmediateUnsigned(size_t registerNumber, size_t immediate) const;
  void setupRd(size_t rdAddress, size_t rdOffset) const;
  void setupRdRs1(size_t rdAddress, size_t rdOffset, size_t rs1Address, size_t rs1Offset) const;
  void setupRdRs1Rs2(
    size_t rdAddress, size_t rdOffset, size_t rs1Address, size_t rs1Offset, size_t rs2Address, size_t rs2Offset) const;

  void emitMemCopyOp(mlir::StringRef opName,
                     size_t rdAddr,
                     size_t rdOffset,
                     size_t rs1Addr,
                     size_t rs1Offset,
                     size_t size,
                     mlir::StringRef sizeFieldName = "size") const;
  void emitCommunicationOp(mlir::StringRef opName, size_t bufferAddr, size_t coreId, size_t size) const;
  void emitMvmOp(size_t groupId, size_t rdAddr, size_t rdOffset, size_t rs1Addr, size_t rs1Offset) const;

public:
  PimCodeGen(PimAcceleratorMemory& memory, llvm::raw_fd_ostream& coreJson)
  : memory(memory), coreFileStream(coreJson) {}

  void codeGenLoadOp(pim::PimMemCopyHostToDevOp loadOp) const;
  void codeGenStoreOp(pim::PimMemCopyDevToHostOp storeOp) const;
  void codeGenLmvOp(pim::PimMemCopyOp lmvOp) const;

  void codeGenReceiveOp(pim::PimReceiveOp receiveOp) const;
  void codeGenSendOp(pim::PimSendOp sendOp) const;

  template <typename MVMTy>
  void codeGenMVMLikeOp(size_t mvmId, MVMTy mvmLikeOp, bool transposeMatrix);

  void codeGenVVAddOp(pim::PimVVAddOp vvaddOp) const;
  void codeGenVVSubOp(pim::PimVVSubOp vvsubOp) const;
  void codeGenVVMulOp(pim::PimVVMulOp vvmulOp) const;
  void codeGenVVMaxOp(pim::PimVVMaxOp vvmaxOp) const;
  void codeGenVVDMulOp(pim::PimVVDMulOp vvdmulOp) const;
  void codeGenVAvgOp(pim::PimVAvgOp vavgOp) const;
  void codeGenVReluOp(pim::PimVReluOp vreluOp) const;
  void codeGenVTanhOp(pim::PimVTanhOp vtanhOp) const;
  void codeGenVSigmOp(pim::PimVSigmOp vsigmOp) const;
  void codeGenTransposeOp(pim::PimTransposeOp transposeOp) const;
};

OnnxMlirCompilerErrorCodes compileToPimJson(mlir::ModuleOp& moduleOpRef, std::string& outputDirName);

} // namespace onnx_mlir
