#include "direct_ipu.hpp"
#include "host_exchange.hpp"
#include "ipu_bootloader_frames.hpp"
#include "ipu_image.hpp"

#include <algorithm>
#include <array>
#include <atomic>
#include <cstdint>
#include <cstring>
#include <elf.h>
#include <fstream>
#include <functional>
#include <iostream>
#include <sstream>
#include <stdexcept>
#include <vector>

#include <signal.h>

namespace {
template <typename T>
T objectAt(const std::vector<std::uint8_t> &bytes, std::size_t offset) {
  if (offset > bytes.size() || sizeof(T) > bytes.size() - offset)
    throw std::runtime_error("truncated bootloader ELF");
  T result;
  std::memcpy(&result, bytes.data() + offset, sizeof(result));
  return result;
}

std::vector<std::uint8_t> readFile(const std::string &path) {
  std::ifstream input(path, std::ios::binary | std::ios::ate);
  if (!input)
    throw std::runtime_error("cannot open " + path);
  const auto size = input.tellg();
  if (size < 0)
    throw std::runtime_error("cannot size " + path);
  std::vector<std::uint8_t> result(static_cast<std::size_t>(size));
  input.seekg(0);
  input.read(reinterpret_cast<char *>(result.data()), result.size());
  if (!input)
    throw std::runtime_error("cannot read " + path);
  return result;
}

std::vector<std::uint8_t> loadBootloader(const std::string &path) {
  const auto file = readFile(path);
  const auto header = objectAt<Elf32_Ehdr>(file, 0);
  if (std::memcmp(header.e_ident, ELFMAG, SELFMAG) != 0 ||
      header.e_ident[EI_CLASS] != ELFCLASS32 ||
      header.e_ident[EI_DATA] != ELFDATA2LSB ||
      header.e_phentsize != sizeof(Elf32_Phdr))
    throw std::runtime_error("unsupported bootloader ELF");
  std::vector<std::uint8_t> image;
  for (unsigned i = 0; i < header.e_phnum; ++i) {
    const auto segment =
        objectAt<Elf32_Phdr>(file, header.e_phoff + i * header.e_phentsize);
    if (segment.p_type != PT_LOAD)
      continue;
    if (segment.p_vaddr != directipu::ipu21::tileMemoryBase ||
        segment.p_filesz > segment.p_memsz || segment.p_offset > file.size() ||
        segment.p_filesz > file.size() - segment.p_offset || !image.empty())
      throw std::runtime_error("unexpected bootloader load segment");
    image.assign(file.begin() + segment.p_offset,
                 file.begin() + segment.p_offset + segment.p_filesz);
  }
  if (image.empty() || image.size() > ipuboot::descriptorAreaSize)
    throw std::runtime_error("invalid bootloader image size");
  image.resize((image.size() + ipuboot::frameSize - 1) &
                   ~std::size_t(ipuboot::frameSize - 1),
               0);
  return image;
}

unsigned logicalToPhysical(unsigned logical) {
  const unsigned pair = logical / 2;
  const unsigned lane = logical & 1u;
  const unsigned block = pair / directipu::ipu21::rows;
  unsigned row = pair % directipu::ipu21::rows;
  if (block & 1u)
    row = directipu::ipu21::rows - 1 - row;
  const unsigned column = (block / 2) * 4 + (block & 1u);
  return row * directipu::ipu21::physicalTilesPerRow + column + lane * 2;
}

void installBootloader(directipu::Device &device,
                       const std::vector<std::uint8_t> &image,
                       unsigned tileCount) {
  device.writeConfig(directipu::pci::autoldCsr, 0);
  for (std::size_t offset = 0; offset < image.size(); offset += 4) {
    std::uint32_t word;
    std::memcpy(&word, image.data() + offset, sizeof(word));
    device.writeConfig(directipu::pci::autoldData, word);
    if ((offset & 127) == 124)
      (void)device.readConfig(directipu::pci::autoldCsr);
  }

  constexpr std::uint32_t autoloadZone = 32u << directipu::pci::autoldZoneShift;
  const std::uint32_t kibibytes = image.size() / 1024;
  const std::uint32_t loadPointer =
      (image.size() / 4 - 1) & directipu::pci::autoldLoadPointerMask;
  device.writeConfig(directipu::pci::autoldTarget,
                     autoloadZone |
                         (kibibytes << directipu::pci::autoldAddressShift));
  device.writeConfig(directipu::pci::autoldCsr,
                     directipu::pci::autoldDataPresent |
                         directipu::pci::autoldGo | loadPointer);
  device.waitAutoloader(std::chrono::seconds(2));

  constexpr std::uint32_t tileMemoryKiB =
      directipu::ipu21::tileMemoryBytes / 1024;
  const std::uint32_t lastPhysicalTile = tileCount - 1;
  device.writeConfig(directipu::pci::autoldTarget,
                     autoloadZone |
                         (tileMemoryKiB << directipu::pci::autoldAddressShift) |
                         kibibytes);
  device.writeConfig(
      directipu::pci::autoldCsr,
      (lastPhysicalTile << directipu::pci::autoldCurrentTileShift) |
          directipu::pci::autoldGo | loadPointer);
  device.waitAutoloader(std::chrono::seconds(2));
}

void blockDeviceInterruptSignals() {
  sigset_t signals;
  if (sigfillset(&signals) != 0 ||
      sigprocmask(SIG_BLOCK, &signals, nullptr) != 0)
    throw directipu::systemError("block device interrupt signals");
}

void initializeDevice(directipu::Device &device) {
  using Message = directipu::MailboxMessage;
  const auto transact = [&](Message request) {
    const auto response = device.transact(request);
    if ((response.word[0] & directipu::icu::responseCommandMask) !=
        (request.word[0] & directipu::icu::responseCommandMask))
      throw std::runtime_error("unexpected ICU mailbox response");
    return response;
  };

  // GraphcoreDeviceAccessICU::newmanryReset for IPU ID 0.
  constexpr Message newmanryResetIpu0{
      {directipu::icu::newmanryResetIpu0, 0, 0, 0, 0}};
  device.control(directipu::stopMonitoring);
  transact(newmanryResetIpu0);
  device.control(directipu::startMonitoring);
  device.control(directipu::accumulateErrors);
  // The kernel uses this as a reset notification; its implementation ignores
  // the ioctl argument.
  device.control(directipu::resetDevice);
  device.setSecondaryId(0);
}

void replayConfiguration(directipu::Device &device, const std::string &path) {
  const auto bytes = readFile(path);
  constexpr std::array<std::uint8_t, 8> magic = {'I', 'P', 'U', 'C',
                                                 'F', 'G', '1', 0};
  if (bytes.size() < 12 ||
      !std::equal(magic.begin(), magic.end(), bytes.begin()))
    throw std::runtime_error("bad IPU configuration image");
  const std::uint32_t count = objectAt<std::uint32_t>(bytes, 8);
  if (count > (bytes.size() - 12) / 8 || bytes.size() != 12 + count * 8ull)
    throw std::runtime_error("invalid IPU configuration image size");
  for (std::uint32_t i = 0; i < count; ++i) {
    const std::uint32_t offset = objectAt<std::uint32_t>(bytes, 12 + i * 8);
    const std::uint32_t value = objectAt<std::uint32_t>(bytes, 16 + i * 8);
    device.writeConfig(offset, value);
  }
}

class PinnedBuffer {
public:
  PinnedBuffer(directipu::Device &device, std::size_t size)
      : device_(device), size_(size) {
    data_ = mmap(nullptr, size_, PROT_READ | PROT_WRITE,
                 MAP_PRIVATE | MAP_ANONYMOUS | MAP_POPULATE, -1, 0);
    if (data_ == MAP_FAILED) {
      data_ = nullptr;
      throw directipu::systemError("allocate bootloader buffer");
    }
    try {
      device_.attachBuffer(0, data_, size_);
      attached_ = true;
    } catch (...) {
      munmap(data_, size_);
      data_ = nullptr;
      throw;
    }
  }

  ~PinnedBuffer() {
    detach();
    if (data_)
      munmap(data_, size_);
  }

  void detach() {
    if (attached_) {
      device_.detachBuffer(0);
      attached_ = false;
    }
  }

  std::uint8_t *data() { return static_cast<std::uint8_t *>(data_); }
  std::size_t size() const { return size_; }

private:
  directipu::Device &device_;
  void *data_ = nullptr;
  std::size_t size_ = 0;
  bool attached_ = false;
};

class ExchangeBufferGuard {
public:
  explicit ExchangeBufferGuard(directipu::Device &device) : device_(device) {
    savedCcsr_ = device_.readConfig(directipu::pci::ccsr);
    for (unsigned instance = 0; instance < operational_.size(); ++instance) {
      const auto address = directipu::xb::dcxcrAddress(instance);
      // The captured configuration contains the SDK's disabled loader state,
      // so reconstruct the normal C600 routing for PCI interfaces 0 and 1.
      const unsigned ipuId = instance / directipu::xb::instancesPerInterface;
      operational_[instance] = directipu::xb::routedDcxControl(ipuId);
      device_.writeConfig(address, directipu::xb::bootloaderDcxControl());
    }
    writeLoaderCcsr();
  }

  void restorePrimary() { restore(0, directipu::xb::instancesPerInterface); }

  void restoreAll() {
    restore(0, operational_.size());
    writeAndFlushCcsr(savedCcsr_);
  }

  ~ExchangeBufferGuard() { restoreAll(); }

private:
  void restore(unsigned begin, unsigned end) {
    for (unsigned instance = begin; instance < end; ++instance)
      device_.writeConfig(directipu::xb::dcxcrAddress(instance),
                          operational_[instance]);
  }

  void writeLoaderCcsr() {
    writeAndFlushCcsr(savedCcsr_ & ~(directipu::pci::noSnoop |
                                     directipu::pci::exchangeLinkEcrcEnable));
  }

  void writeAndFlushCcsr(std::uint32_t value) {
    device_.writeConfig(directipu::pci::ccsr, value);
    (void)device_.readConfig(directipu::pci::ccsr);
  }

  directipu::Device &device_;
  std::array<std::uint32_t, directipu::xb::instanceCount> operational_{};
  std::uint32_t savedCcsr_ = 0;
};

struct TileInput {
  std::uint32_t address;
  std::uint32_t value;
  bool addPhysicalTile;
};

struct ApplicationResult {
  unsigned tile;
  std::uint32_t address;
  std::uint32_t expected;
};

struct LoadOptions {
  unsigned finalMark = 0;
  std::function<void()> beforeLaunch;
  std::function<void()> afterLaunch;
  std::function<void(PinnedBuffer &, ExchangeBufferGuard &)> handoff;
};

void loadImage(directipu::Device &device, ipuimg::Image &image,
               const std::vector<std::uint8_t> &bootloader,
               const LoadOptions &options = {}) {
  constexpr auto syncTimeout = std::chrono::seconds(10);
  const unsigned tileCount = image.header().tileCount;
  if (tileCount == 0 || tileCount % ipuboot::tilesPerBatch != 0 ||
      tileCount - 1 > directipu::pci::autoldCurrentTileMask)
    throw std::runtime_error(
        "image tile count is not representable by complete loader batches");

  installBootloader(device, bootloader, tileCount);
  device.writeConfig(directipu::pci::exchangeWindowBase,
                     directipu::pci::exchangeWindowHexopt);
  PinnedBuffer transport(device, ipuboot::transportSize);
  std::fill(transport.data(), transport.data() + transport.size(), 0);
  std::atomic_thread_fence(std::memory_order_seq_cst);
  ExchangeBufferGuard exchangeBuffers(device);
  exchangeBuffers.restorePrimary();

  unsigned batches = 0;
  for (unsigned first = 0; first < tileCount; first += ipuboot::tilesPerBatch) {
    std::size_t cursor = ipuboot::descriptorAreaSize;
    for (unsigned slot = 0; slot < ipuboot::tilesPerBatch; ++slot) {
      const unsigned tile = first + slot;
      const auto framed = ipuboot::frameTile(tile, image.tile(tile));
      if (framed.size() > transport.size() - cursor)
        throw std::runtime_error("tile batch exceeds bootloader buffer");

      if (framed.size() % ipuboot::frameSize != 0)
        throw std::runtime_error("unaligned framed tile image");
      const std::uint32_t descriptor[2] = {
          static_cast<std::uint32_t>(cursor),
          static_cast<std::uint32_t>(framed.size() / ipuboot::frameSize)};
      std::memcpy(transport.data() + slot * sizeof(descriptor), descriptor,
                  sizeof(descriptor));
      std::memcpy(transport.data() + cursor, framed.data(), framed.size());
      cursor += framed.size();
    }
    std::atomic_thread_fence(std::memory_order_seq_cst);
    device.setMark(1);
    try {
      device.waitForMark(0, syncTimeout);
    } catch (const std::exception &error) {
      throw std::runtime_error("bootloader batch " + std::to_string(batches) +
                               ": " + error.what());
    }
    ++batches;
  }

  if (options.beforeLaunch) {
    exchangeBuffers.restoreAll();
    options.beforeLaunch();
  }

  std::fill(transport.data(), transport.data() + transport.size(),
            ipuboot::executeSentinel);
  std::atomic_thread_fence(std::memory_order_seq_cst);
  try {
    device.waitForMark(0, syncTimeout);
  } catch (const std::exception &error) {
    throw std::runtime_error("before bootloader sentinel: " +
                             std::string(error.what()));
  }
  device.setMark(1);
  try {
    device.waitForMark(0, syncTimeout);
  } catch (const std::exception &error) {
    throw std::runtime_error("bootloader sentinel: " +
                             std::string(error.what()));
  }
  device.setMark(options.finalMark ? options.finalMark : batches);
  if (options.afterLaunch)
    options.afterLaunch();
  try {
    device.waitForMark(0, syncTimeout);
  } catch (const std::exception &error) {
    throw std::runtime_error("application initial sync: " +
                             std::string(error.what()));
  }
  if (options.handoff)
    options.handoff(transport, exchangeBuffers);
}

std::uint32_t readRunningTileWord(directipu::Device &device, unsigned tile,
                                  std::uint32_t address) {
  constexpr unsigned supervisorContext = 0;
  if (device.tileContextQuiescent(tile, supervisorContext))
    return device.readTileMemoryWordStopped(tile, supervisorContext, address);
  const std::uint32_t contextBit = 1u << supervisorContext;
  const std::uint32_t oldRunBreak =
      device.readTileDebug(tile, directipu::tileDebugRunBreak);
  device.writeTileDebug(tile, directipu::tileDebugRunBreak,
                        oldRunBreak | contextBit);
  const auto restore = [&] {
    device.writeTileDebug(tile, directipu::tileDebugRunBreak, oldRunBreak);
    if ((oldRunBreak & contextBit) == 0)
      device.writeTileDebug(tile, directipu::tileDebugExceptionClear,
                            contextBit);
  };
  try {
    const auto deadline =
        std::chrono::steady_clock::now() + std::chrono::milliseconds(100);
    while (device.tileContextState(tile, supervisorContext) !=
           directipu::tileDebug::contextStopped) {
      if (std::chrono::steady_clock::now() >= deadline)
        throw std::runtime_error("application result breakpoint timed out");
    }
    const auto result =
        device.readTileMemoryWordStopped(tile, supervisorContext, address);
    restore();
    return result;
  } catch (...) {
    restore();
    throw;
  }
}

} // namespace

int main(int argc, char **argv) try {
  if (argc == 4 && std::string(argv[1]) == "config-write") {
    const std::uint32_t offset = std::stoul(argv[2], nullptr, 0);
    const std::uint32_t value = std::stoul(argv[3], nullptr, 0);
    directipu::Device device;
    device.writeConfig(offset, value);
    std::cout << "offset=0x" << std::hex << offset << " value=0x" << value
              << " observed=0x" << device.readConfig(offset) << std::dec
              << '\n';
  } else if (argc == 5 && std::string(argv[1]) == "tdi-read") {
    const unsigned tile = std::stoul(argv[2], nullptr, 0);
    const std::uint32_t address = std::stoul(argv[3], nullptr, 0);
    const std::uint32_t expected = std::stoul(argv[4], nullptr, 0);
    directipu::Device device;
    const auto observed = readRunningTileWord(device, tile, address);
    if (observed != expected)
      throw std::runtime_error("TDI memory comparison failed: expected " +
                               std::to_string(expected) + ", observed " +
                               std::to_string(observed));
    std::cout << "tile=" << tile << " address=0x" << std::hex << address
              << " observed=0x" << observed << " expected=0x" << expected
              << std::dec << " compare=PASS\n";
  } else if (argc == 6 && std::string(argv[1]) == "tdi-check-pattern") {
    const unsigned tile = std::stoul(argv[2], nullptr, 0);
    const std::uint32_t address = std::stoul(argv[3], nullptr, 0);
    const std::uint32_t pair = std::stoul(argv[4], nullptr, 0);
    const std::uint32_t count = std::stoul(argv[5], nullptr, 0);
    if (count == 0 || address + std::uint64_t(count) * 4 >
                          directipu::ipu21::tileMemoryBase +
                              directipu::ipu21::tileMemoryBytes)
      throw std::out_of_range("TDI pattern range");
    directipu::Device device;
    for (std::uint32_t index = 0; index < count; ++index) {
      const auto expected = 0x13579bdfu ^ pair * 0x7f4a7c15u ^
                            index * 0x9e3779b9u;
      const auto observed =
          readRunningTileWord(device, tile, address + index * 4);
      if (observed != expected)
        throw std::runtime_error("TDI pattern mismatch at word " +
                                 std::to_string(index) + ": expected " +
                                 std::to_string(expected) + ", observed " +
                                 std::to_string(observed));
    }
    std::cout << "tile=" << tile << " address=0x" << std::hex << address
              << std::dec << " words=" << count << " pair=" << pair
              << " pattern=PASS\n";
  } else if (argc == 5 && std::string(argv[1]) == "tdi-write") {
    const unsigned tile = std::stoul(argv[2], nullptr, 0);
    const unsigned reg = std::stoul(argv[3], nullptr, 0);
    const std::uint32_t value = std::stoul(argv[4], nullptr, 0);
    directipu::Device device;
    device.writeTileDebug(tile, reg, value);
    std::cout << "tile=" << tile << " reg=" << reg << " value=0x"
              << std::hex << value << " observed=0x"
              << device.readTileDebug(tile, reg) << std::dec << '\n';
  } else if (argc == 4 && std::string(argv[1]) == "tdi-nop") {
    const unsigned tile = std::stoul(argv[2], nullptr, 0);
    const unsigned context = std::stoul(argv[3], nullptr, 0);
    // This is the SDK debugger's injectable NOP: `zero $m15`.
    constexpr std::uint32_t nop = 0x44f0f3f1;
    directipu::Device device;
    const auto oldRunBreak =
        device.readTileDebug(tile, directipu::tileDebugRunBreak);
    const auto status = device.executeTileDebugInstruction(tile, context, nop);
    const auto newRunBreak =
        device.readTileDebug(tile, directipu::tileDebugRunBreak);
    if (newRunBreak != oldRunBreak)
      throw std::runtime_error("TDI run-break state was not restored");
    std::cout << "tile=" << tile << " context=" << context << " instruction=0x"
              << std::hex << nop << " injectionStatus=0x" << status
              << " runBreak=0x" << newRunBreak << std::dec << " nop=PASS\n";
  } else if (argc == 4 && std::string(argv[1]) == "tdi-pc") {
    const unsigned tile = std::stoul(argv[2], nullptr, 0);
    const unsigned context = std::stoul(argv[3], nullptr, 0);
    directipu::Device device;
    std::cout << "tile=" << tile << " context=" << context << " pc=0x"
              << std::hex
              << device.readTileProgramCounterStopped(tile, context)
              << std::dec << '\n';
  } else if (argc == 5 && std::string(argv[1]) == "tdi-mreg") {
    const unsigned tile = std::stoul(argv[2], nullptr, 0);
    const unsigned context = std::stoul(argv[3], nullptr, 0);
    const unsigned reg = std::stoul(argv[4], nullptr, 0);
    directipu::Device device;
    std::cout << "tile=" << tile << " context=" << context << " m" << reg
              << "=0x" << std::hex
              << device.readTileMRegisterStopped(tile, context, reg)
              << std::dec << '\n';
  } else if (argc == 5 && std::string(argv[1]) == "tdi-state") {
    const unsigned tile = std::stoul(argv[2], nullptr, 0);
    const unsigned context = std::stoul(argv[3], nullptr, 0);
    const unsigned index = std::stoul(argv[4], nullptr, 0);
    directipu::Device device;
    std::cout << "tile=" << tile << " context=" << context << " state[0x"
              << std::hex << index << "]=0x"
              << device.readTileControlStateStopped(tile, context, index)
              << std::dec << '\n';
  } else if ((argc == 3 || argc == 4) && std::string(argv[1]) == "tdi-dump") {
    const unsigned tile = std::stoul(argv[2], nullptr, 0);
    const unsigned count = argc == 4 ? std::stoul(argv[3], nullptr, 0)
                                     : directipu::tileDebugRegisterCount;
    if (count == 0 || count > directipu::tileDebugRegisterCount)
      throw std::out_of_range("TDI register count");
    directipu::Device device;
    for (unsigned reg = 0; reg < count; ++reg)
      std::cout << "tile=" << tile << " reg=" << reg << " value=0x" << std::hex
                << device.readTileDebug(tile, reg) << std::dec << '\n';
  } else if (argc == 3 && std::string(argv[1]) == "tdi-exception") {
    const unsigned tile = std::stoul(argv[2], nullptr, 0);
    directipu::Device device;
    const auto contextStatus =
        device.readTileDebug(tile, directipu::tileDebugContextStatus);
    const auto exceptionState =
        device.readTileDebug(tile, directipu::tileDebugExceptionState);
    const unsigned context =
        (exceptionState >> directipu::tileDebug::exceptionContextShift) &
        directipu::tileDebug::exceptionContextMask;
    const auto state = device.tileContextState(tile, context);
    std::cout << "tile=" << tile << " contextStatus=0x" << std::hex
              << contextStatus << " exceptionState=0x" << exceptionState
              << " exceptionContext=" << std::dec << context
              << " contextState=" << state
              << " pcMirror="
              << ((exceptionState >>
                   directipu::tileDebug::exceptionPcMirrorShift) &
                  directipu::tileDebug::exceptionPcMirrorMask)
              << " breakOnSync="
              << ((exceptionState >>
                   directipu::tileDebug::exceptionBreakOnSyncShift) &
                  directipu::tileDebug::exceptionBreakOnSyncMask)
              << " exchangeReturnError="
              << ((contextStatus >>
                   directipu::tileDebug::exchangeReturnErrorShift) &
                  directipu::tileDebug::exchangeReturnErrorMask)
              << " exchangeError="
              << ((contextStatus >> directipu::tileDebug::exchangeErrorShift) &
                  1u)
              << " memoryError="
              << ((contextStatus >> directipu::tileDebug::memoryErrorShift) &
                  1u);
    if (state == directipu::tileDebug::contextExceptedDebug ||
        state == directipu::tileDebug::contextExceptedNonDebug) {
      const auto type = device.readTileExceptionTypeStopped(tile, context);
      std::cout << " pc=0x" << std::hex
                << device.readTileProgramCounterStopped(tile, context)
                << " exceptionType=0x" << type << " exceptionName="
                << directipu::tileException::name(type);
    }
    std::cout << std::dec << '\n';
  } else if (argc == 3 && std::string(argv[1]) == "inspect") {
    ipuimg::Image image(argv[2]);
    directipu::Device device;
    std::cout << "tiles=" << image.header().tileCount << " base=0x" << std::hex
              << image.header().baseAddress << " entry=0x"
              << image.header().entryPoint << " autoloadControl=0x"
              << device.readConfig(directipu::pci::autoldCsr) << std::dec
              << " attach=PASS\n";
  } else if (argc == 4 && std::string(argv[1]) == "bootloader") {
    const auto image = loadBootloader(argv[2]);
    const unsigned tileCount = std::stoul(argv[3], nullptr, 0);
    if (tileCount == 0 || tileCount - 1 > directipu::pci::autoldCurrentTileMask)
      throw std::out_of_range("bootloader tile count");
    blockDeviceInterruptSignals();
    directipu::Device device;
    installBootloader(device, image, tileCount);
    std::cout << "bootloaderBytes=" << image.size()
              << " physicalTiles=" << tileCount << " autoloader=PASS\n";
  } else if (argc >= 9 && (argc - 6) % 3 == 0 &&
             std::string(argv[1]) == "host-exchange") {
    ipuimg::Image image(argv[2]);
    const auto bootloader = loadBootloader(argv[3]);
    auto protocol = hostexchange::Protocol::read(argv[5]);
    blockDeviceInterruptSignals();
    directipu::Device device;
    initializeDevice(device);
    replayConfiguration(device, argv[4]);
    hostexchange::Session session(device, std::move(protocol));
    LoadOptions options;
    options.finalMark = session.startupMark();
    options.handoff =
        [&](PinnedBuffer &transport, ExchangeBufferGuard &exchangeBuffers) {
          session.start([&] {
            exchangeBuffers.restoreAll();
            device.writeConfig(directipu::pci::exchangeWindowBase,
                               directipu::pci::exchangeWindowHexopt);
            transport.detach();
          });
        };
    try {
      loadImage(device, image, bootloader, options);
    } catch (const std::exception &error) {
      throw std::runtime_error(
          "host-exchange image startup failed (GS1=" +
          std::to_string(device.readConfig(directipu::pci::hspGs1Control) &
                         directipu::pci::hspMarkMask) +
          ", GS2=" +
          std::to_string(device.readConfig(directipu::pci::hspGs2Control) &
                         directipu::pci::hspMarkMask) +
          ", tile0Context=0x" + [&] {
            std::ostringstream value;
            value << std::hex
                  << device.readTileDebug(
                         0, directipu::tileDebugContextStatus);
            return value.str();
          }() +
          "): " + error.what());
    }
    unsigned calls = 0;
    for (int argument = 6; argument < argc; argument += 3) {
      const std::string name = argv[argument];
      const std::string inputPath = argv[argument + 1];
      const std::string outputPath = argv[argument + 2];
      const std::vector<std::uint8_t> input =
          inputPath == "-" ? std::vector<std::uint8_t>{}
                           : readFile(inputPath);
      const auto output = session.invoke(name, input);
      if (outputPath != "-") {
        std::ofstream file(outputPath, std::ios::binary | std::ios::trunc);
        file.write(reinterpret_cast<const char *>(output.data()), output.size());
        if (!file)
          throw std::runtime_error("cannot write " + outputPath);
      } else if (!output.empty()) {
        throw std::runtime_error(name + " produced " +
                                 std::to_string(output.size()) +
                                 " bytes but no output file was supplied");
      }
      ++calls;
    }
    std::cout << "hostExchangeCalls=" << calls
              << " directHostExchange=PASS\n";
  } else if ((argc == 6 || argc == 9) &&
             (std::string(argv[1]) == "load-bare" ||
              std::string(argv[1]) == "load-bare-startup-trace" ||
              std::string(argv[1]) == "load-bare-result")) {
    ipuimg::Image image(argv[2]);
    const auto bootloader = loadBootloader(argv[3]);
    blockDeviceInterruptSignals();
    directipu::Device device;
    initializeDevice(device);
    replayConfiguration(device, argv[4]);
    const unsigned finalMark = std::stoul(argv[5], nullptr, 0);
    if (finalMark == 0 || finalMark > directipu::pci::hspMarkMask)
      throw std::out_of_range("bare final mark");
    const bool traceStartup = std::string(argv[1]) == "load-bare-startup-trace";
    const bool readResult = std::string(argv[1]) == "load-bare-result";
    if (readResult != (argc == 9))
      throw std::invalid_argument("load-bare-result arguments");
    ApplicationResult applicationResult{};
    const ApplicationResult *result = nullptr;
    if (readResult) {
      applicationResult.tile = std::stoul(argv[6], nullptr, 0);
      applicationResult.address = std::stoul(argv[7], nullptr, 0);
      applicationResult.expected = std::stoul(argv[8], nullptr, 0);
      if (applicationResult.tile >= image.header().tileCount)
        throw std::out_of_range("result tile");
      result = &applicationResult;
    }
    LoadOptions options;
    options.finalMark = finalMark;
    loadImage(device, image, bootloader, options);
    device.writeConfig(directipu::pci::exchangeWindowBase,
                       directipu::pci::exchangeWindowHexopt);
    device.setMark(1);
    std::this_thread::sleep_for(readResult ? std::chrono::milliseconds(250)
                                           : std::chrono::milliseconds(10));
    if (readResult) {
      const std::uint32_t observed =
          readRunningTileWord(device, result->tile, result->address);
      if (observed != result->expected)
        throw std::runtime_error("application result mismatch: expected " +
                                 std::to_string(result->expected) +
                                 ", observed " + std::to_string(observed));
      std::cout << "tile=" << result->tile << " resultAddress=0x" << std::hex
                << result->address << " resultValue=0x" << observed << std::dec
                << " applicationResult=PASS ";
    }
    if (traceStartup) {
      const std::uint32_t startupRunBreak =
          device.readTileDebug(0, directipu::tileDebugRunBreak);
      device.writeTileDebug(0, directipu::tileDebugRunBreak,
                            startupRunBreak | 1u);
      const auto deadline =
          std::chrono::steady_clock::now() + std::chrono::milliseconds(100);
      while (device.tileContextState(0, 0) !=
             directipu::tileDebug::contextStopped) {
        if (std::chrono::steady_clock::now() >= deadline)
          throw std::runtime_error("application startup breakpoint timed out");
      }
      const auto tileBytes = image.tile(0);
      const std::size_t entryOffset =
          image.header().entryPoint - image.header().baseAddress;
      std::uint32_t entryWord;
      std::memcpy(&entryWord, tileBytes.data() + entryOffset + 8,
                  sizeof(entryWord));
      const std::uint32_t target = entryWord & directipu::ipu21::addressMask;
      std::cerr << "startup pc=0x" << std::hex
                << device.readTileProgramCounterStopped(0, 0)
                << " entryPatch=0x"
                << device.readTileMemoryWordStopped(
                       0, 0, image.header().entryPoint + 8)
                << " targetWord=0x"
                << device.readTileMemoryWordStopped(0, 0, target) << std::dec
                << '\n';
      device.writeTileDebug(0, directipu::tileDebugRunBreak, startupRunBreak);
      device.writeTileDebug(0, directipu::tileDebugExceptionClear, 1u);
    }
    const std::uint32_t startupContext =
        device.readTileDebug(0, directipu::tileDebugContextStatus);
    unsigned activeSupervisors = 0;
    for (unsigned tile = 0; tile < image.header().tileCount; ++tile)
      activeSupervisors += device.tileContextState(tile, 0) ==
                           directipu::tileDebug::contextRunning;
    std::cout << "tiles=" << image.header().tileCount << " batches="
              << image.header().tileCount / ipuboot::tilesPerBatch
              << " finalMark=" << finalMark
              << " sentinelHandshake=PASS executionHandshake=PASS"
              << " startupRelease=PASS"
              << " tile0ContextStatus=0x" << std::hex << startupContext
              << std::dec << " activeSupervisorTiles=" << activeSupervisors
              << " execution="
              << (activeSupervisors == image.header().tileCount
                      ? "ACTIVE_ALL_TILES"
                  : activeSupervisors ? "ACTIVE_PARTIAL"
                                      : "UNVERIFIED")
              << '\n';
  } else if ((argc == 5 || argc == 8) && std::string(argv[1]) == "load" &&
             (argc == 5 || std::string(argv[5]) == "--tile-id-input" ||
              std::string(argv[5]) == "--constant-input")) {
    ipuimg::Image image(argv[2]);
    const auto bootloader = loadBootloader(argv[3]);
    blockDeviceInterruptSignals();
    directipu::Device device;
    initializeDevice(device);
    replayConfiguration(device, argv[4]);
    TileInput tileInput{};
    if (argc == 8) {
      tileInput.address = std::stoul(argv[6], nullptr, 0);
      tileInput.value = std::stoul(argv[7], nullptr, 0);
      tileInput.addPhysicalTile = std::string(argv[5]) == "--tile-id-input";
    }
    LoadOptions options;
    if (argc == 8) {
      const auto writeInputs = [&](bool poison) {
        directipu::ExchangeBar exchange;
        for (unsigned tile = 0; tile < image.header().tileCount; ++tile) {
          std::uint32_t value =
              tileInput.value +
              (tileInput.addPhysicalTile ? logicalToPhysical(tile) : 0);
          if (poison)
            value = ~value;
          exchange.writeTileMemory(device, tile, tileInput.address, &value,
                                   sizeof(value));
        }
      };
      options.finalMark = 1;
      options.beforeLaunch = [&] {
        device.control(directipu::setHexoptIdentityTable);
        writeInputs(true);
      };
      options.afterLaunch = [&] {
        std::this_thread::sleep_for(std::chrono::milliseconds(100));
        writeInputs(false);
      };
    }
    loadImage(device, image, bootloader, options);
    std::cout << "tiles=" << image.header().tileCount << " batches="
              << image.header().tileCount / ipuboot::tilesPerBatch
              << " directLoad=PASS\n";
  } else {
    std::cerr << "usage:\n"
              << "  direct_ipu_loader config-write OFFSET VALUE\n"
              << "  direct_ipu_loader tdi-read TILE ADDRESS EXPECTED\n"
              << "  direct_ipu_loader tdi-check-pattern TILE ADDRESS PAIR"
                 " COUNT\n"
              << "  direct_ipu_loader tdi-write TILE REGISTER VALUE\n"
              << "  direct_ipu_loader tdi-nop TILE CONTEXT\n"
              << "  direct_ipu_loader tdi-pc TILE CONTEXT\n"
              << "  direct_ipu_loader tdi-mreg TILE CONTEXT REGISTER\n"
              << "  direct_ipu_loader tdi-state TILE CONTEXT INDEX\n"
              << "  direct_ipu_loader tdi-dump TILE [COUNT]\n"
              << "  direct_ipu_loader tdi-exception TILE\n"
              << "  direct_ipu_loader inspect IMAGE\n"
              << "  direct_ipu_loader bootloader TILE_BOOTLOADER_ELF"
                 " TILE_COUNT\n"
              << "  direct_ipu_loader host-exchange IMAGE TILE_BOOTLOADER_ELF"
                 " CONFIG MANIFEST CALL INPUT|- OUTPUT|- [CALL ...]\n"
              << "  direct_ipu_loader load-bare IMAGE TILE_BOOTLOADER_ELF"
                 " CONFIG FINAL_MARK\n"
              << "  direct_ipu_loader load-bare-startup-trace IMAGE"
                 " TILE_BOOTLOADER_ELF CONFIG FINAL_MARK\n"
              << "  direct_ipu_loader load-bare-result IMAGE"
                 " TILE_BOOTLOADER_ELF CONFIG FINAL_MARK TILE ADDRESS"
                 " EXPECTED\n"
              << "  direct_ipu_loader load IMAGE TILE_BOOTLOADER_ELF CONFIG"
                 " [--tile-id-input ADDRESS BIAS |"
                 " --constant-input ADDRESS VALUE]\n";
    return 2;
  }
  return 0;
} catch (const std::exception &error) {
  std::cerr << "direct_ipu_loader: " << error.what() << '\n';
  return 1;
}