#include <cstdint>
#include <cstring>
#include <dlfcn.h>
#include <iomanip>
#include <iostream>
#include <stdexcept>
#include <string>

namespace {
struct ConstantLocation {
  const char *name;
  std::size_t offset;
  bool hasValidityByte;
};

std::uint32_t readConstant(const std::uint8_t *arch,
                           const ConstantLocation &location) {
  std::uint32_t value;
  std::memcpy(&value, arch + location.offset, sizeof(value));
  if (location.hasValidityByte &&
      arch[location.offset + sizeof(value)] == 0)
    throw std::runtime_error(std::string("invalid architecture constant ") +
                             location.name);
  return value;
}
} // namespace

int main(int argc, char **argv) try {
  if (argc != 2) {
    std::cerr << "usage: ipu21_arch_constants LIBIPU_ARCH_INFO\n";
    return 2;
  }
  void *library = dlopen(argv[1], RTLD_NOW | RTLD_LOCAL);
  if (!library) throw std::runtime_error(dlerror());
  using Lookup = void *(*)(const std::string &);
  auto *rawLookup = dlsym(library, "_Z22ipuArchInfoByName_swigRKNSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEE");
  if (!rawLookup) throw std::runtime_error("ipuArchInfoByName_swig not found");
  const auto lookup = reinterpret_cast<Lookup>(rawLookup);
  const auto *arch = static_cast<const std::uint8_t *>(lookup("ipu21"));
  if (!arch) throw std::runtime_error("ipu21 architecture lookup failed");

  // Offsets are recovered from IPUDebugLLD::getConfigAddressFromTileTdiRegister
  // in Graphcore target-access 3.4. Values remain supplied by the SDK's IPU21
  // architecture module rather than copied into this tool.
  constexpr ConstantLocation constants[] = {
      {"tdiBase", 0x7a18, true},
      {"tdiRegisterStride", 0x0ee0, true},
      {"tdiTileStride", 0x6908, true},
      {"tdiRunBreakIndex", 0xf868, false},
      {"tdiRunBreakThreadShift", 0xf8b0, false},
      {"tdiRunBreakThreadMask", 0xf8b8, false},
      {"tdiRunBreakEnableShift", 0xf8f0, false},
      {"tdiRunBreakEnableMask", 0xf8f8, false},
      {"tdiRunBreakVertexIndex", 0xf970, false},
      {"tdiInstructionIndex", 0xf9f8, false},
      {"tdiInstructionOwnerIndex", 0xfa80, false},
      {"tdiDebugDataIndex", 0xfc98, false},
      {"tdiStatusIndex", 0xfd20, false},
      {"tdiStatusContextNotQuiescentShift", 0xfd68, false},
      {"tdiStatusContextNotQuiescentMask", 0xfd70, false},
      {"tdiStatusInvalidShift", 0xfda8, false},
      {"tdiStatusInvalidMask", 0xfdb0, false},
      {"tdiStatusDoubleShift", 0xfde8, false},
      {"tdiStatusDoubleMask", 0xfdf0, false},
      {"tdiStatusBusyShift", 0xfe28, false},
      {"tdiStatusBusyMask", 0xfe30, false},
      {"tdiStatusClearIndex", 0xfe68, false},
      // Recovered from IPUDebugLLD::enableIBreak.
      {"supervisorDebugExceptionControlIndex", 0xd590, false},
      {"supervisorInstructionBreakControlIndex", 0xd7a8, false},
      {"supervisorInstructionBreakPCIndex", 0xd8b0, false},
      {"supervisorDebugExceptionEpcmShift", 0xd5d8, false},
      {"supervisorDebugExceptionEpcmMask", 0xd5e0, false},
  };
  for (const auto &constant : constants)
    std::cout << constant.name << "=0x" << std::hex
              << readConstant(arch, constant) << std::dec << '\n';

  auto *encoder = *reinterpret_cast<void *const *>(arch + 0x48);
  auto *vtable = *reinterpret_cast<void *const *const *>(encoder);
  using Encode2 = std::uint32_t (*)(void *, unsigned, unsigned);
  using Encode4 =
      std::uint32_t (*)(void *, unsigned, unsigned, unsigned, unsigned);
  const auto encode2 = [&](std::size_t byteOffset, unsigned a, unsigned b) {
    return reinterpret_cast<Encode2>(vtable[byteOffset / sizeof(void *)])(
        encoder, a, b);
  };
  const auto encode4 = [&](std::size_t byteOffset, unsigned a, unsigned b,
                           unsigned c, unsigned d) {
    return reinterpret_cast<Encode4>(vtable[byteOffset / sizeof(void *)])(
        encoder, a, b, c, d);
  };
  constexpr unsigned debugDataCsr = 0x70;
  std::cout << "writeDebugDataToM1Instruction=0x" << std::hex
            << encode2(0x718, 1, debugDataCsr) << '\n'
            << "loadM0FromM1Instruction=0x"
            << encode4(0x770, 0, 15, 1, 0) << '\n'
            << "readM0ToDebugDataInstruction=0x"
            << encode2(0x960, debugDataCsr, 0) << std::dec << '\n';
  dlclose(library);
  return 0;
} catch (const std::exception &error) {
  std::cerr << "ipu21_arch_constants: " << error.what() << '\n';
  return 1;
}