/*
 *  author: xjdrew
 *  date: 2014-06-03 20:38
 */

#include <stdio.h>
#include <stdlib.h>

#include "lauxlib.h"
#include "lua.h"
#include "skiplist.h"

static inline skiplist *
_to_skiplist(lua_State *L) {
    skiplist **sl = lua_touserdata(L, 1);
    if (sl == NULL) {
        luaL_error(L, "must be skiplist object");
    }
    return *sl;
}

static int
_insert(lua_State *L) {
    skiplist *sl = _to_skiplist(L);
    double score = luaL_checknumber(L, 2);
    lua_Integer obj = luaL_checkinteger(L, 3);
    slInsert(sl, score, obj);
    return 0;
}

static int
_delete(lua_State *L) {
    skiplist *sl = _to_skiplist(L);
    double score = luaL_checknumber(L, 2);
    lua_Integer obj = luaL_checkinteger(L, 3);
    lua_pushboolean(L, slDelete(sl, score, obj));
    return 1;
}

static void
_delete_rank_cb(void *ud, int64_t obj) {
    lua_State *L = (lua_State *)ud;
    lua_pushvalue(L, 4);
    lua_pushinteger(L, obj);
    lua_call(L, 1, 0);
}

static int
_delete_by_rank(lua_State *L) {
    skiplist *sl = _to_skiplist(L);
    unsigned int start = luaL_checkinteger(L, 2);
    unsigned int end = luaL_checkinteger(L, 3);
    luaL_checktype(L, 4, LUA_TFUNCTION);
    if (start > end) {
        unsigned int tmp = start;
        start = end;
        end = tmp;
    }

    lua_pushinteger(L, slDeleteByRank(sl, start, end, _delete_rank_cb, L));
    return 1;
}

static int
_get_count(lua_State *L) {
    skiplist *sl = _to_skiplist(L);
    lua_pushinteger(L, sl->length);
    return 1;
}

static int
_get_rank(lua_State *L) {
    skiplist *sl = _to_skiplist(L);
    double score = luaL_checknumber(L, 2);
    lua_Integer obj = luaL_checkinteger(L, 3);

    unsigned long rank = slGetRank(sl, score, obj);
    if (rank == 0) {
        return 0;
    }

    lua_pushinteger(L, rank);
    return 1;
}

static int
_get_rank_range(lua_State *L) {
    skiplist *sl = _to_skiplist(L);
    unsigned long r1 = luaL_checkinteger(L, 2);
    unsigned long r2 = luaL_checkinteger(L, 3);
    int reverse, rangelen;
    if (r1 <= r2) {
        reverse = 0;
        rangelen = r2 - r1 + 1;
    } else {
        reverse = 1;
        rangelen = r1 - r2 + 1;
    }

    skiplistNode *node = slGetNodeByRank(sl, r1);
    lua_createtable(L, rangelen, 0);
    int n = 0;
    while (node && n < rangelen) {
        n++;
        lua_pushinteger(L, node->obj);
        lua_rawseti(L, -2, n);
        node = reverse ? node->backward : node->level[0].forward;
    }
    return 1;
}

static int
_get_rank_score_range(lua_State *L) {
    skiplist *sl = _to_skiplist(L);
    unsigned long r1 = luaL_checkinteger(L, 2);
    unsigned long r2 = luaL_checkinteger(L, 3);
    int reverse, rangelen;
    if (r1 <= r2) {
        reverse = 0;
        rangelen = r2 - r1 + 1;
    } else {
        reverse = 1;
        rangelen = r1 - r2 + 1;
    }

    skiplistNode *node = slGetNodeByRank(sl, r1);
    lua_createtable(L, rangelen, 0);
    int n = 0;
    while (node && n < rangelen) {
        n++;
        lua_createtable(L, 2, 0);
        lua_pushinteger(L, node->obj);
        lua_rawseti(L, -2, 1);
        lua_pushinteger(L, node->score);
        lua_rawseti(L, -2, 2);

        lua_rawseti(L, -2, n);
        node = reverse ? node->backward : node->level[0].forward;
    }
    return 1;
}

static int
_get_score_range(lua_State *L) {
    skiplist *sl = _to_skiplist(L);
    double s1 = luaL_checknumber(L, 2);
    double s2 = luaL_checknumber(L, 3);
    int reverse;
    skiplistNode *node;

    if (s1 <= s2) {
        reverse = 0;
        node = slFirstInRange(sl, s1, s2);
    } else {
        reverse = 1;
        node = slLastInRange(sl, s2, s1);
    }

    lua_newtable(L);
    int n = 0;
    while (node) {
        if (reverse) {
            if (node->score < s2)
                break;
        } else {
            if (node->score > s2)
                break;
        }
        n++;

        lua_pushinteger(L, node->obj);
        lua_rawseti(L, -2, n);

        node = reverse ? node->backward : node->level[0].forward;
    }
    return 1;
}

static int
_get_member_by_rank(lua_State *L){
    skiplist *sl = _to_skiplist(L);
    unsigned long r = luaL_checkinteger(L, 2);
    skiplistNode *node = slGetNodeByRank(sl, r);
    if (node) {
        lua_pushinteger(L, node->obj);
        return 1;
    }
    return 0;
}

static int
_new(lua_State *L) {
    skiplist *psl = slCreate();

    skiplist **sl = (skiplist **)lua_newuserdata(L, sizeof(skiplist *));
    *sl = psl;
    lua_pushvalue(L, lua_upvalueindex(1));
    lua_setmetatable(L, -2);
    return 1;
}

static int
_release(lua_State *L) {
    skiplist *sl = _to_skiplist(L);
    //printf("collect sl:%p\n", sl);
    slFree(sl);
    return 0;
}

LUAMOD_API int
luaopen_skiplist_c(lua_State *L) {
    luaL_checkversion(L);

    luaL_Reg l[] = {
        { "insert", _insert },
        { "delete", _delete },
        { "delete_by_rank", _delete_by_rank },

        { "get_count", _get_count },
        { "get_rank", _get_rank },
        { "get_rank_range", _get_rank_range },
        { "get_score_range", _get_score_range },
        { "get_member_by_rank", _get_member_by_rank},

        { "get_rank_score_range", _get_rank_score_range },

        { NULL, NULL }
    };

    lua_createtable(L, 0, 2);

    luaL_newlib(L, l);
    lua_setfield(L, -2, "__index");
    lua_pushcfunction(L, _release);
    lua_setfield(L, -2, "__gc");

    lua_pushcclosure(L, _new, 1);
    return 1;
}