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

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

#include "third_party/nanobench.h"

constexpr int kNumKeys = 1000000;

constexpr int kOpsPerTx = 100;

constexpr int kPrefixLen = 0;

constexpr int kMvccWindow = 100000;

TrivialSpan makeKey(Arena &arena, int index) {

  uint8_t *buf = new (arena) uint8_t[4 + kPrefixLen];
  auto result = TrivialSpan{buf, 4 + kPrefixLen};
  index = __builtin_bswap32(index);
  memset(buf, 0, kPrefixLen);
  memcpy(buf, &index, 4);

  return result;
}

ConflictSet::ReadRange singleton(Arena &arena, TrivialSpan key) {
  uint8_t *buf = new (arena) uint8_t[key.size() + 1];
  auto r = TrivialSpan(buf, key.size() + 1);
  memcpy(buf, key.data(), key.size());
  buf[key.size()] = 0;
  return {{key.data(), int(key.size())}, {r.data(), int(r.size())}, 0};
}

ConflictSet::ReadRange prefixRange(Arena &arena, TrivialSpan 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);
  }

  uint8_t *buf = new (arena) uint8_t[index + 1];
  auto r = TrivialSpan(buf, index + 1);
  memcpy(buf, key.data(), index + 1);
  buf[r.size() - 1]++;
  return {{key.data(), int(key.size())}, {r.data(), int(r.size())}, 0};
}

void benchConflictSet() {
  ankerl::nanobench::Bench bench;
  bench.minEpochIterations(10000);
  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;
  }

  auto points = set<TrivialSpan, std::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);
    });
  }
}

constexpr int kKeyLenForWorstCase = 50;

ConflictSet worstCaseConflictSetForRadixRangeRead(int cardinality) {
  ConflictSet cs{0};

  for (int i = 0; i < kKeyLenForWorstCase; ++i) {
    for (int j = 0; j < cardinality; ++j) {
      auto b = std::vector<uint8_t>(i, 0);
      b.push_back(j);
      auto e = std::vector<uint8_t>(i, 255);
      e.push_back(255 - j);
      weaselab::ConflictSet::WriteRange w[] = {{
                                                   {b.data(), int(b.size())},
                                                   {nullptr, 0},
                                               },
                                               {
                                                   {e.data(), int(e.size())},
                                                   {nullptr, 0},
                                               }};
      std::sort(std::begin(w), std::end(w),
                [](const auto &lhs, const auto &rhs) {
                  int cl = std::min(lhs.begin.len, rhs.begin.len);
                  if (cl > 0) {
                    int c = memcmp(lhs.begin.p, rhs.begin.p, cl);
                    if (c != 0) {
                      return c < 0;
                    }
                  }
                  return lhs.begin.len < rhs.begin.len;
                });
      cs.addWrites(w, sizeof(w) / sizeof(w[0]), 0);
    }
  }

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

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

  return cs;
}

void benchWorstCaseForRadixRangeRead() {
  ankerl::nanobench::Bench bench;

  std::unique_ptr<ConflictSet> cs[256];
  for (int i = 0; i < 256; ++i) {
    cs[i] =
        std::make_unique<ConflictSet>(worstCaseConflictSetForRadixRangeRead(i));
  }

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

  weaselab::ConflictSet::ReadRange r[] = {
      {{begin.data(), int(begin.size())}, {end.data(), int(end.size())}, 0},
  };
  weaselab::ConflictSet::Result results[sizeof(r) / sizeof(r[0])];
  for (auto &result : results) {
    result = weaselab::ConflictSet::TooOld;
  }
  bench.batch(sizeof(r) / sizeof(r[0]));

  bench.run("worst case for radix tree", [&]() {
    for (int i = 0; i < 256; ++i) {
      cs[i]->check(r, results, sizeof(r) / sizeof(r[0]));
      for (auto result : results) {
        if (result != weaselab::ConflictSet::Commit) {
          abort();
        }
      }
    }
  });

  // for (int i = 0; i < 256; ++i) {
  //   bench.run("worst case for radix tree, span " + std::to_string(i), [&]() {
  //     result = weaselab::ConflictSet::TooOld;
  //     cs[i]->check(&r, &result, 1);
  //     if (result != weaselab::ConflictSet::Commit) {
  //       abort();
  //     }
  //   });
  // }
}

void benchCreateAndDestroy() {
  ankerl::nanobench::Bench bench;

  bench.run("create and destroy", [&]() { ConflictSet cs{0}; });
}

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