#pragma once

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

#include "Common/ValueMap.hpp"
#include "src/Accelerators/PIM/Compiler/PimCompilerUtils.hpp"
#include "src/Accelerators/PIM/Dialect/PIM/PimOps.hpp"

namespace onnx_mlir {

using namespace llvm;
using namespace mlir;
using Value = mlir::Value;
using Type = mlir::Type;
using FunctionType = mlir::FunctionType;

struct MemEntry {
  size_t address;
  size_t size;
};

class PimMemory {
  SmallVector<std::pair<MemEntry, Value>, 32> memEntries;
  SmallDenseMap<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(Value value);
  void allocateMemoryForValue(Value value, MemEntry& memEntry);

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

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

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

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

private:
  SmallDenseMap<size_t, PimMemory> deviceMem;

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

  PimMemory getOrCreateDeviceMem(size_t id);

  size_t getValueAddress(Value value) const;
};

class PimCodeGen {
  PimAcceleratorMemory& memory;
  raw_fd_ostream& coreFileStream;

  static json::Object createEmptyOffset();
  void emitInstruction(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(StringRef opName, size_t rdAddr, size_t rdOffset, size_t rs1Addr, size_t rs1Offset, size_t size) const;
  void emitCommunicationOp(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, 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 codeGenVAddOp(pim::PimVAddOp vaddOp) const;
  void codeGenVMaxOp(pim::PimVMaxOp vmaxOp) const;
  void codeGenVReluOp(pim::PimVReluOp vreluOp) const;
  void codeGenApplyFiltersOp(pim::PimApplyFiltersOp applyFiltersOp) const;
};

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

} // namespace onnx_mlir