#include <ros/common.h>
#include <futex.h>
#include <sys/queue.h>
#include <pthread.h>
#include <parlib.h>
#include <assert.h>
#include <stdio.h>
#include <errno.h>
#include <slab.h>
#include <mcs.h>

static inline int futex_wake(int *uaddr, int count);
static inline int futex_wait(int *uaddr, int val, uint64_t ms_timeout);
static void *timer_thread(void *arg);

struct futex_element {
  TAILQ_ENTRY(futex_element) link;
  pthread_t pthread;
  int *uaddr;
  uint64_t ms_timeout;
  bool timedout;
};
TAILQ_HEAD(futex_queue, futex_element);

struct futex_data {
  struct mcs_pdr_lock lock;
  struct futex_queue queue;
  int timer_enabled;
  pthread_t timer;
  long time;
};
static struct futex_data __futex;

static inline void futex_init()
{
  mcs_pdr_init(&__futex.lock);
  TAILQ_INIT(&__futex.queue);
  __futex.timer_enabled = false;
  pthread_create(&__futex.timer, NULL, timer_thread, NULL);
  __futex.time = 0;
}

static void __futex_block(struct uthread *uthread, void *arg) {
  pthread_t pthread = (pthread_t)uthread;
  struct futex_element *e = (struct futex_element*)arg;

  // Set the remaining properties of the futex element
  e->pthread = pthread;
  e->timedout = false;
  // Adjust the timeout properties on the futex element
  bool enable_timer = false;
  if(e->ms_timeout != (uint64_t)-1) {
    e->ms_timeout += __futex.time;
	// If we are setting the timeout, get ready to
	// enable the timer if it is currently disabled.
    if(__futex.timer_enabled == false) {
      __futex.timer_enabled = true;
      enable_timer = true;
    }
  }
  // Notify the scheduler of the type of yield we did
  __pthread_generic_yield(pthread);
  pthread->state = PTH_BLK_MUTEX;
  // Insert the futex element into the queue
  TAILQ_INSERT_TAIL(&__futex.queue, e, link);
  // Unlock the pdr_lock 
  mcs_pdr_unlock(&__futex.lock);
  // Enable the timer if we need to outside the lock
  if(enable_timer)
    futex_wake(&__futex.timer_enabled, 1);
}

static inline int futex_wait(int *uaddr, int val, uint64_t ms_timeout)
{
  // Atomically do the following...
  mcs_pdr_lock(&__futex.lock);
  // If the value of *uaddr matches val
  if(*uaddr == val) {
    // Create a new futex element and initialize it.
    struct futex_element e;
    e.uaddr = uaddr;
    e.ms_timeout = ms_timeout;
    // Yield the uthread...
    // We set the remaining properties of the futex element, adjust the
    // timeout, set the timer, and unlock the pdr lock on the other side.
    // It is important that we do the unlock on the other side, because (unlike
    // linux, etc.) its possible to get interrupted and drop into vcore context
    // right after releasing the lock.  If that vcore code then calls
    // futex_wake(), we would be screwed.  Doing things this way means we have
    // to hold the lock longer, but its necessary for correctness.
    uthread_yield(TRUE, __futex_block, &e);
    // We are unlocked here!

    // After waking, if we timed out, set the error
    // code appropriately and return
    if(e.timedout) {
      errno = ETIMEDOUT;
      return -1;
    }
  } else {
      mcs_pdr_unlock(&__futex.lock);
  }
  return 0;
}

static inline int futex_wake(int *uaddr, int count)
{
  struct futex_element *e,*n = NULL;
  struct futex_queue q = TAILQ_HEAD_INITIALIZER(q);

  // Atomically grab all relevant futex blockers
  // from the global futex queue
  mcs_pdr_lock(&__futex.lock);
  e = TAILQ_FIRST(&__futex.queue);
  while(e != NULL) {
    if(count > 0) {
      n = TAILQ_NEXT(e, link);
      if(e->uaddr == uaddr) {
        TAILQ_REMOVE(&__futex.queue, e, link);
        TAILQ_INSERT_TAIL(&q, e, link);
        count--;
      }
      e = n;
    }
    else break;
  }
  mcs_pdr_unlock(&__futex.lock);

  // Unblock them outside the lock
  e = TAILQ_FIRST(&q);
  while(e != NULL) {
    n = TAILQ_NEXT(e, link);
    TAILQ_REMOVE(&q, e, link);
    uthread_runnable((struct uthread*)e->pthread);
    e = n;
  }
  return 0;
}

static void *timer_thread(void *arg)
{
  struct futex_element *e,*n = NULL;
  struct futex_queue q = TAILQ_HEAD_INITIALIZER(q);

  // Do this forever...
  for(;;) {
    // Block for 1 millisecond
    sys_block(1000);

    // Then atomically do the following...
    mcs_pdr_lock(&__futex.lock);
    // Up the time
    __futex.time++;

	// Find all futexes that have timed out on this iteration,
	// and count those still waiting
    int waiting = 0;
    e = TAILQ_FIRST(&__futex.queue);
    while(e != NULL) {
      n = TAILQ_NEXT(e, link);
      if(e->ms_timeout == __futex.time) {
        e->timedout = true;
        TAILQ_REMOVE(&__futex.queue, e, link);
        TAILQ_INSERT_TAIL(&q, e, link);
      }
      else if(e->ms_timeout != (uint64_t)-1)
        waiting++;
      e = n;
    }
    // If there are no more waiting, disable the timer
    if(waiting == 0) {
      __futex.time = 0;
      __futex.timer_enabled = false;
    }
    mcs_pdr_unlock(&__futex.lock);

    // Unblock any futexes that have timed out outside the lock
    e = TAILQ_FIRST(&q);
    while(e != NULL) {
      n = TAILQ_NEXT(e, link);
      TAILQ_REMOVE(&q, e, link);
      uthread_runnable((struct uthread*)e->pthread);
      e = n;
    }

    // If we have disabled the timer, park this thread
    futex_wait(&__futex.timer_enabled, false, -1);
  }
}

int futex(int *uaddr, int op, int val,
          const struct timespec *timeout,
          int *uaddr2, int val3)
{
  uint64_t ms_timeout = (uint64_t)-1;
  assert(uaddr2 == NULL);
  assert(val3 == 0);
  if(timeout != NULL) {
    ms_timeout = timeout->tv_sec*1000 + timeout->tv_nsec/1000000L;
    assert(ms_timeout > 0);
  }

  run_once(futex_init());
  switch(op) {
    case FUTEX_WAIT:
      return futex_wait(uaddr, val, ms_timeout);
    case FUTEX_WAKE:
      return futex_wake(uaddr, val);
    default:
      errno = ENOSYS;
      return -1;
  }
  return -1;
}