#include "table.hpp"
#include "source.hpp"

LuaDefine(table_getregistry, "f") {
    LuaArg key;
    LuaRet result;
    LuaStack LS(L, key, result);
    LS.rawget(result, LuaRegistry, key);
    return LS.result();
}

LuaDefine(table_equal, "c") {
    LuaArg t1, t2;
    LuaRet eql;
    LuaStack LS(L, t1, t2, eql);
    LS.checktable(t1);
    LS.checktable(t2);
    lua_pushnil(L);
    int total1 = 0;
    while (lua_next(L, t1.index()) != 0) {
        lua_pushvalue(L, -2);   // k v1 k
        lua_rawget(L, t2.index());  // k v1 v2
        if (!lua_rawequal(L, -1, -2)) {
            LS.set(eql, false);
            return LS.result();
        }
        lua_pop(L, 2);
        total1 += 1;
    }
    int total2 = 0;
    lua_pushnil(L);
    while (lua_next(L, t2.index()) != 0) {
        lua_pop(L, 1);
        total2 += 1;
    }
    LS.set(eql, total1 == total2);
    return LS.result();
}

LuaDefine(table_findremove, "c") {
    luaL_checktype(L, -2, LUA_TTABLE);
    int src = 1;
    int dst = 1;
    while (true) {
        lua_pushinteger(L, src);
        lua_rawget(L, -3);
        if (lua_rawequal(L, -1, -2)) {
            src++;
            lua_pop(L, 1);
        } else if (lua_isnil(L, -1)) {
            lua_pop(L, 1);
            int removed = src - dst;
            while (src > dst) {
                lua_pushinteger(L, dst);
                lua_pushnil(L);
                lua_rawset(L, -4);
                dst++;
            }
            lua_pop(L, 2);
            lua_pushinteger(L, removed);
            return 1;
        } else {
            if (src > dst) {
                lua_pushinteger(L, dst);
                lua_insert(L, lua_gettop(L) - 1);
                lua_rawset(L, -4);
            } else {
                lua_pop(L, 1);
            }
            src++;
            dst++;
        }
    }
}


LuaDefine(table_push, "c") {
    luaL_checktype(L, -2, LUA_TTABLE);
    int len = lua_rawlen(L, -2);
    lua_pushinteger(L, len+1);
    lua_pushvalue(L, -2);
    lua_rawset(L, -4);
    lua_pop(L, 2);
    return 0;
}

LuaDefine(table_find, "c") {
    luaL_checktype(L, -2, LUA_TTABLE);
    for (int i = 1; ; i++) {
        lua_pushinteger(L, i);
        lua_rawget(L, -3);
        if (lua_rawequal(L, -1, -2)) {
            lua_pop(L, 3);
            lua_pushinteger(L, i);
            return 1;
        } else if (lua_isnil(L, -1)) {
            lua_pop(L, 3);
            lua_pushnil(L);
            return 1;
        } else {
            lua_pop(L, 1);
        }
    }
}

LuaDefine(table_empty, "c") {
    luaL_checktype(L, -1, LUA_TTABLE);
    lua_pushnil(L);
    if (lua_next(L, -2) != 0) {
        lua_pop(L, 3);
        lua_pushboolean(L, 0);
        return 1;
    } else {
        lua_pop(L, 1);
        lua_pushboolean(L, 1);
        return 1;
    }
}

LuaDefine(table_count, "c") {
    luaL_checktype(L, -1, LUA_TTABLE);
    lua_pushnil(L);
    lua_Integer total = 0;
    while (lua_next(L, -2) != 0) {
        total += 1;
        lua_pop(L, 1);
    }
    lua_pop(L, 1);
    lua_pushinteger(L, total);
    return 1;
}

LuaDefine(table_clear, "c") {
    LuaArg tab;
    LuaStack LS(L, tab);
    LS.cleartable(tab);
    return LS.result();
}

LuaDefine(queue_create, "c") {
    LuaRet queue;
    LuaStack LS(L, queue);

    LS.newtable(queue);
    LS.rawset(queue, "head", 1000000);
    LS.rawset(queue, "tail", 1000000);
    return LS.result();
}

LuaDefine(queue_push, "c") {
    LuaArg queue, elt;
    lua_Integer head;
    LuaStack LS(L, queue, elt);

    LS.rawget(head, queue, "head");
    LS.rawset(queue, head, elt);
    LS.rawset(queue, "head", head+1);
    return LS.result();
}

LuaDefine(queue_pop, "c") {
    LuaArg queue;
    LuaRet elt;
    lua_Integer head, tail;
    LuaStack LS(L, queue, elt);

    LS.rawget(tail, queue, "tail");
    LS.rawget(head, queue, "head");
    if (head == tail) {
        LS.set(elt, LuaNil);
    } else {
        LS.rawget(elt, queue, tail);
        LS.rawset(queue, tail, LuaNil);
        LS.rawset(queue, "tail", tail + 1);
    }
    return LS.result();
}

LuaDefine(queue_size, "c") {
    LuaArg queue;
    LuaRet size;
    lua_Number head, tail;
    LuaStack LS(L, queue, size);

    LS.rawget(head, queue, "head");
    LS.rawget(tail, queue, "tail");
    LS.set(size, head - tail);
    return LS.result();
}

LuaDefine(queue_nth, "c") {
    LuaArg queue, n;
    LuaRet elt;
    lua_Integer nth, head, tail;
    LuaStack LS(L, queue, n, elt);

    nth = LS.ckinteger(n) - 1;
    LS.rawget(head, queue, "head");
    LS.rawget(tail, queue, "tail");
    if ((nth < 0) || (nth + tail >= head)) {
        luaL_error(L, "index out of range");
    }
    LS.rawget(elt, queue, tail + nth);
    return LS.result();
}

/////////////////////////////////////////////////////////////
//
// table_sortedpairs
//
// Given a table, return all the pairs in the table as a sequence.
// The resulting sequence contains a 1 in the first position,
// followed by the pairs, like so:
//
// sortedpairs({a=1,b=2,c=3}) => {1,"a",1,"b",2,"c",3}
// 
// Note that this function is not directly exposed to lua.
//
/////////////////////////////////////////////////////////////

static void pushkey (lua_State *L, int tab, int i) {
    lua_rawgeti(L, tab, i*2);
}

static void push2keys (lua_State *L, int tab, int i, int j) {
    lua_rawgeti(L, tab, i*2);
    lua_rawgeti(L, tab, j*2);
}

static void pop2keys (lua_State *L, int tab, int i, int j) {
    lua_rawseti(L, tab, i*2);
    lua_rawseti(L, tab, j*2);
}

static void swap2values (lua_State *L, int tab, int i, int j) {
    lua_rawgeti(L, tab, i*2+1);
    lua_rawgeti(L, tab, j*2+1);
    lua_rawseti(L, tab, i*2+1);
    lua_rawseti(L, tab, j*2+1);
}

static void auxsort (lua_State *L, int tab, int l, int u) {
    while (l < u) { /* for tail recursion */
        int i, j;
        /* sort elements a[l], a[(l+u)/2] and a[u] */
        push2keys(L, tab, l, u);
        if (lua_genlt(L, -1, -2)) {
            pop2keys(L, tab, l, u);
            swap2values(L, tab, l, u);
        } else {
            lua_pop(L, 2);
        }
        if (u - l == 1) break; /* only 2 elements */
        i = (l + u) / 2;
        push2keys(L, tab, i, l);
        if (lua_genlt(L, -2, -1)) {
            pop2keys(L, tab, i, l);
            swap2values(L, tab, i, l);
        } else {
            lua_pop(L, 1); /* remove a[l] */
            pushkey(L, tab, u);
            if (lua_genlt(L, -1, -2)) {
                pop2keys(L, tab, i, u);
                swap2values(L, tab, i, u);
            } else {
                lua_pop(L, 2);
            }
        }
        if (u - l == 2) break; /* only 3 elements */
        /* put the pivot value on top of the stack and keep it there */
        pushkey(L, tab, i);
        /* move the pivot from i to u-1 */
        lua_pushvalue(L, -1);
        pushkey(L, tab, u-1);
        pop2keys(L, tab, i, u-1);
        swap2values(L, tab, i, u-1);
        /* a[l] <= P == a[u-1] <= a[u], only need to sort from l+1 to u-2 */
        i = l;
        j = u - 1;
        for (;;) { /* invariant: a[l..i] <= P <= a[j..u] */
            /* repeat ++i until a[i] >= P */
            while (pushkey(L, tab, ++i), lua_genlt(L, -1, -2)) {
                if (i >= u) luaL_error(L, "invalid order function for sorting");
                lua_pop(L, 1); /* remove a[i] */
            }
            /* repeat --j until a[j] <= P */
            while (pushkey(L, tab, --j), lua_genlt(L, -3, -1)) {
                if (j <= l) luaL_error(L, "invalid order function for sorting");
                lua_pop(L, 1); /* remove a[j] */
            }
            if (j < i) {
                lua_pop(L, 3); /* pop pivot, a[i], a[j] */
                break;
            }
            pop2keys(L, tab, i, j);
            swap2values(L, tab, i, j);
        }
        push2keys(L, tab, u-1, i);
        pop2keys(L, tab, u-1, i);
        swap2values(L, tab, u-1, i);
        /* a[l..i-1] <= a[i] == P <= a[i+1..u] */
        /* adjust so that smaller half is in [j..i] and larger one in [l..u] */
        if (i - l < u - i) {
            j = l;
            i = i - 1;
            l = i + 2;
        } else {
            j = i + 1;
            i = u;
            u = j - 2;
        }
        auxsort(L, tab, j, i); /* call recursively the smaller one */
    }                          /* repeat the routine for the larger one */
}

int table_sortedpairs(lua_State *L, bool *unsortable) {
    lua_checkstack(L, 40);
    LuaArg tab;
    LuaVar key, value;
    LuaRet pairs;
    LuaStack LS(L, tab, key, value, pairs);
    if (unsortable != nullptr) {
        *unsortable = false;
    }
    LS.checktable(tab);
    // Count the total number of pairs.
    // TODO: add lua_npairs to make this faster.
    lua_pushnil(L);
    int total = 0;
    while (lua_next(L, tab.index()) != 0) {
        total += 1;
        lua_pop(L, 1);
    }
    // Create the table, store the initial 1.
    LS.createtable(pairs, total * 2 + 1, 0);
    LS.rawseti(pairs, 1, 1);
    // Transfer the pairs into the sequence.
    lua_pushnil(L);
    int offset = 2;
    while (lua_next(L, tab.index()) != 0) {
        int ktype = lua_type(L, -2);
        if (ktype == LUA_TNUMBER || ktype == LUA_TSTRING || ktype == LUA_TBOOLEAN) {
            lua_pushvalue(L, -2); // K V K
            lua_rawseti(L, pairs.index(), offset++);
            lua_rawseti(L, pairs.index(), offset++);
        } else {
            if (unsortable != nullptr) {
                *unsortable = true;
            }
            lua_pop(L, 1);
        }
    }
    auxsort(L, pairs.index(), 1, total);
    return LS.result();
}

/////////////////////////////////////////////////////////////
//
// Given a sortedpairs vector, return the (key, value) pairs
// one by one.  The first element of the vector is used as a
// counter to keep track of our position.
//
/////////////////////////////////////////////////////////////

static int nextsortedpair(lua_State *L) {
    if (lua_gettop(L) < 2) {
        luaL_error(L, "Not enough arguments to nextpair");
    }
    luaL_checktype(L, 1, LUA_TTABLE);
    lua_rawgeti(L, 1, 1);
    lua_Integer i = luaL_checkinteger(L, -1);
    lua_pop(L, 1);
    lua_pushinteger(L, i+1);
    lua_rawseti(L, 1, 1);
    lua_rawgeti(L, 1, i*2);
    if (lua_isnil(L, -1)) {
        return 1;
    } else {
        lua_rawgeti(L, 1, i*2+1);
        return 2;
    }
}

/////////////////////////////////////////////////////////////
//
// Replace the lua 'pairs' and 'next' iterators.
//
/////////////////////////////////////////////////////////////

LuaDefine(table_pairs, "f") {
    LuaArg tab;
    LuaRet closure, rtab, key;
    LuaStack LS(L, tab, closure, rtab, key);
    lua_pushvalue(L, tab.index());
    bool unsortable;
    table_sortedpairs(L, &unsortable);
    if (unsortable) {
        luaL_error(L, "Cannot iterate over a table with unsortable keys");
    }
    lua_replace(L, rtab.index());
    LS.set(closure, nextsortedpair);
    LS.set(key, LuaNil);
    return LS.result();
}

LuaDefine(table_next, "f") {
    luaL_error(L, "The 'next' iterator is not allowed in this version of lua");
    return 0;
}

LuaDefine(table_genlt, "f") {
    LuaArg o1,o2;
    LuaRet lt;
    LuaStack LS(L, o1, o2, lt);
    int ltf = lua_genlt(L, o1.index(), o2.index());
    LS.set(lt, ltf ? true:false);
    return LS.result();
}