#include "ipu_bootloader_frames.hpp"
#include "ipu_image.hpp"

#include <fstream>
#include <iostream>
#include <stdexcept>
#include <cstring>

int main(int argc, char **argv) try {
  if (argc != 5 && argc != 6) {
    std::cerr << "usage: ipu_bootloader_frames IMAGE FIRST_TILE TILE_COUNT"
                 " OUTPUT [--transport]\n";
    return 2;
  }
  const bool transport = argc == 6 && std::string(argv[5]) == "--transport";
  if (argc == 6 && !transport) throw std::runtime_error("unknown option");
  ipuimg::Image image(argv[1]);
  const unsigned firstTile = std::stoul(argv[2]);
  const unsigned tileCount = std::stoul(argv[3]);
  if (firstTile > image.header().tileCount ||
      tileCount > image.header().tileCount - firstTile)
    throw std::runtime_error("tile range exceeds image");
  std::ofstream output(argv[4], std::ios::binary | std::ios::trunc);
  if (!output) throw std::runtime_error("cannot create output file");
  std::uint64_t frames = 0;
  constexpr std::size_t frameOffset = 4096;
  constexpr std::size_t transportSize = 0x2842000;
  std::vector<std::uint8_t> transportBytes;
  std::size_t cursor = frameOffset;
  if (transport) transportBytes.resize(transportSize);
  for (unsigned tile = firstTile; tile < firstTile + tileCount; ++tile) {
    const auto framed = ipuboot::frameTile(tile, image.tile(tile));
    if (transport) {
      const unsigned slot = tile - firstTile;
      const std::uint32_t descriptor[2] = {
          static_cast<std::uint32_t>(cursor),
          static_cast<std::uint32_t>(framed.size() / ipuboot::frameSize)};
      if (slot >= 64 || framed.size() > transportBytes.size() - cursor)
        throw std::runtime_error("transport batch overflow");
      std::memcpy(transportBytes.data() + slot * sizeof(descriptor), descriptor,
                  sizeof(descriptor));
      std::memcpy(transportBytes.data() + cursor, framed.data(), framed.size());
      cursor += framed.size();
    } else {
      output.write(reinterpret_cast<const char *>(framed.data()), framed.size());
    }
    frames += framed.size() / ipuboot::frameSize;
  }
  if (transport)
    output.write(reinterpret_cast<const char *>(transportBytes.data()),
                 transportBytes.size());
  if (!output) throw std::runtime_error("cannot write output file");
  std::cout << "firstPhysicalTile=" << firstTile << " tiles=" << tileCount
            << " frames=" << frames << " bytes="
            << frames * ipuboot::frameSize << '\n';
  return 0;
} catch (const std::exception &error) {
  std::cerr << "ipu_bootloader_frames: " << error.what() << '\n';
  return 1;
}