#ifndef LPX_STACK_MANAGER_HPP
#define LPX_STACK_MANAGER_HPP

extern "C" {
#include "lua.h"
#include "lauxlib.h"
#include "lualib.h"
#include "luajit.h"
}

// LpxArgCheck
//
// This is a template that, when given a lua_CFunction that doesn't
// verify the number of arguments on the stack, returns a new lua_CFunction
// that *does* verify the number of arguments on the stack.
//
// In general, your lua_CFunctions should not verify the number of
// arguments on the stack.  Instead, they should assume that the stack
// contains arbitrary stuff underneath the arguments.  That makes your
// lua_CFunctions callable from both C and Lua.  But it is desirable to
// check arguments (only when calling from lua), so it is useful to use
// LpxArgCheck when turning a lua_CFunction into a lua closure.
//
template<int N, lua_CFunction FN>
int LpxArgCheck(lua_State *L) {
    if (lua_gettop(L) != N) {
        lua_pushfstring(L, "expected %d arguments, got %d", N, lua_gettop(L));
        lua_error(L);
    }
    FN(L);
}

// LpxSlot
//
// An LpxSlot contains a lua stack index.  It is initialized by the
// LpxStackManager and contains the same index until the LpxStackManager
// is destroyed.  You can convert an LpxSlot into a lua stack index
// by simply coercing it to 'int'.
//
// There are three variants of LpxSlot that you can use: LpxArg (function
// argument), LpxRet (function return value), and LpxVar (function local
// variable).
//
class LpxSlot {
private:
    int index_;
    LpxSlot *next_;

public:
    inline operator int() const {
        return index_;
    }

    LpxSlot() {
        index_ = 0;
        next_ = 0;
    }
    friend class LpxStackManager;
};

class LpxArg : public LpxSlot {};
class LpxRet : public LpxSlot {};
class LpxVar : public LpxSlot {};

class LpxStackManager {
private:
    using SSList = LpxSlot*;

    SSList arg_;
    SSList ret_;
    SSList var_;

    int argtop_;
    int gaptop_;
    int vartop_;
    int rettop_;
    int finaltop_;
    
    int narg_;
    int ngap_;
    int nvar_;
    int nret_;

    lua_State *L_;

    static int sscount(SSList l) {
        int total = 0;
        while (l != 0) {
            total += 1;
            l = l->next_;
        }
        return total;
    }

    void ck() {
        if (lua_gettop(L_) != rettop_) {
            lua_pushstring(L_, "LpxStackManager: stack is not right");
            lua_error(L_);
        }
    }

    template<class... SS>
    void record(LpxArg &v, SS & ... stackslots)
    {
        v.next_ = arg_; arg_ = &v;
        record(stackslots...);
    }

    template<class... SS>
    void record(LpxRet &v, SS & ... stackslots)
    {
        v.next_ = ret_; ret_ = &v;
        record(stackslots...);
    }

    template<class... SS>
    void record(LpxVar &v, SS & ... stackslots)
    {
        v.next_ = var_; var_ = &v;
        record(stackslots...);
    }

    void record()
    {
    }

public:
    template<class... SS>
    LpxStackManager(lua_State *L, SS & ... stackslots) {
        L_ = L;
        arg_ = 0;
        ret_ = 0;
        var_ = 0;

        record(stackslots...);
        
        int narg = sscount(arg_);
        int nvar = sscount(var_);
        int nret = sscount(ret_);
        int ngap = (nret - nvar - narg);
        if (ngap < 0) ngap = 0;

        argtop_ = lua_gettop(L_);
        gaptop_ = argtop_ + ngap;
        vartop_ = gaptop_ + nvar;
        rettop_ = vartop_ + nret;
        finaltop_ = argtop_ - narg + nret;
        
        int i = argtop_;
        for (LpxSlot *v = arg_; v != 0; v = v->next_) {
            v->index_ = i;
            i -= 1;
        }
        i = vartop_;
        for (LpxSlot *v = var_; v != 0; v = v->next_) {
            v->index_ = i;
            i -= 1;
        }
        i = rettop_;
        for (LpxSlot *v = ret_; v != 0; v = v->next_) {
            v->index_ = i;
            i -= 1;
        }

        // Initialize vars and rets to NIL.
        // I could conceivably do a Lpx_settop instead.
        for (int i = argtop_; i < rettop_; i++) {
            lua_pushnil(L_);
        }
    }

    ~LpxStackManager() {
        ck();
        int i = finaltop_;
        for (LpxSlot *v = ret_; v != 0; v = v->next_) {
            lua_replace(L_, i);
            i -= 1;
        }
        lua_settop(L_, finaltop_);
    }

    int retval() {
        return rettop_ - vartop_;
    }

    static void reg(lua_State *L, const char *classname, const char *funcname, lua_CFunction fn) {
        luaL_Reg reg;
        reg.name = funcname;
        reg.func = fn;
        luaL_register(L, classname, &reg);
    }

public:
    // The following functions are 'register-machine' variants
    // of the core lua functions.  They do not read from the stack,
    // or leave results on the stack.

    void setnil(int target) {
        lua_pushnil(L_);
        lua_replace(L_, target);
    }

    void setstring(int target, const char *str) {
        lua_pushstring(L_, str);
        lua_replace(L_, target);
    }

    void setlstring(int target, const char *str, int len) {
        lua_pushlstring(L_, str, len);
        lua_replace(L_, target);
    }

    void setnumber(int target, double value) {
        lua_pushnumber(L_, value);
        lua_replace(L_, target);
    }

    int next(int tab, int key, int value) {
        lua_pushvalue(L_, key);
        int ret = lua_next(L_, tab);
        if (ret != 0) {
            lua_replace(L_, value);
            lua_replace(L_, key);
        }
        return ret;
    }
    
    void rawget(int target, int tab, int key) {
        lua_pushvalue(L_, key);
        lua_rawget(L_, tab);
        lua_replace(L_, target);
    }

    void rawset(int tab, int key, int value) {
        lua_pushvalue(L_, key);
        lua_pushvalue(L_, value);
        lua_rawset(L_, tab);
    }

    void setfield(int tab, const char *field, int value) {
        lua_pushvalue(L_, value);
        lua_setfield(L_, tab, field);
    }

    void getfield(int target, int tab, const char *field) {
        lua_getfield(L_, tab, field);
        lua_replace(L_, target);
    }

    void newtable(int target) {
        lua_newtable(L_);
        lua_replace(L_, target);
    }

    void checknometa(int index) {
        if (lua_istable(L_, index)) {
            if (!lua_getmetatable(L_, index)) {
                return;
            }
        }
        luaL_error(L_, "expected simple table with no metatable");
    }
};



#endif // LPX_STACK_MANAGER_HPP