#pragma once

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

#include <array>

#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"

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,
};

inline constexpr size_t kOpcodeCount = static_cast<size_t>(Opcode::sync) + 1;
inline constexpr std::array<llvm::StringLiteral, kOpcodeCount> kOpcodeNames = {
  "nop",   "sldi", "sld",   "sadd",  "ssub", "smul", "saddi", "smuli", "setbw", "mvmul", "vvadd",
  "vvsub", "vvmul", "vvdmul", "vvmax", "vvsll", "vvsra", "vavg", "vrelu", "vtanh", "vsigm", "vsoftmax",
  "vmv",   "vrsu", "vrsl",  "ld",    "st",   "lldi", "lmv",   "send",  "recv",  "wait",  "sync",
};
static_assert(kOpcodeNames.size() == kOpcodeCount);

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;
};

using EncodedInstruction = std::array<char, kRecordSize>;
static_assert(kRecordSize == 20);
static_assert(sizeof(EncodedInstruction) == kRecordSize);

inline EncodedInstruction encodeInstructionRecord(const InstructionRecord& record) {
  EncodedInstruction encoded = {};
  encoded[0] = static_cast<char>(static_cast<uint8_t>(record.opcode));
  encoded[1] = static_cast<char>(record.rd);
  encoded[2] = static_cast<char>(record.r1);
  encoded[3] = static_cast<char>(record.flags);
  llvm::support::endian::write32le(encoded.data() + 4, static_cast<uint32_t>(record.r2OrImm));
  llvm::support::endian::write32le(encoded.data() + 8, static_cast<uint32_t>(record.generic1));
  llvm::support::endian::write32le(encoded.data() + 12, static_cast<uint32_t>(record.generic2));
  llvm::support::endian::write32le(encoded.data() + 16, static_cast<uint32_t>(record.generic3));
  return encoded;
}

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 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 int32_t toI32(int64_t value) { return onnx_mlir::pim::checkedI32OrCrash(value, "binary field"); }

inline uint8_t toU8(int64_t value) {
  return onnx_mlir::pim::checkedU8OrCrash(static_cast<uint64_t>(value), "binary field");
}

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;
}

struct InstructionJsonFormat {
  bool rd;
  bool r1;
  bool offset;
  llvm::StringLiteral r2;
  llvm::StringLiteral generic1;
  llvm::StringLiteral generic2;
  llvm::StringLiteral generic3;
};

inline constexpr std::array<InstructionJsonFormat, kOpcodeCount> kInstructionJsonFormats = {{
  {false, false, false, "",     "",     "",      ""    }, // nop
  {true,  false, false, "imm",  "",     "",      ""    }, // sldi
  {true,  true,  true,  "",     "",     "",      ""    }, // sld
  {true,  true,  false, "rs2",  "",     "",      ""    }, // sadd
  {true,  true,  false, "rs2",  "",     "",      ""    }, // ssub
  {true,  true,  false, "rs2",  "",     "",      ""    }, // smul
  {true,  true,  false, "imm",  "",     "",      ""    }, // saddi
  {true,  true,  false, "imm",  "",     "",      ""    }, // smuli
  {false, false, false, "",     "ibiw", "obiw",  ""    }, // setbw
  {true,  true,  false, "mbiw", "relu", "group", ""    }, // mvmul
  {true,  true,  true,  "rs2",  "",     "",      "len" }, // vvadd
  {true,  true,  true,  "rs2",  "",     "",      "len" }, // vvsub
  {true,  true,  true,  "rs2",  "",     "",      "len" }, // vvmul
  {true,  true,  true,  "rs2",  "",     "",      "len" }, // vvdmul
  {true,  true,  true,  "rs2",  "",     "",      "len" }, // vvmax
  {true,  true,  true,  "rs2",  "",     "",      "len" }, // vvsll
  {true,  true,  true,  "rs2",  "",     "",      "len" }, // vvsra
  {true,  true,  true,  "rs2",  "",     "",      "len" }, // vavg
  {true,  true,  true,  "",     "",     "",      "len" }, // vrelu
  {true,  true,  true,  "",     "",     "",      "len" }, // vtanh
  {true,  true,  true,  "",     "",     "",      "len" }, // vsigm
  {true,  true,  true,  "",     "",     "",      "len" }, // vsoftmax
  {true,  true,  true,  "rs2",  "",     "",      "len" }, // vmv
  {true,  true,  true,  "rs2",  "",     "",      "len" }, // vrsu
  {true,  true,  true,  "rs2",  "",     "",      "len" }, // vrsl
  {true,  true,  true,  "",     "",     "",      "size"}, // ld
  {true,  true,  true,  "",     "",     "",      "size"}, // st
  {true,  false, true,  "imm",  "",     "",      "len" }, // lldi
  {true,  true,  true,  "",     "",     "",      "len" }, // lmv
  {true,  false, true,  "core", "",     "",      "size"}, // send
  {true,  false, true,  "core", "",     "",      "size"}, // recv
  {false, false, false, "",     "event_register", "wait_value", ""}, // wait
  {false, false, false, "core", "event_register", "", ""}, // sync
}};
static_assert(kInstructionJsonFormats.size() == kOpcodeCount);

inline Opcode opcodeFromString(llvm::StringRef opName) {
  for (auto [index, name] : llvm::enumerate(kOpcodeNames))
    if (opName == name)
      return static_cast<Opcode>(index);
  llvm_unreachable("Unsupported PIM binary opcode");
}

inline llvm::StringRef opcodeToString(Opcode opcode) {
  size_t index = static_cast<size_t>(opcode);
  assert(index < kOpcodeNames.size() && "Unsupported PIM binary opcode");
  return kOpcodeNames[index];
}

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);
  const auto& format = kInstructionJsonFormats[static_cast<size_t>(record.opcode)];
  if (format.rd)
    record.rd = toU8(getOptionalInt(instruction, "rd"));
  if (format.r1)
    record.r1 = toU8(getOptionalInt(instruction, "rs1"));
  if (!format.r2.empty())
    record.r2OrImm = getOptionalInt(instruction, format.r2);
  if (!format.generic1.empty())
    record.generic1 = getOptionalInt(instruction, format.generic1);
  if (!format.generic2.empty())
    record.generic2 = getOptionalInt(instruction, format.generic2);
  if (format.offset) {
    if (auto* offsetValue = instruction.getObject("offset")) {
      record.generic1 = getOptionalInt(*offsetValue, "offset_select");
      record.generic2 = getOptionalInt(*offsetValue, "offset_value");
    }
  }
  if (!format.generic3.empty())
    record.generic3 = getOptionalInt(instruction, format.generic3);
  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);
  };
  const auto& format = kInstructionJsonFormats[static_cast<size_t>(record.opcode)];
  if (format.rd)
    instruction["rd"] = static_cast<int64_t>(record.rd);
  if (format.r1)
    instruction["rs1"] = static_cast<int64_t>(record.r1);
  if (!format.r2.empty())
    instruction[format.r2] = record.r2OrImm;
  if (!format.generic1.empty())
    instruction[format.generic1] = record.generic1;
  if (!format.generic2.empty())
    instruction[format.generic2] = record.generic2;
  if (format.offset)
    addOffset(record.generic1, record.generic2);
  if (!format.generic3.empty())
    instruction[format.generic3] = record.generic3;
  return instruction;
}

} // namespace onnx_mlir::pim_binary