#pragma once
#include "bag.h"
#include "machine.h"
#include <cassert>
#include <emscripten/atomic.h>
#include <emscripten/wasm_worker.h>
#include <functional>
#include <map>

// User must ensure that an instance of DeadlineWorkQ outlives all Future
// instances it returns
class DeadlineWorkQ {
  struct Node {
    Node *next{};
    size_t deadline{};
    std::function<void()> work;
    bool closed{};
    uint32_t finished;
    uint32_t refcount;
  };

  friend class IntrusiveBag<Node>;
  IntrusiveBag<Node> freeNodes{[] { return false; }};
  IntrusiveBag<Node> submitted{[] { return false; }};

  emscripten_lock_t pqLock{};
  std::multimap<size_t, Node *> pq;
  alignas(CACHE_LINE_SIZE) uint32_t closed{};

  Node *waitForNode() {
    emscripten_lock_waitinf_acquire(&pqLock);
    for (;;) {
      // add all submitted unprioritized work to pq
      auto root = submitted.removeAll();
      while (root != nullptr) {
        auto next = root->next;
        root->next = nullptr;
        pq.insert(std::make_pair(root->deadline, root));
        root = next;
      }

      // get the earliest-deadline work
      if (!pq.empty()) {
        auto it = pq.begin();
        auto node = it->second;
        pq.erase(it);
        // if this is a close node, put it back and return nullptr
        // to exit the calling worker - this will cause all workers to exit
        if (node->closed) {
          submitted.add(node);
          submitted.notifyHead();
          emscripten_lock_release(&pqLock);
          return nullptr;
        }
        emscripten_lock_release(&pqLock);
        return node;
      }

      submitted.waitNonempty();
    }
  }

  Node *getFreeNode() {
    auto n = freeNodes.remove();
    if (n == nullptr) {
      n = new Node{};
    }
    return n;
  }

public:
  class Future {
    DeadlineWorkQ *q;
    Node *node;
    Future(DeadlineWorkQ *q, Node *n) : q{q}, node{n} {}
    friend class DeadlineWorkQ;

  public:
    Future() : q{nullptr}, node{nullptr} {}
    Future(const Future &) = delete;
    Future &operator=(const Future &) = delete;
    Future(Future &&other) noexcept : q{other.q}, node{other.node} {
      other.node = nullptr;
    }
    Future &operator=(Future &&other) noexcept {
      if (this != &other) {
        this->q = other.q;
        this->node = other.node;
        other.node = nullptr;
      }
      return *this;
    }
    ~Future() {
      if (node != nullptr) {
        auto prevval = emscripten_atomic_sub_u32(&node->refcount, 1);
        if (prevval == 1) {
          q->freeNodes.add(node);
        }
      }
    }

    bool isFinished() const {
      if (node == nullptr) {
        return true;
      }
      return emscripten_atomic_load_u32(&node->finished) != 0;
    }

    void wait() const {
      if (node == nullptr) {
        return;
      }
      emscripten_atomic_wait_u32(&node->finished, 0,
                                 ATOMICS_WAIT_DURATION_INFINITE);
    }
  };

  DeadlineWorkQ(const DeadlineWorkQ &) = delete;
  DeadlineWorkQ &operator=(const DeadlineWorkQ &) = delete;
  DeadlineWorkQ(DeadlineWorkQ &&) = delete;
  DeadlineWorkQ &operator=(DeadlineWorkQ &&) = delete;

  DeadlineWorkQ(size_t preallocateNodes = 100) {
    emscripten_lock_init(&pqLock);
    for (size_t i = 0; i < preallocateNodes; i++) {
      freeNodes.add(new Node{});
    }
  }
  ~DeadlineWorkQ() {
    assert(emscripten_atomic_load_u32(&closed) == 1);
    auto root = freeNodes.removeAll();
    while (root != nullptr) {
      auto next = root->next;
      delete root;
      root = next;
    }
    root = submitted.removeAll();
    while (root != nullptr) {
      auto next = root->next;
      delete root;
      root = next;
    }
  }

  Future submit(std::function<void()> &&f, size_t deadline) {
    assert(emscripten_atomic_load_u32(&closed) == 0);
    auto n = getFreeNode();
    n->work = std::move(f);
    n->deadline = deadline;
    n->closed = false;
    emscripten_atomic_store_u32(&n->finished, 0);
    emscripten_atomic_store_u32(&n->refcount, 2);
    submitted.add(n);
    submitted.notifyHead();
    return Future(this, n);
  }

  bool doWork() {
    auto n = waitForNode();
    if (n == nullptr) {
      return false;
    }
    n->work();
    if (emscripten_atomic_sub_u32(&n->refcount, 1) == 1) {
      freeNodes.add(n);
    } else {
      emscripten_atomic_store_u32(&n->finished, 1);
      emscripten_atomic_notify(&n->finished, EMSCRIPTEN_NOTIFY_ALL_WAITERS);
    }
    return true;
  }

  void close() {
    if (emscripten_atomic_load_u32(&closed) == 1) {
      return;
    }
    emscripten_atomic_store_u32(&closed, 1);
    auto n = getFreeNode();
    n->deadline = std::numeric_limits<typeof(n->deadline)>::max();
    n->closed = true;
    submitted.add(n);
    submitted.notifyHead();
  }
};

class DeadlineThreadPool {
  struct WorkerPayload {
    uint32_t sema{};
    DeadlineWorkQ &q;
    WorkerPayload(DeadlineWorkQ &q) : q{q} {}
  };

  WorkerPayload payload;
  static void workq_worker_start(int payloadPtr);
  std::vector<emscripten_wasm_worker_t> workers;

public:
  DeadlineThreadPool(const DeadlineThreadPool &other) = delete;
  DeadlineThreadPool &operator=(const DeadlineThreadPool &other) = delete;
  DeadlineThreadPool(DeadlineThreadPool &&other) = delete;
  DeadlineThreadPool &operator=(DeadlineThreadPool &&other) = delete;

  DeadlineThreadPool(int numThreads, DeadlineWorkQ &q);
  ~DeadlineThreadPool();
};