#include <Windows.h>
#include <assert.h>
#include <stdio.h>

// Since we're not really switching contexts, we only actually need to save
// the call saved registers!
struct SavedContextFrame
{
  void *r15;
  void *r14;
  void *r13;
  void *r12;
  void *rdi;
  void *rsi;
  void *rbp;
  void *rbx;
  void (*rip)();
};

// This has a few extra entries at the bottom of
// the stack to handle a thread returning
struct NewFrame
{
  SavedContextFrame context;

  void *threadexit;
  void *zero;
};

struct Task
{
  // rsp at beginning of structure for convenience from asm
  SavedContextFrame *rsp;

  enum State
  {
    INVALID,
    ROOT,
    CREATED,
    EXITED,
  };
  State state;

  void *stack;
  size_t stacksize;

  Task *next, *prev;

  Task()
    : rsp(0)
    , state(INVALID)
    , stack(0)
  {
    next = prev = this;
  }
};

// Circular lists to track runnable and waiting tasks
static Task root_runnable;
static Task root_waiting;
static Task *task = &root_runnable;
extern "C" Task *current_task;

Task *current_task;

extern "C" void TaskBootup(void *rsp);
extern "C" void SwitchToTask(void *rsp, void **old_task_rsp);
extern "C" void SwitchToNextTask();
extern "C" void StartNewTask();

extern "C" SavedContextFrame *ScheduleTask(SavedContextFrame *rsp)
{
  // Save outgoing task's frame
  current_task->rsp = rsp;

  // See if outgoing task needs self destruct
  if (current_task->state == Task::EXITED)
  {
    // Capture node being deleted
    auto *tmp = current_task;

    // Advance to next
    current_task = current_task->next;

    // Unlink
    tmp->next->prev = tmp->prev;
    tmp->prev->next = tmp->next;

    // Wipe dangling pointers
    tmp->next = 0;
    tmp->prev = 0;

    // Free memory
    delete tmp;
  }

  for ( ; root_runnable.next != root_runnable.prev; )
  {
    // Go to next task
    current_task = current_task->next;

    // Skip over first node
    if (current_task == &root_runnable)
      continue;

    return current_task->rsp;
  }

  // No more tasks to run, switch back to root task
  current_task = &root_runnable;

  assert(current_task->rsp != 0);

  return current_task->rsp;
}

static void task_exit(void *a1, void *a2, void *a3, void *a4)
{
  current_task->state = Task::EXITED;
  SwitchToNextTask();
}

void test_thread(void *a1, void *a2, void *a3, void *a4)
{
  for (auto i = 0; i < 4; ++i)
  {
    printf("0x%p 0x%p 0x%p 0x%p\n", a1, a2, a3, a4);
    SwitchToNextTask();
  }
  printf("0x%p 0x%p 0x%p 0x%p - done\n", a1, a2, a3, a4);
}

void CreateTask(void (*f)(void*,void*,void*,void*),
  void *a0, void *a1, void *a2, void *a3)
{
  Task *task = new Task;

  // Allocate a new stack and put the context at the end of it
  task->stacksize = 1<<20;
  task->stack = malloc(task->stacksize);

  auto frame = (NewFrame*)((char*)task->stack + task->stacksize) - 1;
  memset(frame, 0, sizeof(*frame));

  frame->context.rip = StartNewTask;
  frame->threadexit = task_exit;
  frame->zero = 0;

  // When we're creating a new task, we own all the registers
  // in the frame, so we put all the information into those
  // and make it start at a task creation thunk

  frame->context.rbx = task;
  frame->context.rbp = f;
  frame->context.r12 = a0;
  frame->context.r13 = a1;
  frame->context.r14 = a2;
  frame->context.r15 = a3;

  task->rsp = &frame->context;

  // Link into start of circular list
  task->prev = root_runnable.prev;
  task->next = &root_runnable;
  task->prev->next = task;
  task->next->prev = task;

  task->state = Task::CREATED;

  if (current_task)
    return;

  // We only reach here first time
  current_task = &root_runnable;
  current_task->state = Task::ROOT;

  current_task = &root_runnable;
}

class FiberGroup
{
};

int main()
{
  void *tasks[16];

  for (auto i = (char *)0; i < (char *)8; ++i)
    tasks[i] = CreateFiber(0, test_fiber, 0);

  ConvertThreadToFiber(0);
  SwitchToFiber(tasks[0]);

  for (auto i = (char *)0; i < (char *)8; ++i)
    CreateTask(test_thread, i + 1, i + 2, i + 3, i + 4);

  while (root_runnable.next != root_runnable.prev)
    SwitchToNextTask();

  printf("Root done\n");

  return 0;
}