#include "luastack.hpp"
#include "table.hpp"
#include "idalloc.hpp"

LuaDefine(idalloc_initmaster, "c") {
    LuaArg allocator, queuefill;
    LuaVar salvqueue;
    LuaStack LS(L, allocator, queuefill, salvqueue);

    // Use the registry by default.
    if (LS.isnil(allocator)) LS.set(allocator, LuaRegistry);

    LS.checktable(allocator);
    LS.checknumber(queuefill);
    LS.call(salvqueue, queue_create);
    LS.setfield(allocator, "id_salvaged", salvqueue);
    LS.setfield(allocator, "id_nextbatch", 0x0001000000000000);
    LS.setfield(allocator, "id_nextid", 0x0010000000000000);
    LS.setfield(allocator, "id_queuefill", queuefill);
    return LS.result();
}

LuaDefine(idalloc_initsynch, "c") {
    LuaArg allocator, queuefill;
    LuaStack LS(L, allocator, queuefill);

    // Use the registry by default.
    if (LS.isnil(allocator)) LS.set(allocator, LuaRegistry);

    LS.checktable(allocator);
    LS.checknumber(queuefill);
    LS.setfield(allocator, "id_salvaged", LuaNil);
    LS.setfield(allocator, "id_nextbatch", LuaNil);
    LS.setfield(allocator, "id_nextid", 0x001E000000000000);
    LS.setfield(allocator, "id_queuefill", queuefill);
    return LS.result();
}

LuaDefine(idalloc_refill, "c") {
    LuaArg allocator, queue;
    LuaVar salvaged, batch;
    lua_Integer qhead, qtail, shead, stail, nextb, queuefill;
    LuaStack LS(L, allocator, queue, salvaged, batch);

    // Use the registry by default.
    if (LS.isnil(allocator)) LS.set(allocator, LuaRegistry);

    // Get salvaged batch table.  If there is none, we're in donotpredict mode.
    LS.getfield(salvaged, allocator, "id_salvaged");
    if (LS.isnil(salvaged)) {
        return LS.result();
    }

    // Get the head and tail of the queue.
    LS.checktable(queue);
    LS.getfield(qhead, queue, "head");
    LS.getfield(qtail, queue, "tail");
    lua_Integer oqhead = qhead;

    // Try using salvaged batches.
    LS.getfield(queuefill, allocator, "id_queuefill");
    if (qhead - qtail < queuefill) {
        // Grab the head and tail of the salvaged batches.
        LS.checktable(salvaged);
        LS.getfield(shead, salvaged, "head");
        LS.getfield(stail, salvaged, "tail");

        // Get salvaged batches where possible.
        while ((stail < shead) && (qhead - qtail < queuefill)) {
            LS.rawget(batch, salvaged, stail);
            LS.rawset(salvaged, stail, LuaNil);
            stail += 1;
            LS.checknumber(batch);
            LS.rawset(queue, qhead, batch);
            qhead += 1;
        }

        // Update the head and tail of the salvaged batches.
        LS.setfield(salvaged, "head", shead);
        LS.setfield(salvaged, "tail", stail);
    }

    // Try using newly-created batches.
    if (qhead - qtail < queuefill) {
        // Grab the next batch counter.
        LS.getfield(nextb, allocator, "id_nextbatch");

        // Get newly-allocated batches to fill the rest.
        while (qhead - qtail < queuefill) {
            LS.rawset(queue, qhead, nextb);
            nextb += 256;
            qhead += 1;
        }

        // Update the counter in the registry.
        LS.setfield(allocator, "id_nextbatch", nextb);
    }

    // Update the head of the queue.
    if (oqhead != qhead) {
        LS.setfield(queue, "head", qhead);
    }

    return LS.result();
}

LuaDefine(idalloc_unqueue, "c") {
    LuaArg allocator, queue;
    LuaVar salvaged, batch;
    lua_Integer qhead, qtail, shead, stail;
    LuaStack LS(L, allocator, queue, salvaged, batch);

    // Use the registry by default.
    if (LS.isnil(allocator)) LS.set(allocator, LuaRegistry);

    // Get the head and tail of the queue.
    LS.checktable(queue);
    LS.getfield(qhead, queue, "head");
    LS.getfield(qtail, queue, "tail");

    // Grab the table of salvaged batches.
    // If there is none, we're in donotpredict mode.  In that
    // case, just empty the queue and dump the batches.
    LS.getfield(salvaged, allocator, "id_salvaged");
    if (LS.isnil(salvaged)) {
        while (qhead > qtail) {
            LS.rawset(queue, qtail, LuaNil);
            qtail += 1;
        }
        LS.setfield(queue, "tail", qtail);
        return LS.result();
    }

    // We're in master mode.  Transfer batches from the queue
    // into the salvaged batches table.
    LS.checktable(salvaged);
    LS.getfield(shead, salvaged, "head");
    LS.getfield(stail, salvaged, "tail");

    while (qhead > qtail) {
        LS.rawget(batch, queue, qtail);
        LS.rawset(queue, qtail, LuaNil);
        qtail += 1;

        LS.checknumber(batch);
        LS.rawset(salvaged, shead, batch);
        shead += 1;
    }

    // Update the queue pointers.
    LS.setfield(queue, "tail", qtail);
    LS.setfield(salvaged, "head", shead);

    return LS.result();
}

LuaDefine(idalloc_preparethread, "c") {
    LuaArg queue, thread;
    LuaVar batch;
    LuaStack LS(L, queue, thread, batch);

    // Get the thread.
    lua_State *TH = LS.ckthread(thread);

    // Pop a batch from the queue.  If there's nothing in
    // the queue, just leave the thread unprepped.
    LS.call(batch, queue_pop, queue);
    if (LS.isnil(batch)) {
        return LS.result();
    }

    // Store the batch into the thread.
    lua_setnextid(TH, LS.ckinteger(batch));
    return LS.result();
}

LuaDefine(idalloc_salvagethread, "c") {
    LuaArg allocator, thread;
    LuaVar salvaged;
    LuaStack LS(L, allocator, thread, salvaged);
    
    lua_State *TH = LS.ckthread(thread);
    lua_Integer idbatch = lua_getnextid(TH);
    lua_setnextid(TH, 0);
    if (idbatch == 0) {
        return LS.result();
    }
    if ((idbatch & 0xFF) >= 128) {
        return LS.result();
    }

    // Use the registry by default.
    if (LS.isnil(allocator)) LS.set(allocator, LuaRegistry);

    // Push the batch onto the queue of salvaged batches.
    // If the table of salvaged batches is nil, we're in donotpredict
    // mode.  In that case, don't bother salvaging.
    LS.getfield(salvaged, allocator, "id_salvaged");
    if (LS.isnil(salvaged)) {
        return LS.result();
    }
    LS.call(queue_push, salvaged, idbatch);
    return LS.result();
}

LuaDefine(idalloc_allocid, "c") {
    LuaArg allocator;
    LuaRet result;
    LuaStack LS(L, allocator, result);
    
    lua_Integer id = lua_getnextid(L);
    if (id == 0) {
        // Use the registry by default.
        if (LS.isnil(allocator)) LS.set(allocator, LuaRegistry);
        LS.getfield(id, allocator, "id_nextid");
        LS.setfield(allocator, "id_nextid", id + 1);
    } else {
        lua_Integer next = id + 1;
        if ((next & 0xFF) == 0) {
            next = 0;
        }
        lua_setnextid(L, next);
    }
    LS.set(result, id);
    return LS.result();
}

LuaDefine(idalloc_getthreadbatch, "c") {
    LuaArg thread;
    LuaRet result;
    LuaStack LS(L, thread, result);

    lua_State *TH = LS.ckthread(thread);
    lua_Integer id = lua_getnextid(TH);
    LS.set(result, id);
    return LS.result();
}