#include "ConflictSet.h"
#include "Internal.h"
#include <byteswap.h>
#include <cstdint>

#define ANKERL_NANOBENCH_IMPLEMENT
#include "third_party/nanobench.h"

std::string toKey(int n) {
  std::string result;
  result.resize(32);

  for (int i = 0; i < 32; ++i) {
    result[i] = n & (1 << (31 - i)) ? '1' : '0';
  }
  return result;
}

constexpr int kNumKeys = 100000;
// A range read, a point read, and a point write. Range writes can erase
// keys, and we don't want to change the number of keys stored in the
// conflict set.
constexpr int kOpsPerTx = 100;

constexpr int kPrefixLen = 0;

std::span<const uint8_t> makeKey(Arena &arena, int index) {
  auto result = std::span<uint8_t>{new (arena) uint8_t[4], 4};
  index = __builtin_bswap32(index);
  memcpy(result.data(), &index, 4);
  return result;
}

template <class ConflictSet_> void benchConflictSet() {
  ankerl::nanobench::Bench bench;
  ConflictSet_ cs{0};

  bench.batch(kOpsPerTx);

  int64_t version = 0;

  // Populate conflict set
  Arena arena;
  {
    std::vector<ConflictSet::WriteRange> writes;
    writes.reserve(kNumKeys);
    for (int i = 0; i < kNumKeys; ++i) {
      auto key = makeKey(arena, i);
      ConflictSet::WriteRange conflict;
      conflict.begin.p = key.data();
      conflict.begin.len = key.size();
      conflict.end.len = 0;
      conflict.writeVersion = version + 1;
      writes.push_back(conflict);
    }
    cs.addWrites(writes.data(), writes.size());
    ++version;
  }

  // I don't know why std::less didn't work /shrug
  struct Less {
    bool operator()(const std::span<const uint8_t> &lhs,
                    const std::span<const uint8_t> &rhs) const {
      return lhs < rhs;
    }
  };
  auto points = set<std::span<const uint8_t>, Less>(arena);

  // Two points for each range read, one for each point read, and one for each
  // point write
  while (points.size() < kOpsPerTx * 2 + 1) {
    // TODO don't use rand?
    points.insert(makeKey(arena, rand() % kNumKeys));
  }

  // Make short-circuiting non-trivial
  {
    std::vector<ConflictSet::WriteRange> writes;
    writes.reserve(kNumKeys);
    for (int i = 0; i < kNumKeys; ++i) {
      auto key = makeKey(arena, i);
      if (points.find(key) != points.end()) {
        continue;
      }
      ConflictSet::WriteRange conflict;
      conflict.begin.p = key.data();
      conflict.begin.len = key.size();
      conflict.end.len = 0;
      conflict.writeVersion = version + 1;
      writes.push_back(conflict);
    }
    cs.addWrites(writes.data(), writes.size());
    ++version;
  }

  {
    std::vector<ConflictSet::ReadRange> reads;
    auto iter = points.begin();
    for (int i = 0; i < kOpsPerTx; ++i) {
      ConflictSet::ReadRange r;
      r.begin.p = iter->data();
      r.begin.len = iter->size();
      r.end.len = 0;
      r.readVersion = version - 1;
      reads.push_back(r);
      ++iter;
    }

    auto *results = new (arena) ConflictSet::Result[kOpsPerTx];

    bench.run("radix tree (point reads)",
              [&]() { cs.check(reads.data(), results, kOpsPerTx); });
  }

  {
    std::vector<ConflictSet::ReadRange> reads;
    auto iter = points.begin();
    for (int i = 0; i < kOpsPerTx; ++i) {
      auto begin = *iter++;
      auto end = *iter++;
      ConflictSet::ReadRange r;
      r.begin.p = begin.data();
      r.begin.len = begin.size();
      r.end.p = end.data();
      r.end.len = end.size();
      r.readVersion = version - 1;
      reads.push_back(r);
    }

    auto *results = new (arena) ConflictSet::Result[kOpsPerTx];

    bench.run("radix tree (range reads)",
              [&]() { cs.check(reads.data(), results, kOpsPerTx); });
  }

  {
    std::vector<ConflictSet::WriteRange> writes;
    auto iter = points.begin();
    for (int i = 0; i < kOpsPerTx; ++i) {
      ConflictSet::WriteRange w;
      w.begin.p = iter->data();
      w.begin.len = iter->size();
      w.end.len = 0;
      writes.push_back(w);
      ++iter;
    }

    bench.run("radix tree (point writes)", [&]() {
      auto v = ++version;
      for (auto &w : writes) {
        w.writeVersion = v;
      }
      cs.addWrites(writes.data(), writes.size());
    });
  }

  {
    std::vector<ConflictSet::WriteRange> writes;
    auto iter = points.begin();
    for (int i = 0; i < kOpsPerTx - 1; ++i) {
      auto begin = *iter++;
      auto end = *iter++;
      ConflictSet::WriteRange w;
      w.begin.p = begin.data();
      w.begin.len = begin.size();
      w.end.p = end.data();
      w.end.len = end.size();
      writes.push_back(w);
    }

    bench.run("radix tree (range writes)", [&]() {
      auto v = ++version;
      for (auto &w : writes) {
        w.writeVersion = v;
      }
      cs.addWrites(writes.data(), writes.size());
    });
  }
}

int main(void) { benchConflictSet<ConflictSet>(); }