/*
 * This module exports three classes, and each instance of those classes has its
 * own private registry for temporary reference storage(keeping state between
 * calls). A private registry makes managing memory much easier since all we
 * have to do is call luaL_unref passing the registry reference when the
 * instance is collected by the __gc metamethod.
 *
 * This private registry is manipulated with `lmpack_ref` / `lmpack_unref` /
 * `lmpack_geti`, which are analogous to `luaL_ref` / `luaL_unref` /
 * `lua_rawgeti` but operate on the private registry passed as argument.
 *
 * In order to simplify debug registry leaks during normal operation(with the
 * leak_test.lua script), these `lmpack_*` registry functions will target the
 * normal lua registry when MPACK_DEBUG_REGISTRY_LEAK is defined during
 * compilation.
 */
#define LUA_LIB
/* for snprintf */
#ifdef __APPLE__
// Apple uses C99 source, not XOPEN
#define _C99_SOURCE 1
#else
#define _XOPEN_SOURCE 500
#endif
#include <limits.h>
#include <stdlib.h>
#include <string.h>
#include <stdio.h>

#include <lauxlib.h>
#include <lua.h>
#include <luaconf.h>

#ifdef MPACK_USE_SYSTEM
# include <mpack.h>
#else
# define MPACK_API static
# include "mpack-src/src/mpack.c"
#endif

#define UNPACKER_META_NAME "mpack.Unpacker"
#define PACKER_META_NAME "mpack.Packer"
#define SESSION_META_NAME "mpack.Session"
#define NIL_NAME "mpack.NIL"

/* 
 * TODO(tarruda): When targeting lua 5.3 and being compiled with `long long`
 * support(not -ansi), we should make use of lua 64 bit integers for
 * representing msgpack integers, since `double` can't represent the full range.
 */

#ifndef luaL_reg
/* Taken from Lua5.1's lauxlib.h */
#define luaL_reg    luaL_Reg
#endif

#if LUA_VERSION_NUM > 501
#ifndef luaL_register
#define luaL_register(L,n,f) luaL_setfuncs(L,f,0)
#endif
#endif

typedef struct {
  lua_State *L;
  mpack_parser_t *parser;
  int reg, ext, unpacking;
  char *string_buffer;
} Unpacker;

typedef struct {
  lua_State *L;
  mpack_parser_t *parser;
  int reg, ext, root, packing;
  int is_bin, is_bin_fn;
} Packer;

typedef struct {
  lua_State *L;
  int reg;
  mpack_rpc_session_t *session;
  struct {
    int type;
    mpack_rpc_message_t msg;
    int method_or_error;
    int args_or_result;
  } unpacked;
  int unpacker;
} Session;

static int lmpack_ref(lua_State *L, int reg)
{
#ifdef MPACK_DEBUG_REGISTRY_LEAK
  return luaL_ref(L, LUA_REGISTRYINDEX);
#else
  int rv;
  lua_rawgeti(L, LUA_REGISTRYINDEX, reg);
  lua_pushvalue(L, -2);
  rv = luaL_ref(L, -2);
  lua_pop(L, 2);
  return rv;
#endif
}

static void lmpack_unref(lua_State *L, int reg, int ref)
{
#ifdef MPACK_DEBUG_REGISTRY_LEAK
  luaL_unref(L, LUA_REGISTRYINDEX, ref);
#else
  lua_rawgeti(L, LUA_REGISTRYINDEX, reg);
  luaL_unref(L, -1, ref);
  lua_pop(L, 1);
#endif
}

static void lmpack_geti(lua_State *L, int reg, int ref)
{
#ifdef MPACK_DEBUG_REGISTRY_LEAK
  lua_rawgeti(L, LUA_REGISTRYINDEX, ref);
#else
  lua_rawgeti(L, LUA_REGISTRYINDEX, reg);
  lua_rawgeti(L, -1, ref);
  lua_replace(L, -2);
#endif
}

/* make a shallow copy of the table on stack and remove it after the copy is
 * done */
static void lmpack_shallow_copy(lua_State *L)
{
  lua_newtable(L);
  lua_pushnil(L);
  while (lua_next(L, -3)) {
    lua_pushvalue(L, -2);
    lua_insert(L, -2);
    lua_settable(L, -4);
  }
  lua_remove(L, -2);
}

static mpack_parser_t *lmpack_grow_parser(mpack_parser_t *parser)
{
  mpack_parser_t *old = parser;
  mpack_uint32_t new_capacity = old->capacity * 2;
  parser = malloc(MPACK_PARSER_STRUCT_SIZE(new_capacity));
  if (!parser) goto end;
  mpack_parser_init(parser, new_capacity);
  mpack_parser_copy(parser, old);
  free(old);
end:
  return parser;
}

static mpack_rpc_session_t *lmpack_grow_session(mpack_rpc_session_t *session)
{
  mpack_rpc_session_t *old = session;
  mpack_uint32_t new_capacity = old->capacity * 2;
  session = malloc(MPACK_RPC_SESSION_STRUCT_SIZE(new_capacity));
  if (!session) goto end;
  mpack_rpc_session_init(session, new_capacity);
  mpack_rpc_session_copy(session, old);
  free(old);
end:
  return session;
}

static Unpacker *lmpack_check_unpacker(lua_State *L, int index)
{
  return luaL_checkudata(L, index, UNPACKER_META_NAME);
}

static Packer *lmpack_check_packer(lua_State *L, int index)
{
  return luaL_checkudata(L, index, PACKER_META_NAME);
}

static Session *lmpack_check_session(lua_State *L, int index)
{
  return luaL_checkudata(L, index, SESSION_META_NAME);
}

static int lmpack_isnil(lua_State *L, int index)
{
  int rv;
  if (!lua_isuserdata(L, index)) return 0;
  lua_getfield(L, LUA_REGISTRYINDEX, NIL_NAME);
  rv = lua_rawequal(L, -1, -2);
  lua_pop(L, 1);
  return rv;
}

static int lmpack_isunpacker(lua_State *L, int index)
{
  int rv;
  if (!lua_isuserdata(L, index) || !lua_getmetatable(L, index)) return 0;
  luaL_getmetatable(L, UNPACKER_META_NAME);
  rv = lua_rawequal(L, -1, -2);
  lua_pop(L, 2);
  return rv;
}

static void lmpack_pushnil(lua_State *L)
{
  lua_getfield(L, LUA_REGISTRYINDEX, NIL_NAME);
}

/* adapted from
 * https://github.com/antirez/lua-cmsgpack/blob/master/lua_cmsgpack.c */
static mpack_uint32_t lmpack_objlen(lua_State *L, int *is_array)
{
  size_t len, max;
  int isarr, type;
  lua_Number n;
#ifndef NDEBUG
  int top = lua_gettop(L);
  assert(top);
#endif

  if ((type = lua_type(L, -1)) != LUA_TTABLE) {
#if LUA_VERSION_NUM >= 502
    len = lua_rawlen(L, -1);
#elif LUA_VERSION_NUM == 501
    len = lua_objlen(L, -1);
#else
    #error You have either broken or too old Lua installation. This library requires Lua>=5.1
#endif
    goto end;
  }

  /* count the number of keys and determine if it is an array */
  len = 0;
  max = 0;
  isarr = 1;
  lua_pushnil(L);

  while (lua_next(L, -2)) {
    lua_pop(L, 1);  /* pop value */
    isarr = isarr
      && lua_isnumber(L, -1)            /* lua number */
      && (n = lua_tonumber(L, -1)) > 0  /* greater than 0 */
      && (size_t)n == n;                /* and integer */
    max = isarr && (size_t)n > max ? (size_t)n : max;
    len++;
  }

  *is_array = isarr && max == len;

end:
  if ((size_t)-1 > (mpack_uint32_t)-1 && len > (mpack_uint32_t)-1)
    /* msgpack spec doesn't allow lengths > 32 bits */
    len = (mpack_uint32_t)-1;
  assert(top == lua_gettop(L));
  return (mpack_uint32_t)len;
}

static int lmpack_unpacker_new(lua_State *L)
{
  Unpacker *rv;

  if (lua_gettop(L) > 1)
    return luaL_error(L, "expecting at most 1 table argument"); 

  rv = lua_newuserdata(L, sizeof(*rv));
  rv->parser = malloc(sizeof(*rv->parser));
  if (!rv->parser) return luaL_error(L, "Failed to allocate memory");
  mpack_parser_init(rv->parser, 0);
  rv->parser->data.p = rv;
  rv->string_buffer = NULL;
  rv->L = L;
  rv->unpacking = 0;
  luaL_getmetatable(L, UNPACKER_META_NAME);
  lua_setmetatable(L, -2);

#ifndef MPACK_DEBUG_REGISTRY_LEAK
  lua_newtable(L);
  rv->reg = luaL_ref(L, LUA_REGISTRYINDEX);
#endif
  rv->ext = LUA_NOREF;

  if (lua_istable(L, 1)) {
    /* parse options */
    lua_getfield(L, 1, "ext");
    if (!lua_isnil(L, -1)) {
      if (!lua_istable(L, -1))
        return luaL_error(L, "\"ext\" option must be a table"); 
      lmpack_shallow_copy(L);
    }
    rv->ext = lmpack_ref(L, rv->reg);
  }

  return 1;
}

static int lmpack_unpacker_delete(lua_State *L)
{
  Unpacker *unpacker = lmpack_check_unpacker(L, 1);
  if (unpacker->ext != LUA_NOREF)
    lmpack_unref(L, unpacker->reg, unpacker->ext);
#ifndef MPACK_DEBUG_REGISTRY_LEAK
  luaL_unref(L, LUA_REGISTRYINDEX, unpacker->reg);
#endif
  free(unpacker->parser);
  return 0;
}

static void lmpack_parse_enter(mpack_parser_t *parser, mpack_node_t *node)
{
  Unpacker *unpacker = parser->data.p;
  lua_State *L = unpacker->L;

  switch (node->tok.type) {
    case MPACK_TOKEN_NIL:
      lmpack_pushnil(L); break;
    case MPACK_TOKEN_BOOLEAN:
      lua_pushboolean(L, (int)mpack_unpack_boolean(node->tok)); break;
    case MPACK_TOKEN_UINT:
    case MPACK_TOKEN_SINT:
    case MPACK_TOKEN_FLOAT:
      lua_pushnumber(L, mpack_unpack_number(node->tok)); break;
    case MPACK_TOKEN_CHUNK:
      assert(unpacker->string_buffer);
      memcpy(unpacker->string_buffer + MPACK_PARENT_NODE(node)->pos,
          node->tok.data.chunk_ptr, node->tok.length);
      break;
    case MPACK_TOKEN_BIN:
    case MPACK_TOKEN_STR:
    case MPACK_TOKEN_EXT:
      unpacker->string_buffer = malloc(node->tok.length);
      if (!unpacker->string_buffer) luaL_error(L, "Failed to allocate memory");
      break;
    case MPACK_TOKEN_ARRAY:
    case MPACK_TOKEN_MAP:
      lua_newtable(L);
      node->data[0].i = lmpack_ref(L, unpacker->reg);
      break;
  }
}

static void lmpack_parse_exit(mpack_parser_t *parser, mpack_node_t *node)
{
  Unpacker *unpacker = parser->data.p;
  lua_State *L = unpacker->L;
  mpack_node_t *parent = MPACK_PARENT_NODE(node);

  switch (node->tok.type) {
    case MPACK_TOKEN_BIN:
    case MPACK_TOKEN_STR:
    case MPACK_TOKEN_EXT:
      lua_pushlstring(L, unpacker->string_buffer, node->tok.length);
      free(unpacker->string_buffer);
      unpacker->string_buffer = NULL;
      if (node->tok.type == MPACK_TOKEN_EXT && unpacker->ext != LUA_NOREF) {
        /* check if there's a handler for this type */
        lmpack_geti(L, unpacker->reg, unpacker->ext);
        lua_rawgeti(L, -1, node->tok.data.ext_type);
        if (lua_isfunction(L, -1)) {
          /* stack:
           *
           * -1: ext unpacker function
           * -2: ext unpackers table 
           * -3: ext string 
           *
           * We want to call the ext unpacker function with the type and string
           * as arguments, so push those now
           */
          lua_pushinteger(L, node->tok.data.ext_type);
          lua_pushvalue(L, -4);
          lua_call(L, 2, 1);
          /* stack:
           *
           * -1: returned object
           * -2: ext unpackers table
           * -3: ext string 
           */
          lua_replace(L, -3);
        } else {
          /* the last lua_rawgeti should have pushed nil on the stack,
           * remove it */
          lua_pop(L, 1);
        }
        /* pop the ext unpackers table */
        lua_pop(L, 1);
      }
      break;
    case MPACK_TOKEN_ARRAY:
    case MPACK_TOKEN_MAP:
      lmpack_geti(L, unpacker->reg, (int)node->data[0].i);
      lmpack_unref(L, unpacker->reg, (int)node->data[0].i);
      break;
    default:
      break;
  }

  if (parent && parent->tok.type < MPACK_TOKEN_BIN) {
    /* At this point the parsed object is on the stack. Add it to the parent
     * container. First put the container on the stack. */
    lmpack_geti(L, unpacker->reg, (int)parent->data[0].i);

    if (parent->tok.type == MPACK_TOKEN_ARRAY) {
      /* Array, save the value on key equal to `parent->pos` */
      lua_pushnumber(L, (lua_Number)parent->pos);
      lua_pushvalue(L, -3);
      lua_settable(L, -3);
    } else {
      assert(parent->tok.type == MPACK_TOKEN_MAP);
      if (parent->key_visited) {
        /* save the key on the registry */ 
        lua_pushvalue(L, -2);
        parent->data[1].i = lmpack_ref(L, unpacker->reg);
      } else {
        /* set the key/value pair */
        lmpack_geti(L, unpacker->reg, (int)parent->data[1].i);
        lmpack_unref(L, unpacker->reg, (int)parent->data[1].i);
        lua_pushvalue(L, -3);
        lua_settable(L, -3);
      }
    }
    lua_pop(L, 2);  /* pop the container/object */
  }
}

static int lmpack_unpacker_unpack_str(lua_State *L, Unpacker *unpacker,
    const char **str, size_t *len)
{
  int rv;

  if (unpacker->unpacking) {
    return luaL_error(L, "Unpacker instance already working. Use another "
                         "Unpacker or the module's \"unpack\" function if you "
                         "need to unpack from the ext handler");
  }
  
  do {
    unpacker->unpacking = 1;
    rv = mpack_parse(unpacker->parser, str, len, lmpack_parse_enter,
        lmpack_parse_exit);
    unpacker->unpacking = 0;

    if (rv == MPACK_NOMEM) {
      unpacker->parser = lmpack_grow_parser(unpacker->parser);
      if (!unpacker->parser) {
        unpacker->unpacking = 0;
        return luaL_error(L, "failed to grow Unpacker capacity");
      }
    }
  } while (rv == MPACK_NOMEM);

  if (rv == MPACK_ERROR)
    return luaL_error(L, "invalid msgpack string");

  return rv;
}

static int lmpack_unpacker_unpack(lua_State *L)
{
  int result, argc;
  lua_Number startpos;
  size_t len, offset;
  const char *str, *str_init;
  Unpacker *unpacker;
  
  if ((argc = lua_gettop(L)) > 3 || argc < 2)
    return luaL_error(L, "expecting between 2 and 3 arguments"); 

  unpacker = lmpack_check_unpacker(L, 1);
  unpacker->L = L;

  str_init = str = luaL_checklstring(L, 2, &len);
  startpos = lua_gettop(L) == 3 ? luaL_checknumber(L, 3) : 1;

  luaL_argcheck(L, startpos > 0, 3,
      "start position must be greater than zero");
  luaL_argcheck(L, (size_t)startpos == startpos, 3,
      "start position must be an integer");
  luaL_argcheck(L, (size_t)startpos <= len, 3,
      "start position must be less than or equal to the input string length");

  offset = (size_t)startpos - 1 ;
  str += offset;
  len -= offset;
  result = lmpack_unpacker_unpack_str(L, unpacker, &str, &len);

  if (result == MPACK_EOF)
    /* if we hit EOF, return nil as the object */
    lua_pushnil(L);

  /* also return the new position in the input string */
  lua_pushinteger(L, str - str_init + 1);
  assert(lua_gettop(L) == argc + 2);
  return 2;
}

static int lmpack_packer_new(lua_State *L)
{
  Packer *rv;

  if (lua_gettop(L) > 1)
    return luaL_error(L, "expecting at most 1 table argument"); 

  rv = lua_newuserdata(L, sizeof(*rv));
  rv->parser = malloc(sizeof(*rv->parser));
  if (!rv->parser) return luaL_error(L, "failed to allocate parser memory");
  mpack_parser_init(rv->parser, 0);
  rv->parser->data.p = rv;
  rv->L = L;
  rv->packing = 0;
  rv->is_bin = 0;
  rv->is_bin_fn = LUA_NOREF;
  luaL_getmetatable(L, PACKER_META_NAME);
  lua_setmetatable(L, -2);

#ifndef MPACK_DEBUG_REGISTRY_LEAK
  lua_newtable(L);
  rv->reg = luaL_ref(L, LUA_REGISTRYINDEX);
#endif
  rv->ext = LUA_NOREF;

  if (lua_istable(L, 1)) {
    /* parse options */
    lua_getfield(L, 1, "ext");
    if (!lua_isnil(L, -1)) {
      if (!lua_istable(L, -1))
        return luaL_error(L, "\"ext\" option must be a table"); 
      lmpack_shallow_copy(L);
    }
    rv->ext = lmpack_ref(L, rv->reg);
    lua_getfield(L, 1, "is_bin");
    if (!lua_isnil(L, -1)) {
      if (!lua_isboolean(L, -1) && !lua_isfunction(L, -1))
        return luaL_error(L,
            "\"is_bin\" option must be a boolean or function"); 
      rv->is_bin = lua_toboolean(L, -1);
      if (lua_isfunction(L, -1)) rv->is_bin_fn = lmpack_ref(L, rv->reg);
      else lua_pop(L, 1);
    } else {
      lua_pop(L, 1);
    }

  }

  return 1;
}

static int lmpack_packer_delete(lua_State *L)
{
  Packer *packer = lmpack_check_packer(L, 1);
  if (packer->ext != LUA_NOREF)
    lmpack_unref(L, packer->reg, packer->ext);
#ifndef MPACK_DEBUG_REGISTRY_LEAK
  luaL_unref(L, LUA_REGISTRYINDEX, packer->reg);
#endif
  free(packer->parser);
  return 0;
}

static void lmpack_unparse_enter(mpack_parser_t *parser, mpack_node_t *node)
{
  int type;
  Packer *packer = parser->data.p;
  lua_State *L = packer->L;
  mpack_node_t *parent = MPACK_PARENT_NODE(node);

  if (parent) {
    /* get the parent */
    lmpack_geti(L, packer->reg, (int)parent->data[0].i);

    if (parent->tok.type > MPACK_TOKEN_MAP) {
      /* strings are a special case, they are packed as single child chunk
       * node */
      const char *str = lua_tolstring(L, -1, NULL);
      node->tok = mpack_pack_chunk(str, parent->tok.length);
      lua_pop(L, 1);
      return;
    }

    if (parent->tok.type == MPACK_TOKEN_ARRAY) {
      /* push the next index */
      lua_pushnumber(L, (lua_Number)(parent->pos + 1));
      /* push the element */
      lua_gettable(L, -2);
    } else if (parent->tok.type == MPACK_TOKEN_MAP) {
      int result;
      /* push the previous key */
      lmpack_geti(L, packer->reg, (int)parent->data[1].i);
      /* push the pair */
      result = lua_next(L, -2);
      assert(result);  /* should not be here if the map was fully processed */
      if (parent->key_visited) {
        /* release the current key */
        lmpack_unref(L, packer->reg, (int)parent->data[1].i);
        /* push key to the top */
        lua_pushvalue(L, -2);
        /* set the key for the next iteration, leaving value on top */
        parent->data[1].i = lmpack_ref(L, packer->reg);
        /* replace key by the value */
        lua_replace(L, -2);
      } else {
        /* pop value */
        lua_pop(L, 1);
      }
    }
    /* remove parent, leaving only the object which will be serialized */
    lua_remove(L, -2);
  } else {
    /* root object */
    lmpack_geti(L, packer->reg, packer->root);
  }

  type = lua_type(L, -1);

  switch (type) {
    case LUA_TBOOLEAN:
      node->tok = mpack_pack_boolean((unsigned)lua_toboolean(L, -1));
      break;
    case LUA_TNUMBER:
      node->tok = mpack_pack_number(lua_tonumber(L, -1));
      break;
    case LUA_TSTRING: {
      int is_bin = packer->is_bin;
      if (is_bin && packer->is_bin_fn != LUA_NOREF) {
        lmpack_geti(L, packer->reg, packer->is_bin_fn);
        lua_pushvalue(L, -2);
        lua_call(L, 1, 1);
        is_bin = lua_toboolean(L, -1);
        lua_pop(L, 1);
      }
      if (is_bin) node->tok = mpack_pack_bin(lmpack_objlen(L, NULL));
      else node->tok = mpack_pack_str(lmpack_objlen(L, NULL));
      break;
    }
    case LUA_TTABLE: {
      int is_array;
      mpack_uint32_t len;
      mpack_node_t *n;

      if (packer->ext != LUA_NOREF && lua_getmetatable(L, -1)) {
        /* check if there's a handler for this metatable */
        lmpack_geti(L, packer->reg, packer->ext);
        lua_pushvalue(L, -2);
        lua_gettable(L, -2);
        if (lua_isfunction(L, -1)) {
          lua_Number ext = -1;
          /* stack:
           *
           * -1: ext packer function
           * -2: ext packers table
           * -3: metatable
           * -4: original object
           *
           * We want to call the ext packer function with the original object as
           * argument, so push it on the top
           */
          lua_pushvalue(L, -4);
          /* handler should return type code and string */
          lua_call(L, 1, 2);
          if (!lua_isnumber(L, -2) || (ext = lua_tonumber(L, -2)) < 0
              || ext > 127 || (int)ext != ext)
            luaL_error(L,
                "the first result from ext packer must be an integer "
                "between 0 and 127");
          if (!lua_isstring(L, -1))
            luaL_error(L,
                "the second result from ext packer must be a string");
          node->tok = mpack_pack_ext((int)ext, lmpack_objlen(L, NULL));
          /* stack: 
           *
           * -1: ext string
           * -2: ext type
           * -3: ext packers table
           * -4: metatable
           * -5: original table 
           *
           * We want to leave only the returned ext string, so
           * replace -5 with the string and pop 3
           */
          lua_replace(L, -5);
          lua_pop(L, 3);
          break;  /* done */
        } else {
          /* stack: 
           *
           * -1: ext packers table
           * -2: metatable
           * -3: original table 
           *
           * We want to leave only the original table since it will be handled
           * below, so pop 2
           */
          lua_pop(L, 2);
        }
      }

      /* check for cycles */
      n = node;
      while ((n = MPACK_PARENT_NODE(n))) {
        lmpack_geti(L, packer->reg, (int)n->data[0].i);
        if (lua_rawequal(L, -1, -2)) {
          /* break out of cycles with NIL  */
          node->tok = mpack_pack_nil();
          lua_pop(L, 2);
          lmpack_pushnil(L);
          goto end;
        }
        lua_pop(L, 1);
      }

      len = lmpack_objlen(L, &is_array);
      if (is_array) {
        node->tok = mpack_pack_array(len);
      } else {
        node->tok = mpack_pack_map(len);
        /* save nil as the previous key to start iteration */
        node->data[1].i = LUA_REFNIL;
      }
      break;
    }
    case LUA_TUSERDATA:
      if (lmpack_isnil(L, -1)) {
        node->tok = mpack_pack_nil();
        break;
      }
    /* Fallthrough */
    default:
	  {
		/* #define FMT */
		char errmsg[50];
		snprintf(errmsg, 50, "can't serialize object of type %d", type);
		luaL_error(L, errmsg);
	  }
  }

end:
  node->data[0].i = lmpack_ref(L, packer->reg);
}

static void lmpack_unparse_exit(mpack_parser_t *parser, mpack_node_t *node)
{
  Packer *packer = parser->data.p;
  lua_State *L = packer->L;
  if (node->tok.type != MPACK_TOKEN_CHUNK) {
    /* release the object */
    lmpack_unref(L, packer->reg, (int)node->data[0].i);
    if (node->tok.type == MPACK_TOKEN_MAP)
      lmpack_unref(L, packer->reg, (int)node->data[1].i);
  }
}

static int lmpack_packer_pack(lua_State *L)
{
  char *b;
  size_t bl;
  int result, argc;
  Packer *packer;
  luaL_Buffer buffer;

  if ((argc = lua_gettop(L)) != 2)
    return luaL_error(L, "expecting exactly 2 arguments"); 

  packer = lmpack_check_packer(L, 1);
  packer->L = L;
  packer->root = lmpack_ref(L, packer->reg);
  luaL_buffinit(L, &buffer);
  b = luaL_prepbuffer(&buffer);
  bl = LUAL_BUFFERSIZE;

  if (packer->packing) {
    return luaL_error(L, "Packer instance already working. Use another Packer "
                         "or the module's \"pack\" function if you need to "
                         "pack from the ext handler");
  }

  do {
    size_t bl_init = bl;
    packer->packing = 1;
    result = mpack_unparse(packer->parser, &b, &bl, lmpack_unparse_enter,
        lmpack_unparse_exit);
    packer->packing = 0;

    if (result == MPACK_NOMEM) {
      packer->parser = lmpack_grow_parser(packer->parser);
      if (!packer->parser) {
        packer->packing = 0;
        return luaL_error(L, "Failed to grow Packer capacity");
      }
    }

    luaL_addsize(&buffer, bl_init - bl);

    if (!bl) {
      /* buffer empty, resize */
      b = luaL_prepbuffer(&buffer);
      bl = LUAL_BUFFERSIZE;
    }
  } while (result == MPACK_EOF || result == MPACK_NOMEM);

  lmpack_unref(L, packer->reg, packer->root);
  luaL_pushresult(&buffer);
  assert(lua_gettop(L) == argc);
  return 1;
}

static int lmpack_session_new(lua_State *L)
{
  Session *rv = lua_newuserdata(L, sizeof(*rv));
  rv->session = malloc(sizeof(*rv->session));
  if (!rv->session) return luaL_error(L, "Failed to allocate memory");
  mpack_rpc_session_init(rv->session, 0);
  rv->L = L;
  luaL_getmetatable(L, SESSION_META_NAME);
  lua_setmetatable(L, -2);
#ifndef MPACK_DEBUG_REGISTRY_LEAK
  lua_newtable(L);
  rv->reg = luaL_ref(L, LUA_REGISTRYINDEX);
#endif
  rv->unpacker = LUA_REFNIL;
  rv->unpacked.args_or_result = LUA_NOREF;
  rv->unpacked.method_or_error = LUA_NOREF;
  rv->unpacked.type = MPACK_EOF;

  if (lua_istable(L, 1)) {
    /* parse options */
    lua_getfield(L, 1, "unpack");
    if (!lmpack_isunpacker(L, -1)) {
      return luaL_error(L,
          "\"unpack\" option must be a " UNPACKER_META_NAME " instance"); 
    }
    rv->unpacker = lmpack_ref(L, rv->reg);
  }

  return 1;
}

static int lmpack_session_delete(lua_State *L)
{
  Session *session = lmpack_check_session(L, 1);
  lmpack_unref(L, session->reg, session->unpacker);
#ifndef MPACK_DEBUG_REGISTRY_LEAK
  luaL_unref(L, LUA_REGISTRYINDEX, session->reg);
#endif
  free(session->session);
  return 0;
}

static int lmpack_session_receive(lua_State *L)
{
  int argc, done, rcount = 3;
  lua_Number startpos;
  size_t len;
  const char *str, *str_init;
  Session *session;
  Unpacker *unpacker = NULL;

  if ((argc = lua_gettop(L)) > 3 || argc < 2)
    return luaL_error(L, "expecting between 2 and 3 arguments"); 

  session = lmpack_check_session(L, 1);
  str_init = str = luaL_checklstring(L, 2, &len);
  startpos = lua_gettop(L) == 3 ? luaL_checknumber(L, 3) : 1;

  luaL_argcheck(L, startpos > 0, 3,
      "start position must be greater than zero");
  luaL_argcheck(L, (size_t)startpos == startpos, 3,
      "start position must be an integer");
  luaL_argcheck(L, (size_t)startpos <= len, 3,
      "start position must be less than or equal to the input string length");

  str += (size_t)startpos - 1;

  if (session->unpacker != LUA_REFNIL) {
    lmpack_geti(L, session->reg, session->unpacker);
    unpacker = lmpack_check_unpacker(L, -1);
    unpacker->L = L;
    rcount += 2;
    lua_pop(L, 1);
  }

  for (;;) {
    int result;

    if (session->unpacked.type == MPACK_EOF) {
      session->unpacked.type =
        mpack_rpc_receive(session->session, &str, &len, &session->unpacked.msg);

      if (!unpacker || session->unpacked.type == MPACK_EOF)
        break;
    }
    
    result = lmpack_unpacker_unpack_str(L, unpacker, &str, &len);

    if (result == MPACK_EOF) break;

    if (session->unpacked.method_or_error == LUA_NOREF) {
      session->unpacked.method_or_error = lmpack_ref(L, session->reg);
    } else {
      session->unpacked.args_or_result = lmpack_ref(L, session->reg);
      break;
    }
  }

  done = session->unpacked.type != MPACK_EOF
    && (session->unpacked.args_or_result != LUA_NOREF || !unpacker);

  if (!done) {
    lua_pushnil(L);
    lua_pushnil(L);
    if (unpacker) {
      lua_pushnil(L);
      lua_pushnil(L);
    }
    goto end;
  }

  switch (session->unpacked.type) {
    case MPACK_RPC_REQUEST:
      lua_pushstring(L, "request");
      lua_pushnumber(L, session->unpacked.msg.id);
      break;
    case MPACK_RPC_RESPONSE:
      lua_pushstring(L, "response");
      lmpack_geti(L, session->reg, (int)session->unpacked.msg.data.i);
      lmpack_unref(L, session->reg, (int)session->unpacked.msg.data.i);
      session->unpacked.msg.data.i = LUA_NOREF;
      break;
    case MPACK_RPC_NOTIFICATION:
      lua_pushstring(L, "notification");
      lua_pushnil(L);
      break;
    default:
      /* In most cases the only sane thing to do when receiving invalid
       * msgpack-rpc is to close the connection, so handle all errors with
       * this generic message. Later may add more detailed information. */
      return luaL_error(L, "invalid msgpack-rpc string");
  }

  session->unpacked.type = MPACK_EOF;

  if (unpacker) {
    lmpack_geti(L, session->reg, session->unpacked.method_or_error);
    lmpack_geti(L, session->reg, session->unpacked.args_or_result);
    lmpack_unref(L, session->reg, session->unpacked.method_or_error);
    lmpack_unref(L, session->reg, session->unpacked.args_or_result);
    session->unpacked.method_or_error = LUA_NOREF;
    session->unpacked.args_or_result = LUA_NOREF;
  }

end:
  lua_pushinteger(L, str - str_init + 1);
  return rcount;
}

static int lmpack_session_request(lua_State *L)
{
  int result;
  char buf[16], *b = buf;
  size_t bl = sizeof(buf);
  Session *session;
  mpack_data_t data;

  if (lua_gettop(L) > 2 || lua_gettop(L) < 1)
    return luaL_error(L, "expecting 1 or 2 arguments"); 

  session = lmpack_check_session(L, 1);
  data.i = lua_isnoneornil(L, 2) ? LUA_NOREF : lmpack_ref(L, session->reg);
  do {
    result = mpack_rpc_request(session->session, &b, &bl, data);
    if (result == MPACK_NOMEM) {
      session->session = lmpack_grow_session(session->session);
      if (!session->session)
        return luaL_error(L, "Failed to grow Session capacity");
    }
  } while (result == MPACK_NOMEM);

  assert(result == MPACK_OK);
  lua_pushlstring(L, buf, sizeof(buf) - bl);
  return 1;
}

static int lmpack_session_reply(lua_State *L)
{
  int result;
  char buf[16], *b = buf;
  size_t bl = sizeof(buf);
  Session *session;
  lua_Number id;

  if (lua_gettop(L) != 2)
    return luaL_error(L, "expecting exactly 2 arguments"); 

  session = lmpack_check_session(L, 1);
  id = lua_tonumber(L, 2);
  luaL_argcheck(L, ((size_t)id == id && id >= 0 && id <= 0xffffffff), 2,
      "invalid request id");
  result = mpack_rpc_reply(session->session, &b, &bl, (mpack_uint32_t)id);
  assert(result == MPACK_OK);
  lua_pushlstring(L, buf, sizeof(buf) - bl);
  return 1;
}

static int lmpack_session_notify(lua_State *L)
{
  int result;
  char buf[16], *b = buf;
  size_t bl = sizeof(buf);
  Session *session;

  if (lua_gettop(L) != 1)
    return luaL_error(L, "expecting exactly 1 argument"); 

  session = lmpack_check_session(L, 1);
  result = mpack_rpc_notify(session->session, &b, &bl);
  assert(result == MPACK_OK);
  lua_pushlstring(L, buf, sizeof(buf) - bl);
  return 1;
}

static int lmpack_nil_tostring(lua_State* L)
{
  lua_pushfstring(L, NIL_NAME, lua_topointer(L, 1));
  return 1;
}

static int lmpack_unpack(lua_State *L)
{
  int result;
  size_t len;
  const char *str;
  Unpacker unpacker;
  mpack_parser_t parser;

  if (lua_gettop(L) != 1)
    return luaL_error(L, "expecting exactly 1 argument"); 

  str = luaL_checklstring(L, 1, &len);

  /* initialize unpacker */
  lua_newtable(L);
  unpacker.reg = luaL_ref(L, LUA_REGISTRYINDEX);
  unpacker.ext = LUA_NOREF;
  unpacker.parser = &parser;
  mpack_parser_init(unpacker.parser, 0);
  unpacker.parser->data.p = &unpacker;
  unpacker.string_buffer = NULL;
  unpacker.L = L;

  result = mpack_parse(&parser, &str, &len, lmpack_parse_enter,
      lmpack_parse_exit);

  luaL_unref(L, LUA_REGISTRYINDEX, unpacker.reg);

  if (result == MPACK_NOMEM)
    return luaL_error(L, "object was too deep to unpack");
  else if (result == MPACK_EOF)
    return luaL_error(L, "incomplete msgpack string");
  else if (result == MPACK_ERROR)
    return luaL_error(L, "invalid msgpack string");
  else if (result == MPACK_OK && len)
    return luaL_error(L, "trailing data in msgpack string");

  assert(result == MPACK_OK);
  return 1;
}

static int lmpack_pack(lua_State *L)
{
  char *b;
  size_t bl;
  int result;
  Packer packer;
  mpack_parser_t parser;
  luaL_Buffer buffer;

  if (lua_gettop(L) != 1)
    return luaL_error(L, "expecting exactly 1 argument"); 

  /* initialize packer */
  lua_newtable(L);
  packer.reg = luaL_ref(L, LUA_REGISTRYINDEX);
  packer.ext = LUA_NOREF;
  packer.parser = &parser;
  mpack_parser_init(packer.parser, 0);
  packer.parser->data.p = &packer;
  packer.is_bin = 0;
  packer.L = L;
  packer.root = lmpack_ref(L, packer.reg);

  luaL_buffinit(L, &buffer);
  b = luaL_prepbuffer(&buffer);
  bl = LUAL_BUFFERSIZE;

  do {
    size_t bl_init = bl;
    result = mpack_unparse(packer.parser, &b, &bl, lmpack_unparse_enter,
        lmpack_unparse_exit);

    if (result == MPACK_NOMEM) {
      lmpack_unref(L, packer.reg, packer.root);
      luaL_unref(L, LUA_REGISTRYINDEX, packer.reg);
      return luaL_error(L, "object was too deep to pack");
    }

    luaL_addsize(&buffer, bl_init - bl);

    if (!bl) {
      /* buffer empty, resize */
      b = luaL_prepbuffer(&buffer);
      bl = LUAL_BUFFERSIZE;
    }
  } while (result == MPACK_EOF);

  lmpack_unref(L, packer.reg, packer.root);
  luaL_unref(L, LUA_REGISTRYINDEX, packer.reg);
  luaL_pushresult(&buffer);
  return 1;
}

static const luaL_reg unpacker_methods[] = {
  {"__call", lmpack_unpacker_unpack},
  {"__gc", lmpack_unpacker_delete},
  {NULL, NULL}
};

static const luaL_reg packer_methods[] = {
  {"__call", lmpack_packer_pack},
  {"__gc", lmpack_packer_delete},
  {NULL, NULL}
};

static const luaL_reg session_methods[] = {
  {"receive", lmpack_session_receive},
  {"request", lmpack_session_request},
  {"reply", lmpack_session_reply},
  {"notify", lmpack_session_notify},
  {"__gc", lmpack_session_delete},
  {NULL, NULL}
};

static const luaL_reg mpack_functions[] = {
  {"Unpacker", lmpack_unpacker_new},
  {"Packer", lmpack_packer_new},
  {"Session", lmpack_session_new},
  {"unpack", lmpack_unpack},
  {"pack", lmpack_pack},
  {NULL, NULL}
};

int luaopen_mpack(lua_State *L)
{
  /* Unpacker */
  luaL_newmetatable(L, UNPACKER_META_NAME);
  lua_pushvalue(L, -1);
  lua_setfield(L, -2, "__index");
  luaL_register(L, NULL, unpacker_methods);
  /* Packer */
  luaL_newmetatable(L, PACKER_META_NAME);
  lua_pushvalue(L, -1);
  lua_setfield(L, -2, "__index");
  luaL_register(L, NULL, packer_methods);
  /* Session */
  luaL_newmetatable(L, SESSION_META_NAME);
  lua_pushvalue(L, -1);
  lua_setfield(L, -2, "__index");
  luaL_register(L, NULL, session_methods);
  /* NIL */
  /* Check if NIL is already stored in the registry */
  lua_getfield(L, LUA_REGISTRYINDEX, NIL_NAME);
  /* If it isn't, create it */
  if (lua_isnil(L, -1)) {
    /* Use a constant userdata to represent NIL */
    (void)lua_newuserdata(L, sizeof(void *));
    /* Create a metatable for NIL userdata */
    lua_createtable(L, 0, 1);
    lua_pushstring(L, "__tostring");
    lua_pushcfunction(L, lmpack_nil_tostring);
    lua_settable(L, -3);
    /* Assign the metatable to the userdata object */
    lua_setmetatable(L, -2);
    /* Save NIL on the registry so we can access it easily from other functions */
    lua_setfield(L, LUA_REGISTRYINDEX, NIL_NAME);
  }
  /* module */
  lua_newtable(L);
  luaL_register(L, NULL, mpack_functions);
  /* save NIL on the module */
  lua_getfield(L, LUA_REGISTRYINDEX, NIL_NAME);
  lua_setfield(L, -2, "NIL");
  return 1;
}