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

#if SHOW_MEMORY
void showMemory(const ConflictSet &cs);
#endif

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

constexpr int kNumKeys = 1000000;

constexpr int kOpsPerTx = 100;

constexpr int kPrefixLen = 0;

constexpr int kMvccWindow = 100000;

std::span<const uint8_t> makeKey(Arena &arena, int index) {

  auto result =
      std::span<uint8_t>{new (arena) uint8_t[4 + kPrefixLen], 4 + kPrefixLen};
  index = __builtin_bswap32(index);
  memset(result.data(), 0, kPrefixLen);
  memcpy(result.data() + kPrefixLen, &index, 4);

  return result;
}

ConflictSet::ReadRange singleton(Arena &arena, std::span<const uint8_t> key) {
  auto r =
      std::span<uint8_t>(new (arena) uint8_t[key.size() + 1], key.size() + 1);
  memcpy(r.data(), key.data(), key.size());
  r[key.size()] = 0;
  return {{key.data(), int(key.size())}, {r.data(), int(r.size())}, 0};
}

ConflictSet::ReadRange prefixRange(Arena &arena, std::span<const uint8_t> key) {
  int index;
  for (index = key.size() - 1; index >= 0; index--)
    if ((key[index]) != 255)
      break;

  // Must not be called with a string that consists only of zero or more '\xff'
  // bytes.
  if (index < 0) {
    assert(false);
  }

  auto r = std::span<uint8_t>(new (arena) uint8_t[index + 1], index + 1);
  memcpy(r.data(), key.data(), index + 1);
  r[r.size() - 1]++;
  return {{key.data(), int(key.size())}, {r.data(), int(r.size())}, 0};
}

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 w;
      auto r = singleton(arena, key);
      w.begin.p = r.begin.p;
      w.begin.len = r.begin.len;
      w.end.p = r.end.p;
      w.end.len = 0;
      writes.push_back(w);
    }
    cs.addWrites(writes.data(), writes.size(), version + 1);
    ++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);

  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;
    auto iter = points.begin();
    ++iter; // Complement of the set we'll be reading with range reads. Almost.
    for (int i = 0; i < kOpsPerTx; ++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);
    }
    cs.addWrites(writes.data(), kOpsPerTx, version + 1);
    ++version;
  }

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

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

    bench.run("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 r = prefixRange(arena, *iter);
      r.readVersion = version - 1;
      reads.push_back(r);
      ++iter;
    }

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

    bench.run("prefix 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("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;
      auto r = singleton(arena, *iter);
      w.begin.p = r.begin.p;
      w.begin.len = r.begin.len;
      w.end.p = r.end.p;
      w.end.len = 0;
      writes.push_back(w);
      ++iter;
    }

    while (version < kMvccWindow) {
      auto v = ++version;
      cs.addWrites(writes.data(), 1, v);
    }

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

  {
    std::vector<ConflictSet::WriteRange> writes;
    auto iter = points.begin();
    for (int i = 0; i < kOpsPerTx; ++i) {
      ConflictSet::WriteRange w;
      auto r = prefixRange(arena, *iter);
      w.begin.p = r.begin.p;
      w.begin.len = r.begin.len;
      w.end.p = r.end.p;
      w.end.len = r.end.len;
      writes.push_back(w);
      ++iter;
    }

    bench.run("prefix writes", [&]() {
      auto v = ++version;
      cs.addWrites(writes.data(), writes.size(), v);
    });
  }

  {
    std::vector<ConflictSet::WriteRange> writes;
    auto iter = points.begin();
    for (int i = 0; i < kOpsPerTx; ++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("range writes", [&]() {
      auto v = ++version;
      cs.addWrites(writes.data(), writes.size(), v);
    });
  }

  bench.batch(1);

  bench.warmup(10000);

  {
    bench.run("monotonic increasing point writes", [&]() {
      auto v = ++version;
      ConflictSet::WriteRange w;

      uint8_t b[9];
      b[8] = 0;
      auto x = __builtin_bswap64(version);
      memcpy(b, &x, 8);

      w.begin.p = b;
      w.begin.len = 8;
      w.end.len = 0;
      w.end.p = b;
      cs.addWrites(&w, 1, v);
      cs.setOldestVersion(version - kMvccWindow);
    });
  }
}

void benchWorstCaseForRadixRangeRead() {
  ankerl::nanobench::Bench bench;
  ConflictSet cs{0};

  int64_t version = 0;

  constexpr int kKeyLength = 50;

  for (int i = 0; i < kKeyLength; ++i) {
    for (int j = 0; j < 256; ++j) {
      auto b = std::vector<uint8_t>(i, 0);
      b.push_back(j);
      auto e = std::vector<uint8_t>(i, 255);
      e.push_back(j);
      weaselab::ConflictSet::WriteRange w[] = {{
                                                   {b.data(), int(b.size())},
                                                   {nullptr, 0},
                                               },
                                               {
                                                   {e.data(), int(e.size())},
                                                   {nullptr, 0},
                                               }};
      cs.addWrites(w, sizeof(w) / sizeof(w[0]), version);
    }
  }

  // Defeat short-circuiting on the left
  {
    auto k = std::vector<uint8_t>(kKeyLength, 0);
    weaselab::ConflictSet::WriteRange w[] = {
        {
            {k.data(), int(k.size())},
            {nullptr, 0},
        },
    };
    cs.addWrites(w, sizeof(w) / sizeof(w[0]), 2);
  }

  // Defeat short-circuiting on the right
  {
    auto k = std::vector<uint8_t>(kKeyLength, 255);
    weaselab::ConflictSet::WriteRange w[] = {
        {
            {k.data(), int(k.size())},
            {nullptr, 0},
        },
    };
    cs.addWrites(w, sizeof(w) / sizeof(w[0]), 2);
  }

  auto begin = std::vector<uint8_t>(kKeyLength - 1, 0);
  begin.push_back(1);
  auto end = std::vector<uint8_t>(kKeyLength - 1, 255);
  end.push_back(254);

  weaselab::ConflictSet::Result result;
  weaselab::ConflictSet::ReadRange r{
      {begin.data(), int(begin.size())}, {end.data(), int(end.size())}, 1};

  bench.run("worst case for radix tree", [&]() {
    result = weaselab::ConflictSet::TooOld;
    cs.check(&r, &result, 1);
    if (result != weaselab::ConflictSet::Commit) {
      abort();
    }
  });
}

int main(void) {
  benchConflictSet();
  benchWorstCaseForRadixRangeRead();
}