#pragma once

#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Endian.h"
#include "llvm/Support/JSON.h"
#include "llvm/Support/raw_ostream.h"

#include <array>
#include <cassert>
#include <limits>

namespace onnx_mlir::pim_binary {

inline constexpr char kMagic[4] = {'P', 'I', 'M', 'B'};
inline constexpr uint32_t kVersion = 1;
inline constexpr uint64_t kCountOffset = 8;
inline constexpr size_t kHeaderSize = 12;
inline constexpr size_t kRecordSize = 20;

enum class Opcode : uint32_t {
  nop = 0,
  sldi = 1,
  sld = 2,
  sadd = 3,
  ssub = 4,
  smul = 5,
  saddi = 6,
  smuli = 7,
  setbw = 8,
  mvmul = 9,
  vvadd = 10,
  vvsub = 11,
  vvmul = 12,
  vvdmul = 13,
  vvmax = 14,
  vvsll = 15,
  vvsra = 16,
  vavg = 17,
  vrelu = 18,
  vtanh = 19,
  vsigm = 20,
  vsoftmax = 21,
  vmv = 22,
  vrsu = 23,
  vrsl = 24,
  ld = 25,
  st = 26,
  lldi = 27,
  lmv = 28,
  send = 29,
  recv = 30,
  wait = 31,
  sync = 32,
};

struct InstructionRecord {
  Opcode opcode = Opcode::nop;
  uint8_t rd = 0;
  uint8_t r1 = 0;
  int32_t r2OrImm = 0;
  int32_t generic1 = 0;
  int32_t generic2 = 0;
  int32_t generic3 = 0;
  uint8_t flags = 0;
};

inline void writeUint32LE(llvm::raw_ostream& os, uint32_t value) {
  std::array<char, sizeof(uint32_t)> bytes;
  llvm::support::endian::write32le(bytes.data(), value);
  os.write(bytes.data(), bytes.size());
}

inline void writeInt32LE(llvm::raw_ostream& os, int32_t value) { writeUint32LE(os, static_cast<uint32_t>(value)); }

inline void writeHeader(llvm::raw_ostream& os) {
  os.write(kMagic, sizeof(kMagic));
  writeUint32LE(os, kVersion);
  writeUint32LE(os, 0);
}

inline void patchInstructionCount(llvm::raw_pwrite_stream& os, uint32_t instructionCount) {
  std::array<char, sizeof(uint32_t)> bytes;
  llvm::support::endian::write32le(bytes.data(), instructionCount);
  os.pwrite(bytes.data(), bytes.size(), kCountOffset);
}

inline void writeInstructionRecord(llvm::raw_ostream& os, const InstructionRecord& record) {
  os << static_cast<char>(static_cast<uint8_t>(record.opcode));
  os << static_cast<char>(record.rd);
  os << static_cast<char>(record.r1);
  os << static_cast<char>(record.flags);
  writeInt32LE(os, record.r2OrImm);
  writeInt32LE(os, record.generic1);
  writeInt32LE(os, record.generic2);
  writeInt32LE(os, record.generic3);
}

inline int32_t toI32(int64_t value) {
  assert(value >= std::numeric_limits<int32_t>::min() && value <= std::numeric_limits<int32_t>::max()
         && "PIM binary field out of int32 range");
  return static_cast<int32_t>(value);
}

inline uint8_t toU8(int64_t value) {
  assert(value >= 0 && value <= std::numeric_limits<uint8_t>::max() && "PIM binary field out of uint8 range");
  return static_cast<uint8_t>(value);
}

inline int32_t getOptionalInt(const llvm::json::Object& object, llvm::StringRef key, int32_t defaultValue = 0) {
  if (std::optional<int64_t> value = object.getInteger(key))
    return toI32(*value);
  return defaultValue;
}

inline Opcode opcodeFromString(llvm::StringRef opName) {
  if (opName == "nop")
    return Opcode::nop;
  if (opName == "sldi")
    return Opcode::sldi;
  if (opName == "sld")
    return Opcode::sld;
  if (opName == "sadd")
    return Opcode::sadd;
  if (opName == "ssub")
    return Opcode::ssub;
  if (opName == "smul")
    return Opcode::smul;
  if (opName == "saddi")
    return Opcode::saddi;
  if (opName == "smuli")
    return Opcode::smuli;
  if (opName == "setbw")
    return Opcode::setbw;
  if (opName == "mvmul")
    return Opcode::mvmul;
  if (opName == "vvadd")
    return Opcode::vvadd;
  if (opName == "vvsub")
    return Opcode::vvsub;
  if (opName == "vvmul")
    return Opcode::vvmul;
  if (opName == "vvdmul")
    return Opcode::vvdmul;
  if (opName == "vvmax")
    return Opcode::vvmax;
  if (opName == "vvsll")
    return Opcode::vvsll;
  if (opName == "vvsra")
    return Opcode::vvsra;
  if (opName == "vavg")
    return Opcode::vavg;
  if (opName == "vrelu")
    return Opcode::vrelu;
  if (opName == "vtanh")
    return Opcode::vtanh;
  if (opName == "vsigm")
    return Opcode::vsigm;
  if (opName == "vsoftmax")
    return Opcode::vsoftmax;
  if (opName == "vmv")
    return Opcode::vmv;
  if (opName == "vrsu")
    return Opcode::vrsu;
  if (opName == "vrsl")
    return Opcode::vrsl;
  if (opName == "ld")
    return Opcode::ld;
  if (opName == "st")
    return Opcode::st;
  if (opName == "lldi")
    return Opcode::lldi;
  if (opName == "lmv")
    return Opcode::lmv;
  if (opName == "send")
    return Opcode::send;
  if (opName == "recv")
    return Opcode::recv;
  if (opName == "wait")
    return Opcode::wait;
  if (opName == "sync")
    return Opcode::sync;
  llvm_unreachable("Unsupported PIM binary opcode");
}

inline llvm::StringRef opcodeToString(Opcode opcode) {
  switch (opcode) {
  case Opcode::nop:      return "nop";
  case Opcode::sldi:     return "sldi";
  case Opcode::sld:      return "sld";
  case Opcode::sadd:     return "sadd";
  case Opcode::ssub:     return "ssub";
  case Opcode::smul:     return "smul";
  case Opcode::saddi:    return "saddi";
  case Opcode::smuli:    return "smuli";
  case Opcode::setbw:    return "setbw";
  case Opcode::mvmul:    return "mvmul";
  case Opcode::vvadd:    return "vvadd";
  case Opcode::vvsub:    return "vvsub";
  case Opcode::vvmul:    return "vvmul";
  case Opcode::vvdmul:   return "vvdmul";
  case Opcode::vvmax:    return "vvmax";
  case Opcode::vvsll:    return "vvsll";
  case Opcode::vvsra:    return "vvsra";
  case Opcode::vavg:     return "vavg";
  case Opcode::vrelu:    return "vrelu";
  case Opcode::vtanh:    return "vtanh";
  case Opcode::vsigm:    return "vsigm";
  case Opcode::vsoftmax: return "vsoftmax";
  case Opcode::vmv:      return "vmv";
  case Opcode::vrsu:     return "vrsu";
  case Opcode::vrsl:     return "vrsl";
  case Opcode::ld:       return "ld";
  case Opcode::st:       return "st";
  case Opcode::lldi:     return "lldi";
  case Opcode::lmv:      return "lmv";
  case Opcode::send:     return "send";
  case Opcode::recv:     return "recv";
  case Opcode::wait:     return "wait";
  case Opcode::sync:     return "sync";
  }
  llvm_unreachable("Unsupported PIM binary opcode");
}

inline InstructionRecord makeInstructionRecord(const llvm::json::Object& instruction) {
  InstructionRecord record;
  std::optional<llvm::StringRef> opName = instruction.getString("op");
  assert(opName && "Missing op field in PIM instruction");
  record.opcode = opcodeFromString(*opName);
  record.rd = toU8(getOptionalInt(instruction, "rd"));
  record.r1 = toU8(getOptionalInt(instruction, "rs1"));

  switch (record.opcode) {
  case Opcode::sldi:
  case Opcode::saddi:
  case Opcode::smuli:
  case Opcode::lldi:  record.r2OrImm = getOptionalInt(instruction, "imm"); break;
  case Opcode::mvmul:
    record.r2OrImm = getOptionalInt(instruction, "mbiw");
    record.generic1 = getOptionalInt(instruction, "relu");
    record.generic2 = getOptionalInt(instruction, "group");
    break;
  case Opcode::setbw:
    record.generic1 = getOptionalInt(instruction, "ibiw");
    record.generic2 = getOptionalInt(instruction, "obiw");
    break;
  case Opcode::send:
  case Opcode::recv:
    record.r2OrImm = getOptionalInt(instruction, "core");
    record.generic3 = getOptionalInt(instruction, "size");
    break;
  default: record.r2OrImm = getOptionalInt(instruction, "rs2"); break;
  }

  if (record.opcode != Opcode::mvmul && record.opcode != Opcode::setbw) {
    if (auto* offsetValue = instruction.getObject("offset")) {
      record.generic1 = getOptionalInt(*offsetValue, "offset_select");
      record.generic2 = getOptionalInt(*offsetValue, "offset_value");
    }
  }

  if (instruction.get("len"))
    record.generic3 = getOptionalInt(instruction, "len");
  else if (instruction.get("size") && record.opcode != Opcode::send && record.opcode != Opcode::recv)
    record.generic3 = getOptionalInt(instruction, "size");

  return record;
}

inline llvm::json::Object makeInstructionJson(const InstructionRecord& record) {
  llvm::json::Object instruction;
  instruction["op"] = opcodeToString(record.opcode).str();

  auto addOffset = [&](int32_t offsetSelect, int32_t offsetValue) {
    llvm::json::Object offset;
    offset["offset_select"] = offsetSelect;
    offset["offset_value"] = offsetValue;
    instruction["offset"] = std::move(offset);
  };

  switch (record.opcode) {
  case Opcode::sldi:
    instruction["rd"] = static_cast<int64_t>(record.rd);
    instruction["imm"] = record.r2OrImm;
    break;
  case Opcode::sld:
    instruction["rd"] = static_cast<int64_t>(record.rd);
    instruction["rs1"] = static_cast<int64_t>(record.r1);
    addOffset(record.generic1, record.generic2);
    break;
  case Opcode::sadd:
  case Opcode::ssub:
  case Opcode::smul:
    instruction["rd"] = static_cast<int64_t>(record.rd);
    instruction["rs1"] = static_cast<int64_t>(record.r1);
    instruction["rs2"] = record.r2OrImm;
    break;
  case Opcode::saddi:
  case Opcode::smuli:
    instruction["rd"] = static_cast<int64_t>(record.rd);
    instruction["rs1"] = static_cast<int64_t>(record.r1);
    instruction["imm"] = record.r2OrImm;
    break;
  case Opcode::setbw:
    instruction["ibiw"] = record.generic1;
    instruction["obiw"] = record.generic2;
    break;
  case Opcode::mvmul:
    instruction["rd"] = static_cast<int64_t>(record.rd);
    instruction["rs1"] = static_cast<int64_t>(record.r1);
    instruction["mbiw"] = record.r2OrImm;
    instruction["relu"] = record.generic1;
    instruction["group"] = record.generic2;
    break;
  case Opcode::vvadd:
  case Opcode::vvsub:
  case Opcode::vvmul:
  case Opcode::vvdmul:
  case Opcode::vvmax:
  case Opcode::vvsll:
  case Opcode::vvsra:
  case Opcode::vavg:
  case Opcode::vmv:
  case Opcode::vrsu:
  case Opcode::vrsl:
    instruction["rd"] = static_cast<int64_t>(record.rd);
    instruction["rs1"] = static_cast<int64_t>(record.r1);
    instruction["rs2"] = record.r2OrImm;
    addOffset(record.generic1, record.generic2);
    instruction["len"] = record.generic3;
    break;
  case Opcode::vrelu:
  case Opcode::vtanh:
  case Opcode::vsigm:
  case Opcode::vsoftmax:
    instruction["rd"] = static_cast<int64_t>(record.rd);
    instruction["rs1"] = static_cast<int64_t>(record.r1);
    addOffset(record.generic1, record.generic2);
    instruction["len"] = record.generic3;
    break;
  case Opcode::ld:
  case Opcode::st:
    instruction["rd"] = static_cast<int64_t>(record.rd);
    instruction["rs1"] = static_cast<int64_t>(record.r1);
    addOffset(record.generic1, record.generic2);
    instruction["size"] = record.generic3;
    break;
  case Opcode::lldi:
    instruction["rd"] = static_cast<int64_t>(record.rd);
    instruction["imm"] = record.r2OrImm;
    addOffset(record.generic1, record.generic2);
    instruction["len"] = record.generic3;
    break;
  case Opcode::lmv:
    instruction["rd"] = static_cast<int64_t>(record.rd);
    instruction["rs1"] = static_cast<int64_t>(record.r1);
    addOffset(record.generic1, record.generic2);
    instruction["len"] = record.generic3;
    break;
  case Opcode::send:
  case Opcode::recv:
    instruction["rd"] = static_cast<int64_t>(record.rd);
    instruction["core"] = record.r2OrImm;
    addOffset(record.generic1, record.generic2);
    instruction["size"] = record.generic3;
    break;
  case Opcode::wait:
  case Opcode::sync:
  case Opcode::nop:  break;
  }

  return instruction;
}

} // namespace onnx_mlir::pim_binary