#include "ipu_image.hpp"
#include "baremetal_protocol.hpp"
#include "exchange_plan.hpp"
#include "ipu_image_writer.hpp"

#include <algorithm>
#include <array>
#include <cstdint>
#include <elf.h>
#include <filesystem>
#include <fstream>
#include <iostream>
#include <limits>
#include <stdexcept>
#include <string>
#include <vector>

namespace fs = std::filesystem;

namespace {
unsigned logicalToPhysical(unsigned logical) {
  const unsigned pair = logical / 2;
  const unsigned lane = logical & 1u;
  const unsigned block = pair / 23;
  unsigned row = pair % 23;
  if (block & 1u)
    row = 22 - row;
  const unsigned column = (block / 2) * 4 + (block & 1u);
  return row * 64 + column + lane * 2;
}

struct LoadImage {
  std::uint32_t base = 0;
  std::uint32_t entry = 0;
  std::vector<std::uint8_t> bytes;
};

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

std::vector<std::uint8_t> readFile(const fs::path &path) {
  std::ifstream input(path, std::ios::binary | std::ios::ate);
  if (!input)
    throw std::runtime_error("cannot open " + path.string());
  const auto size = input.tellg();
  if (size < 0)
    throw std::runtime_error("cannot size " + path.string());
  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.string());
  return result;
}

LoadImage loadElf(const fs::path &path) {
  const auto file = readFile(path);
  const Elf32_Ehdr eh = readObject<Elf32_Ehdr>(file, 0);
  if (std::memcmp(eh.e_ident, ELFMAG, SELFMAG) != 0 ||
      eh.e_ident[EI_CLASS] != ELFCLASS32 ||
      eh.e_ident[EI_DATA] != ELFDATA2LSB || eh.e_type != ET_EXEC ||
      eh.e_phentsize != sizeof(Elf32_Phdr))
    throw std::runtime_error(path.string() + ": unsupported ELF");

  std::vector<Elf32_Phdr> loads;
  std::uint64_t low = std::numeric_limits<std::uint32_t>::max();
  std::uint64_t high = 0;
  for (unsigned i = 0; i < eh.e_phnum; ++i) {
    const auto ph = readObject<Elf32_Phdr>(
        file, eh.e_phoff + static_cast<std::size_t>(i) * eh.e_phentsize);
    if (ph.p_type != PT_LOAD)
      continue;
    if (ph.p_filesz > ph.p_memsz || ph.p_offset > file.size() ||
        ph.p_filesz > file.size() - ph.p_offset)
      throw std::runtime_error(path.string() + ": invalid PT_LOAD");
    low = std::min<std::uint64_t>(low, ph.p_vaddr);
    high =
        std::max<std::uint64_t>(high, std::uint64_t(ph.p_vaddr) + ph.p_memsz);
    loads.push_back(ph);
  }
  if (loads.empty() || high > std::numeric_limits<std::uint32_t>::max() ||
      high - low > std::numeric_limits<std::uint32_t>::max())
    throw std::runtime_error(path.string() + ": invalid load address range");

  LoadImage result;
  result.base = static_cast<std::uint32_t>(low);
  result.entry = eh.e_entry;
  result.bytes.resize(static_cast<std::size_t>(high - low), 0);
  for (const auto &ph : loads)
    std::copy_n(file.begin() + ph.p_offset, ph.p_filesz,
                result.bytes.begin() + (ph.p_vaddr - result.base));
  return result;
}

void pack(const fs::path &directory, const fs::path &outputPath,
          unsigned tileCount, unsigned templateTile) {
  if (tileCount == 0 || templateTile >= tileCount)
    throw std::runtime_error("invalid tile count or template tile");
  std::vector<LoadImage> images;
  images.reserve(tileCount);
  for (unsigned tile = 0; tile < tileCount; ++tile)
    images.push_back(
        loadElf(directory / ("t_" + std::to_string(tile) + ".elf")));
  const LoadImage &base = images[templateTile];
  for (unsigned tile = 0; tile < tileCount; ++tile) {
    if (images[tile].base != base.base || images[tile].entry != base.entry ||
        images[tile].bytes.size() != base.bytes.size())
      throw std::runtime_error("tile ELF load layouts do not match");
  }

  const std::uint32_t baseAddress = base.base;
  const std::uint32_t entryPoint = base.entry;
  const std::size_t imageSize = base.bytes.size();
  std::vector<std::vector<std::uint8_t>> bytes;
  bytes.reserve(images.size());
  for (auto &image : images)
    bytes.push_back(std::move(image.bytes));
  const auto stats = ipuimg::write(outputPath.string(), bytes, baseAddress,
                                   entryPoint, templateTile);
  std::cout << "tiles=" << tileCount << " base=0x" << std::hex << base.base
            << " entry=0x" << base.entry << std::dec
            << " bytesPerTile=" << imageSize
            << " changedBytes=" << stats.changedBytes
            << " patchRuns=" << stats.patchRuns
            << " outputBytes=" << stats.outputBytes << '\n';
}

void verify(const fs::path &imagePath, const fs::path *directory) {
  ipuimg::Image image(imagePath.string());
  for (unsigned tile = 0; tile < image.header().tileCount; ++tile) {
    const auto actual = image.tile(tile);
    if (directory &&
        actual !=
            loadElf(*directory / ("t_" + std::to_string(tile) + ".elf")).bytes)
      throw std::runtime_error("tile " + std::to_string(tile) +
                               " differs from source ELF");
  }
  std::cout << "tiles=" << image.header().tileCount << " base=0x" << std::hex
            << image.header().baseAddress << " entry=0x"
            << image.header().entryPoint << std::dec
            << " checksums=PASS sourceCompare="
            << (directory ? "PASS" : "skipped") << '\n';
}

void extract(const fs::path &imagePath, unsigned tile,
             const fs::path &outputPath) {
  ipuimg::Image image(imagePath.string());
  const auto bytes = image.tile(tile);
  std::ofstream output(outputPath, std::ios::binary | std::ios::trunc);
  output.write(reinterpret_cast<const char *>(bytes.data()), bytes.size());
  if (!output)
    throw std::runtime_error("cannot write " + outputPath.string());
  std::cout << "tile=" << tile << " address=0x" << std::hex
            << image.header().baseAddress << std::dec
            << " bytes=" << bytes.size() << '\n';
}

void redirect(const fs::path &imagePath, const fs::path &outputPath,
              std::uint32_t firstMarker, std::uint32_t secondMarker) {
  ipuimg::Image image(imagePath.string());
  if (image.header().entryPoint < image.header().baseAddress)
    throw std::runtime_error("entry precedes image base");
  const std::size_t entryOffset =
      image.header().entryPoint - image.header().baseAddress;
  std::vector<std::vector<std::uint8_t>> tiles;
  tiles.reserve(image.header().tileCount);
  std::uint32_t minimumTarget = std::numeric_limits<std::uint32_t>::max();
  std::uint32_t maximumTarget = 0;
  for (unsigned tile = 0; tile < image.header().tileCount; ++tile) {
    auto bytes = image.tile(tile);
    std::size_t markerOffset = bytes.size();
    for (std::size_t offset = 0; offset + 8 <= bytes.size(); offset += 4) {
      std::uint32_t words[2];
      std::memcpy(words, bytes.data() + offset, sizeof(words));
      if (words[0] != firstMarker || words[1] != secondMarker)
        continue;
      if (markerOffset != bytes.size())
        throw std::runtime_error("bare-metal marker is not unique");
      markerOffset = offset;
    }
    if (markerOffset == bytes.size())
      throw std::runtime_error("bare-metal marker not found on tile " +
                               std::to_string(tile));
    // Preserve the startup nop/sync3 rendezvous. The marked function repeats
    // marker0, marker1, sync3, so enter immediately after its 12-byte prefix.
    const std::uint32_t tileTarget =
        (image.header().baseAddress + markerOffset + 12 + 7) & ~7u;
    minimumTarget = std::min(minimumTarget, tileTarget);
    maximumTarget = std::max(maximumTarget, tileTarget);

    const std::uint32_t startup[2] = {
        0x19000000u | tileTarget, // setzi $m0, target
        0x43000000u,              // br $m0
    };
    if (entryOffset > bytes.size() || bytes.size() - entryOffset < 16)
      throw std::runtime_error("tile image is too small for entry redirect");
    std::memcpy(bytes.data() + entryOffset + 8, startup, sizeof(startup));
    tiles.push_back(std::move(bytes));
  }
  const auto stats =
      ipuimg::write(outputPath.string(), tiles, image.header().baseAddress,
                    image.header().entryPoint, image.header().templateTile);
  std::cout << "tiles=" << tiles.size() << " targetRange=0x" << std::hex
            << minimumTarget << "..0x" << maximumTarget << std::dec
            << " outputBytes=" << stats.outputBytes << " redirect=PASS\n";
}

void injectPhysical(const fs::path &imagePath, const fs::path &outputPath,
                    std::uint32_t marker, std::uint32_t bias, bool logicalSlots,
                    bool constant) {
  ipuimg::Image image(imagePath.string());
  std::vector<std::vector<std::uint8_t>> tiles;
  tiles.reserve(image.header().tileCount);
  for (unsigned tile = 0; tile < image.header().tileCount; ++tile) {
    auto bytes = image.tile(tile);
    std::size_t markerOffset = bytes.size();
    for (std::size_t offset = 0; offset + 4 <= bytes.size(); offset += 4) {
      std::uint32_t word;
      std::memcpy(&word, bytes.data() + offset, sizeof(word));
      if (word != marker)
        continue;
      if (markerOffset != bytes.size())
        throw std::runtime_error("input marker is not unique");
      markerOffset = offset;
    }
    if (markerOffset == bytes.size())
      throw std::runtime_error("input marker not found on tile " +
                               std::to_string(tile));
    const std::uint32_t value =
        constant ? bias
                 : (logicalSlots ? logicalToPhysical(tile) : tile) + bias;
    std::memcpy(bytes.data() + markerOffset, &value, sizeof(value));
    tiles.push_back(std::move(bytes));
  }
  const auto stats =
      ipuimg::write(outputPath.string(), tiles, image.header().baseAddress,
                    image.header().entryPoint, image.header().templateTile);
  std::cout << "tiles=" << tiles.size() << " marker=0x" << std::hex << marker
            << " bias=0x" << bias << std::dec << " mapping="
            << (constant ? "constant"
                         : (logicalSlots ? "logical-to-physical" : "index"))
            << " outputBytes=" << stats.outputBytes << " inject=PASS\n";
}

void injectExchangePlan(const fs::path &imagePath, const fs::path &outputPath,
                        unsigned senderLogical, unsigned receiverLogical,
                        unsigned count) {
  constexpr std::array<std::uint32_t, 9> marker = {
      0x0bad0000u, 0x0bad0001u, 0x0bad0002u, 0x0bad0003u, 0x0bad0004u,
      0x0bad0005u, 0x0bad0006u, 0x0bad0007u, 0x0bad0008u};
  ipuimg::Image image(imagePath.string());
  if (senderLogical >= image.header().tileCount ||
      receiverLogical >= image.header().tileCount)
    throw std::out_of_range("exchange-plan logical tile");
  constexpr std::uint32_t nop = 0x44f0f3f1u; // zero $m15
  std::array<std::uint32_t, 9> sender = {nop, 0x43a00000u};
  std::array<std::uint32_t, 9> receiver = sender;
  if (count != 0) {
    const auto transfer = xcom::assemble(senderLogical, receiverLogical, count);
    sender = transfer.sender;
    std::copy(transfer.receiver.begin() + 1, transfer.receiver.end(),
              receiver.begin());
    receiver[1] ^= xcom::logicalToPhysical(senderLogical);
    sender[0] = nop;
    receiver[0] = nop;
  }

  const unsigned senderPhysical = logicalToPhysical(senderLogical);
  std::vector<std::vector<std::uint8_t>> tiles;
  tiles.reserve(image.header().tileCount);
  for (unsigned tile = 0; tile < image.header().tileCount; ++tile) {
    auto bytes = image.tile(tile);
    const auto &row = tile == senderPhysical ? sender : receiver;
    unsigned replacements = 0;
    for (std::size_t offset = 0; offset + sizeof(marker) <= bytes.size();
         offset += sizeof(std::uint32_t)) {
      if (std::memcmp(bytes.data() + offset, marker.data(), sizeof(marker)) !=
          0)
        continue;
      std::memcpy(bytes.data() + offset, row.data(), sizeof(row));
      ++replacements;
    }
    if (replacements != 2)
      throw std::runtime_error("expected two exchange plan slots on tile " +
                               std::to_string(tile) + ", found " +
                               std::to_string(replacements));
    tiles.push_back(std::move(bytes));
  }
  const auto stats =
      ipuimg::write(outputPath.string(), tiles, image.header().baseAddress,
                    image.header().entryPoint, image.header().templateTile);
  std::cout << "tiles=" << tiles.size() << " sender=" << senderLogical
            << " receiver=" << receiverLogical << " count=" << count
            << " slotsPerTile=2 outputBytes=" << stats.outputBytes
            << " exchangePlanInject=PASS\n";
}

void injectHostInputs(const fs::path &imagePath, const fs::path &outputPath,
                      std::uint32_t receiverValue, std::uint32_t senderValue) {
  ipuimg::Image image(imagePath.string());
  const unsigned receiverPhysical =
      logicalToPhysical(bareproto::receiverLogicalTile);
  const unsigned senderPhysical =
      logicalToPhysical(bareproto::senderLogicalTile);
  std::vector<std::vector<std::uint8_t>> tiles;
  tiles.reserve(image.header().tileCount);
  for (unsigned tile = 0; tile < image.header().tileCount; ++tile) {
    auto bytes = image.tile(tile);
    const std::uint32_t value = tile == receiverPhysical ? receiverValue
                                : tile == senderPhysical ? senderValue
                                                         : 0;
    unsigned replacements = 0;
    unsigned roleReplacements = 0;
    for (std::size_t offset = 0; offset + sizeof(value) <= bytes.size();
         offset += sizeof(value)) {
      std::uint32_t word;
      std::memcpy(&word, bytes.data() + offset, sizeof(word));
      if (word == bareproto::imageInputMarker) {
        std::memcpy(bytes.data() + offset, &value, sizeof(value));
        ++replacements;
      } else if (word == 0x0face000u) {
        const std::uint32_t setPhysicalTile = 0x19300000u | tile;
        std::memcpy(bytes.data() + offset, &setPhysicalTile,
                    sizeof(setPhysicalTile));
        ++roleReplacements;
      }
    }
    if (replacements != 1 || roleReplacements != 1)
      throw std::runtime_error("unexpected host input count on tile " +
                               std::to_string(tile) + ": found " +
                               std::to_string(replacements) + " input, " +
                               std::to_string(roleReplacements) + " role");
    tiles.push_back(std::move(bytes));
  }
  const auto stats =
      ipuimg::write(outputPath.string(), tiles, image.header().baseAddress,
                    image.header().entryPoint, image.header().templateTile);
  std::cout << "tiles=" << tiles.size() << " receiverValue=0x" << std::hex
            << receiverValue << " senderValue=0x" << senderValue << std::dec
            << " outputBytes=" << stats.outputBytes
            << " hostInputInject=PASS\n";
}

void patchWord(const fs::path &imagePath, const fs::path &outputPath,
               unsigned tile, std::uint32_t address, std::uint32_t expected,
               std::uint32_t replacement) {
  ipuimg::Image image(imagePath.string());
  if (tile >= image.header().tileCount ||
      address < image.header().baseAddress || (address & 3u) != 0)
    throw std::out_of_range("word patch location");
  const std::size_t offset = address - image.header().baseAddress;
  std::vector<std::vector<std::uint8_t>> tiles;
  tiles.reserve(image.header().tileCount);
  for (unsigned index = 0; index < image.header().tileCount; ++index) {
    auto bytes = image.tile(index);
    if (index == tile) {
      if (offset > bytes.size() || sizeof(expected) > bytes.size() - offset)
        throw std::out_of_range("word patch address");
      std::uint32_t observed;
      std::memcpy(&observed, bytes.data() + offset, sizeof(observed));
      if (observed != expected)
        throw std::runtime_error("word patch expected value mismatch");
      std::memcpy(bytes.data() + offset, &replacement, sizeof(replacement));
    }
    tiles.push_back(std::move(bytes));
  }
  const auto stats =
      ipuimg::write(outputPath.string(), tiles, image.header().baseAddress,
                    image.header().entryPoint, image.header().templateTile);
  std::cout << "tile=" << tile << " address=0x" << std::hex << address
            << " old=0x" << expected << " new=0x" << replacement << std::dec
            << " outputBytes=" << stats.outputBytes << " patchWord=PASS\n";
}

void redirectD2hSource(const fs::path &imagePath, const fs::path &outputPath,
                       unsigned logicalTile, std::uint32_t byteOffset) {
  constexpr std::uint32_t sendOpcodeMask = 0xf8000000u;
  constexpr std::uint32_t sendOpcode = 0x78000000u;
  constexpr std::uint32_t sendAddressMask = 0x001ffff8u;
  constexpr std::uint32_t delayOpcodeMask = 0xfff80000u;
  constexpr std::uint32_t delayOpcode = 0x40a00000u;
  constexpr std::uint32_t syncReceive = 0x41800000u;
  constexpr std::uint32_t setDowncountOne = 0x19800001u;
  constexpr std::uint32_t putDowncount = 0x438080a6u;
  if ((byteOffset & 3u) != 0)
    throw std::invalid_argument("D2H source offset must be word aligned");

  ipuimg::Image image(imagePath.string());
  if (logicalTile >= image.header().tileCount)
    throw std::out_of_range("D2H logical tile");
  const unsigned physicalTile = logicalToPhysical(logicalTile);
  std::vector<std::vector<std::uint8_t>> tiles;
  tiles.reserve(image.header().tileCount);
  std::uint32_t oldAddress = 0;
  std::uint32_t newAddress = 0;
  unsigned matches = 0;
  for (unsigned tile = 0; tile < image.header().tileCount; ++tile) {
    auto bytes = image.tile(tile);
    if (tile == physicalTile) {
      for (std::size_t offset = 0; offset + 7 * sizeof(std::uint32_t) <= bytes.size();
           offset += sizeof(std::uint32_t)) {
        std::array<std::uint32_t, 7> words;
        std::memcpy(words.data(), bytes.data() + offset, sizeof(words));
        const auto isSend = [&](unsigned index, unsigned count) {
          return (words[index] & sendOpcodeMask) == sendOpcode &&
                 ((words[index] >> 21) & 0x3fu) + 1 == count;
        };
        if (words[0] != setDowncountOne || words[1] != putDowncount ||
            !isSend(2, 2) || !isSend(3, 16) ||
            (words[4] & delayOpcodeMask) != delayOpcode || !isSend(5, 2) ||
            words[6] != syncReceive)
          continue;
        oldAddress = ((words[3] & sendAddressMask) >> 3) * 4;
        newAddress = oldAddress + byteOffset;
        words[3] = xcom::detail::send(15, words[3] & 7u, newAddress >> 2);
        std::memcpy(bytes.data() + offset, words.data(), sizeof(words));
        ++matches;
      }
    }
    tiles.push_back(std::move(bytes));
  }
  if (matches != 1)
    throw std::runtime_error("expected one 64-byte D2H target operation, found " +
                             std::to_string(matches));
  const auto stats =
      ipuimg::write(outputPath.string(), tiles, image.header().baseAddress,
                    image.header().entryPoint, image.header().templateTile);
  std::cout << "logicalTile=" << logicalTile << " physicalTile=" << physicalTile
            << " oldSource=0x" << std::hex << oldAddress << " newSource=0x"
            << newAddress << std::dec << " outputBytes=" << stats.outputBytes
            << " redirectD2hSource=PASS\n";
}
} // namespace

int main(int argc, char **argv) try {
  if (argc >= 2 && std::string(argv[1]) == "pack" && (argc == 5 || argc == 6)) {
    pack(argv[2], argv[3], std::stoul(argv[4]),
         argc == 6 ? std::stoul(argv[5]) : 0);
  } else if (argc >= 2 && std::string(argv[1]) == "verify" &&
             (argc == 3 || argc == 4)) {
    const fs::path directory = argc == 4 ? fs::path(argv[3]) : fs::path();
    verify(argv[2], argc == 4 ? &directory : nullptr);
  } else if (argc == 5 && std::string(argv[1]) == "extract") {
    extract(argv[2], std::stoul(argv[3]), argv[4]);
  } else if (argc == 6 && std::string(argv[1]) == "redirect") {
    redirect(argv[2], argv[3], std::stoul(argv[4], nullptr, 0),
             std::stoul(argv[5], nullptr, 0));
  } else if (argc == 6 && std::string(argv[1]) == "inject-physical") {
    injectPhysical(argv[2], argv[3], std::stoul(argv[4], nullptr, 0),
                   std::stoul(argv[5], nullptr, 0), false, false);
  } else if (argc == 6 && std::string(argv[1]) == "inject-tile-id") {
    injectPhysical(argv[2], argv[3], std::stoul(argv[4], nullptr, 0),
                   std::stoul(argv[5], nullptr, 0), true, false);
  } else if (argc == 6 && std::string(argv[1]) == "inject-constant") {
    injectPhysical(argv[2], argv[3], std::stoul(argv[4], nullptr, 0),
                   std::stoul(argv[5], nullptr, 0), false, true);
  } else if (argc == 7 && std::string(argv[1]) == "inject-exchange-plan") {
    injectExchangePlan(argv[2], argv[3], std::stoul(argv[4], nullptr, 0),
                       std::stoul(argv[5], nullptr, 0),
                       std::stoul(argv[6], nullptr, 0));
  } else if (argc == 6 && std::string(argv[1]) == "inject-host-inputs") {
    injectHostInputs(argv[2], argv[3], std::stoul(argv[4], nullptr, 0),
                     std::stoul(argv[5], nullptr, 0));
  } else if (argc == 8 && std::string(argv[1]) == "patch-word") {
    patchWord(argv[2], argv[3], std::stoul(argv[4], nullptr, 0),
              std::stoul(argv[5], nullptr, 0),
              std::stoul(argv[6], nullptr, 0),
              std::stoul(argv[7], nullptr, 0));
  } else if (argc == 6 && std::string(argv[1]) == "redirect-d2h-source") {
    redirectD2hSource(argv[2], argv[3], std::stoul(argv[4], nullptr, 0),
                      std::stoul(argv[5], nullptr, 0));
  } else {
    std::cerr << "usage:\n"
              << "  ipu_image pack ELF_DIR OUTPUT TILE_COUNT [TEMPLATE_TILE]\n"
              << "  ipu_image verify IMAGE [ELF_DIR]\n"
              << "  ipu_image extract IMAGE TILE OUTPUT_RAW\n"
              << "  ipu_image redirect IMAGE OUTPUT MARKER0 MARKER1\n"
              << "  ipu_image inject-physical IMAGE OUTPUT MARKER BIAS\n"
              << "  ipu_image inject-tile-id IMAGE OUTPUT MARKER BIAS\n"
              << "  ipu_image inject-constant IMAGE OUTPUT MARKER VALUE\n"
              << "  ipu_image inject-exchange-plan IMAGE OUTPUT SENDER"
                 " RECEIVER COUNT\n"
              << "  ipu_image inject-host-inputs IMAGE OUTPUT RECEIVER_VALUE"
                 " SENDER_VALUE\n"
              << "  ipu_image patch-word IMAGE OUTPUT TILE ADDRESS EXPECTED"
                 " REPLACEMENT\n"
              << "  ipu_image redirect-d2h-source IMAGE OUTPUT LOGICAL_TILE"
                 " BYTE_OFFSET\n";
    return 2;
  }
  return 0;
} catch (const std::exception &error) {
  std::cerr << "ipu_image: " << error.what() << '\n';
  return 1;
}