#include "minilua.h"
#include <stdio.h>

lua_State* uNewThread(lua_State* M) {
  lua_State* L = lua_newthread(M);
  printf("NEW THREAD %p\n", L);
  return L;
};

// don't forget to pop values before a resume
int uResumeThread(lua_State* L, lua_State* M, int n) {
  int nres = 0;
  if (n == -1) {
    n = lua_gettop(L);
  }

  int err = lua_resume(L, M, n, &nres);
  if (err != LUA_YIELD) {
    printf("THREAD END : %s\n", lua_gettop(L) ? lua_tostring(L, -1) : "(NIL)");
    return err == LUA_OK ? 0 : -1;
  }

  return 1; // do some stuff with yielded values
};

#define THREADPOOL_MAX 10
struct lua_threadpool {
  lua_State* M;
  lua_State* arr[THREADPOOL_MAX];
  size_t     size; 
};

void pool_add(struct lua_threadpool* pool, lua_State* L) {
  if (L == pool->M || L == NULL) return;
  if (pool->size >= THREADPOOL_MAX) return;
  pool->arr[pool->size++] = L;
}

int pool_add_code(struct lua_threadpool* pool, const char* source) {
  lua_State* L = uNewThread(pool->M);
  if (luaL_loadstring(L, source) != LUA_OK) {
    printf("SOURCE LOAD ERROR : %s\n", lua_gettop(L) ? lua_tostring(L, -1) : "(NIL)");
    return 1;
  }
  
  pool_add(pool, L);
  return 0;
}

lua_State* pool_pop(struct lua_threadpool* pool) {
  if (!pool->size) return NULL;
  return pool->arr[--pool->size];
}

int main() {
  lua_State* M = luaL_newstate();
  luaL_openlibs(M);

  printf("Heloo!\n");
  
  struct lua_threadpool pool = {M, 0};
  
  pool_add_code(&pool, "print('heloo'); print(coroutine.yield()); print('end');");
  pool_add_code(&pool, "print('heloo'); i = 10; while i > 0 do print(coroutine.yield(10, 20)); i = i - 1; end;");
  pool_add_code(&pool, "local i; print('test'); i = 6; while i > 0 do print(coroutine.yield(10, 20)); i = i - 1; end; \
coroutine.yield('aha'); print(debug.traceback()); coroutine.yield('ohno');");

  int numiter = 0;
  while (pool.size > 0) {
    //printf("LOOP AGAIN!\n");
    for (size_t i = 0; i < pool.size; i++) {
      lua_State* L = pool.arr[i];
      lua_pushnumber(L, numiter);
      int status = uResumeThread(L, pool.M, 1);
      if (status < 1) { // remove thread, it's done
        pool.arr[i] = pool_pop(&pool);
        lua_resetthread(L);
        continue;
      }
      lua_settop(L, 0); // remove values to be resumed.
    }
    numiter++;
  };
  printf("It's done! %i iterations!\n", numiter);
  lua_close(M);
  return 0;
}