#include "circle/schema_generated.h"

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

namespace
{

std::vector<uint8_t> build_model(uint64_t external_offset, uint64_t external_size)
{
  flatbuffers::FlatBufferBuilder fbb;

  std::vector<flatbuffers::Offset<circle::OperatorCode>> opcode_offsets;
  auto opcodes = fbb.CreateVector(opcode_offsets);

  std::vector<flatbuffers::Offset<circle::Tensor>> tensor_offsets;
  std::vector<int32_t> io_indices;
  std::vector<flatbuffers::Offset<circle::Operator>> operator_offsets;
  auto subgraph = circle::CreateSubGraphDirect(fbb, &tensor_offsets, &io_indices, &io_indices,
                                               &operator_offsets, "empty");
  std::vector<flatbuffers::Offset<circle::SubGraph>> subgraph_offsets{subgraph};
  auto subgraphs = fbb.CreateVector(subgraph_offsets);

  auto empty_data = fbb.CreateVector(std::vector<uint8_t>{});
  auto default_buffer = circle::CreateBuffer(fbb, empty_data, 0, 0);
  auto external_buffer = circle::CreateBuffer(fbb, 0, external_offset, external_size);
  std::vector<flatbuffers::Offset<circle::Buffer>> buffer_offsets{default_buffer, external_buffer};
  auto buffers = fbb.CreateVector(buffer_offsets);

  auto desc = fbb.CreateString("external buffer offset/size boundary probe");
  auto model = circle::CreateModel(fbb, 3, opcodes, subgraphs, desc, buffers);
  circle::FinishModelBuffer(fbb, model);

  const auto *ptr = fbb.GetBufferPointer();
  return std::vector<uint8_t>(ptr, ptr + fbb.GetSize());
}

bool write_file(const std::string &path, const std::vector<uint8_t> &data)
{
  std::ofstream out(path, std::ios::binary);
  if (!out)
    return false;
  out.write(reinterpret_cast<const char *>(data.data()), static_cast<std::streamsize>(data.size()));
  return static_cast<bool>(out);
}

} // namespace

int main(int argc, char **argv)
{
  if (argc != 2)
  {
    std::cerr << "usage: " << argv[0] << " out.circle\n";
    return 2;
  }

  auto first = build_model(2, 0x1000);
  auto final_model = build_model(first.size() - 1, 0x1000);

  if (!write_file(argv[1], final_model))
  {
    std::cerr << "failed to write " << argv[1] << "\n";
    return 1;
  }

  std::cout << "wrote " << argv[1] << " bytes=" << final_model.size()
            << " external_offset=" << (final_model.size() - 1)
            << " external_size=4096\n";
  return 0;
}