#include "ipu_image_writer.hpp"

#include <algorithm>
#include <cstdint>
#include <fstream>
#include <iostream>
#include <limits>
#include <stdexcept>
#include <string>
#include <vector>

#include <poplar/Executable.hpp>
#include <poplar/Module.hpp>

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

int main(int argc, char **argv) try {
  if (argc != 3 && argc != 4) {
    std::cerr
        << "usage: popef_to_ipu_image EXECUTABLE OUTPUT [TEMPLATE_TILE]\n";
    return 2;
  }
  const unsigned templateTile = argc == 4 ? std::stoul(argv[3]) : 0;
  std::ifstream input(argv[1], std::ios::binary);
  if (!input)
    throw std::runtime_error("cannot open executable");
  poplar::Executable executable = poplar::Executable::deserialize(input);
  poplar::Module module(executable.getImpl());

  struct Segment {
    std::size_t tile;
    std::uint32_t address;
    std::vector<std::uint8_t> bytes;
  };
  std::vector<Segment> segments;
  std::size_t tileCount = 0;
  std::uint64_t low = std::numeric_limits<std::uint32_t>::max();
  std::uint64_t high = 0;
  module.forEachLoadableSegment([&](std::size_t tile, std::size_t address,
                                    std::size_t size, const char *data) {
    if (address > std::numeric_limits<std::uint32_t>::max() ||
        size > std::numeric_limits<std::uint32_t>::max() ||
        address + size > std::numeric_limits<std::uint32_t>::max())
      throw std::runtime_error("loadable segment exceeds 32-bit address space");
    Segment segment{tile, static_cast<std::uint32_t>(address),
                    std::vector<std::uint8_t>(data, data + size)};
    segments.push_back(std::move(segment));
    tileCount = std::max(tileCount, tile + 1);
    low = std::min<std::uint64_t>(low, address);
    high = std::max<std::uint64_t>(high, address + size);
  });
  if (segments.empty() || templateTile >= tileCount)
    throw std::runtime_error("executable has no usable tile images");

  const std::uint64_t firstSegment = low;
  // Poplar omits the zero-fill .secondaryLower PT_LOAD segment. It remains an
  // address reservation: the secondary loader installs the first enumerated
  // byte at its linked address rather than at the start of tile SRAM.
  constexpr std::uint32_t ipu21LoadBase = 0x4c000;
  if (low < ipu21LoadBase)
    throw std::runtime_error("loadable segment precedes IPU21 tile memory");
  const auto loadBase = static_cast<std::uint32_t>(firstSegment);

  std::vector<std::vector<std::uint8_t>> images(
      tileCount, std::vector<std::uint8_t>(high - firstSegment, 0));
  for (const auto &segment : segments) {
    const unsigned physicalTile = logicalToPhysical(segment.tile);
    std::copy(segment.bytes.begin(), segment.bytes.end(),
              images[physicalTile].begin() + (segment.address - firstSegment));
  }

  const auto stats =
      ipuimg::write(argv[2], images, loadBase, loadBase, templateTile);
  std::cout << "tiles=" << tileCount << " base=0x" << std::hex << loadBase
            << " entry=0x" << loadBase << std::dec
            << " bytesPerTile=" << images[0].size()
            << " changedBytes=" << stats.changedBytes
            << " patchRuns=" << stats.patchRuns
            << " outputBytes=" << stats.outputBytes << '\n';
  return 0;
} catch (const std::exception &error) {
  std::cerr << "popef_to_ipu_image: " << error.what() << '\n';
  return 1;
}