#include <cstdint>
#include <cstdlib>
#include <fstream>
#include <iostream>
#include <string>
#include <vector>

#include <poplar/DeviceManager.hpp>
#include <poplar/Engine.hpp>
#include <poplar/Graph.hpp>
#include <poplar/Program.hpp>
#include <poplar/Target.hpp>

int main(int argc, char **argv) {
  if (argc != 2 && argc != 3) {
    std::cerr
        << "usage: build_baremetal OUTPUT_EXECUTABLE"
           " [signal|result|result-alt|host-exchange|constant|image-input|"
           "image-constant]\n";
    return 2;
  }
  const bool signalOnly = argc == 3 && std::string(argv[2]) == "signal";
  const bool result = argc == 3 && std::string(argv[2]) == "result";
  const bool alternateResult =
      argc == 3 && std::string(argv[2]) == "result-alt";
  const bool hostExchange =
      argc == 3 && std::string(argv[2]) == "host-exchange";
  const bool constantCheck = argc == 3 && std::string(argv[2]) == "constant";
  const bool imageInput = argc == 3 && std::string(argv[2]) == "image-input";
  const bool imageConstant =
      argc == 3 && std::string(argv[2]) == "image-constant";
  if (argc == 3 && !signalOnly && !result && !alternateResult &&
      !hostExchange && !constantCheck && !imageInput && !imageConstant) {
    std::cerr << "unknown bare-metal mode " << argv[2] << '\n';
    return 2;
  }

  auto target = poplar::Target::createIPUTarget(1, "C600");
  std::vector<poplar::Device> devices;
  if (std::getenv("BAREMETAL_DEVICE_TARGET")) {
    poplar::DeviceManager manager;
    devices = manager.getDevices(poplar::TargetType::IPU, 1);
    if (devices.empty() || !devices[0].attach()) {
      std::cerr << "failed to attach an IPU for its target description\n";
      return 3;
    }
    target = devices[0].getTarget();
  }
  poplar::Graph graph(target);
  graph.addCodelets(std::vector<std::string>{"baremetal_codelets.cpp"}, "-O2",
                    "ipu21");
  const char *vertexName = signalOnly        ? "BareSignal"
                           : result          ? "BareResult"
                           : alternateResult ? "BareResultAlternate"
                           : hostExchange    ? "BareHostExchange"
                           : imageInput      ? "BareImageInputCheck"
                           : imageConstant   ? "BareImageConstantCheck"
                           : constantCheck   ? "BareHostConstantCheck"
                                             : "BareHostWriteCheck";
  auto computeSet = graph.addComputeSet(vertexName);
  for (unsigned tile = 0; tile < target.getNumTiles(); ++tile) {
    auto vertex = graph.addVertex(computeSet, vertexName);
    graph.setTileMapping(vertex, tile);
  }

  poplar::Engine engine(graph, poplar::program::Execute(computeSet));
  std::ofstream output(argv[1], std::ios::binary);
  engine.serializeExecutable(output);
  if (!output) {
    std::cerr << "failed to write " << argv[1] << '\n';
    return 1;
  }
  std::cout << "tiles=" << target.getNumTiles() << " executable=PASS\n";
}