#include "bag.h"
#include <catch2/catch.hpp>
#include <emscripten/atomic.h>
#include <emscripten/wasm_worker.h>
#include <random>
#include <vector>

struct TestNode {
  TestNode(int v) : value(v), next(nullptr) {}
  int value;
  TestNode *next;
};

TEST_CASE("Bag single thread", "[concurrent]") {
  IntrusiveBag<TestNode> b([] { return false; });
  REQUIRE(b.remove() == nullptr);
  for (int i = 0; i < 100; i++) {
    b.add(new TestNode(i));
  }
  for (int i = 0; i < 100; i++) {
    auto n = b.remove();
    REQUIRE(n->value == 99 - i);
    delete n;
  }
  REQUIRE(b.remove() == nullptr);
}

static emscripten_wasm_worker_t fakeRealtimeThread;

struct WorkerPayload {
  uint32_t sema{};
  IntrusiveBag<TestNode> bag{
      [] { return emscripten_wasm_worker_self_id() == fakeRealtimeThread; }};
};

static constexpr const int NUM_THREADS = 8;
static constexpr const int NUM_ELEMS = 100;
static constexpr const int NUM_ITERS = 10000;

static void bag_multi_worker_start(int payloadPtr) {
  auto payload = reinterpret_cast<WorkerPayload *>(payloadPtr);

  // do test
  std::mt19937 gen(emscripten_wasm_worker_self_id());
  std::uniform_int_distribution<> dis(0, NUM_ELEMS / NUM_THREADS);
  std::uniform_real_distribution<> realDis(0, 1);
  for (int i = 0; i < NUM_ITERS; i++) {
    std::vector<TestNode *> mine;
    if (realDis(gen) < 0.005) {
      // removeAll
      mine.reserve(NUM_ELEMS);
      auto root = payload->bag.removeAll();
      while (root) {
        mine.push_back(root);
        root = root->next;
      }
    } else {
      // remove
      const auto removeCount = dis(gen);
      mine.reserve(removeCount);
      for (int j = 0; j < removeCount; j++) {
        auto n = payload->bag.remove();
        if (n != nullptr) {
          mine.push_back(n);
        }
      }
    }
    for (auto n : mine) {
      payload->bag.add(n);
    }
  }

  emscripten_atomic_add_u32(&payload->sema, 1);
  emscripten_atomic_notify(&payload->sema, EMSCRIPTEN_NOTIFY_ALL_WAITERS);
}

TEST_CASE("Bag multi thread", "[.][concurrent][noasan]") {
#if defined(__has_feature)
#if __has_feature(address_sanitizer)
  assert(false);
#endif
#endif
  WorkerPayload w;

  for (int i = 0; i < NUM_ELEMS; i++) {
    w.bag.add(new TestNode(i));
  }

  std::vector<emscripten_wasm_worker_t> workers;
  workers.reserve(NUM_THREADS);
  for (int i = 0; i < NUM_THREADS; ++i) {
    workers.push_back(emscripten_malloc_wasm_worker(3670016));
  }
  fakeRealtimeThread = workers[0];

  // test with concurrency in 1..NUM_THREADS
  for (int i = 1; i < NUM_THREADS; i++) {
    w.sema = 0;

    for (int j = 0; j < i; j++) {
      emscripten_wasm_worker_post_function_vi(
          workers[j], bag_multi_worker_start, reinterpret_cast<int>(&w));
    }

    for (;;) {
      auto v = emscripten_atomic_load_u32(&w.sema);
      if (v == i) {
        break;
      }
      emscripten_atomic_wait_u32(&w.sema, v, ATOMICS_WAIT_DURATION_INFINITE);
    }

    std::vector<TestNode *> res;
    for (;;) {
      auto n = w.bag.remove();
      if (n == nullptr) {
        break;
      }
      res.push_back(n);
    }

    std::sort(res.begin(), res.end(), [](const TestNode *a, const TestNode *b) {
      return a->value < b->value;
    });
    REQUIRE(res.size() == NUM_ELEMS);
    for (int i = 0; i < res.size(); i++) {
      REQUIRE(i == res[i]->value);
    }

    for (auto &n : res) {
      w.bag.add(n);
    }
  }

  for (auto &worker : workers) {
    emscripten_terminate_wasm_worker(worker);
  }

  for (;;) {
    auto n = w.bag.remove();
    if (n == nullptr) {
      break;
    }
    delete n;
  }
}