////////////////////////////////////////////////////////////////////
//
// This file contains the code to pair up tables in the synchronous
// model with tables in the master model.  Here are the top-level
// routines in this file:
//
//  World::number_lua_tables
//  World::pair_lua_tables
//  World::number_remaining_tables
//  World::create_new_tables
//  World::unnumber_lua_tables
//
////////////////////////////////////////////////////////////////////

#include "luastack.hpp"
#include "streambuffer.hpp"
#include "world.hpp"

int World::number_lua_tables(const IdVector &basis) {
    // This is conceptually recursive, but we're going to use an
    // explicit stack (the lua stack).
    lua_State *L = state();
    int nextid = 1;
    {
        LuaVar tnmap, ntmap, tangibles, tab, key, val, xid;
        LuaExtStack LS(L, tnmap, ntmap, tangibles, tab, key, val, xid);
        LS.set(tnmap, LuaNewTable);
        LS.set(ntmap, LuaNewTable);
        LS.rawset(LuaRegistry, "tnmap", tnmap);
        LS.rawset(LuaRegistry, "ntmap", ntmap);
        LS.rawget(tangibles, LuaRegistry, "tangibles");
        int top = lua_gettop(L);

        // Push all subtables onto the stack.  Note that we may push
        // the same table twice, that's OK.
        for (int64_t id : basis) {
            LS.rawget(tab, tangibles, id);
            assert(LS.istable(tab));
            // Traverse subtables.
            LS.set(key, LuaNil);
            while (LS.next(tab, key, val)) {
                if (LS.istable(val) && LS.gettabletype(val)==LUA_TT_GENERAL) {
                    lua_checkstack(L, 10);
                    lua_pushvalue(L, val.index());
                }
            }
        }

        // Pop tables from the stack one by one.  If the table is not
        // already numbered, number it and push subtables onto the stack.
        while (lua_gettop(L) > top) {
            lua_replace(L, tab.index());
            LS.rawget(xid, tnmap, tab);
            if (LS.isnil(xid)) {
                int id = nextid++;
                LS.rawset(tnmap, tab, id);
                LS.rawset(ntmap, id, tab);
                // Traverse the metatable.
                LS.getmetatable(val, tab);
                if (LS.istable(val) && LS.gettabletype(val)==LUA_TT_GENERAL) {
                    lua_checkstack(L, 10);
                    lua_pushvalue(L, val.index());
                }
                // Traverse the subtables.
                LS.set(key, LuaNil);
                while (LS.next(tab, key, val)) {
                    if (LS.istable(val) && LS.gettabletype(val)==LUA_TT_GENERAL) {
                        lua_checkstack(L, 10);
                        lua_pushvalue(L, val.index());
                    }
                }
            }
        }
    }

    assert(stack_is_clear());
    return nextid - 1;
}

void World::pair_lua_tables(const IdVector &basis, lua_State *master) {
    lua_State *synch = state();
    LuaVar stangibles, mtangibles, sntmap, mntmap, stnmap, mtnmap, stab, mtab, skey, mkey, sval, mval, sidx, midx;
    LuaExtStack SLS(synch, stangibles, stab, skey, sval, sntmap, stnmap, sidx);
    LuaExtStack MLS(master, mtangibles, mtab, mkey, mval, mntmap, mtnmap, midx);

    // Fetch the tangible databases
    SLS.rawget(stangibles, LuaRegistry, "tangibles");
    MLS.rawget(mtangibles, LuaRegistry, "tangibles");

    // Fetch the synchronous model tnmap and ntmap
    SLS.rawget(stnmap, LuaRegistry, "tnmap");
    SLS.rawget(sntmap, LuaRegistry, "ntmap");
    assert(SLS.istable(stnmap));
    assert(SLS.istable(sntmap));

    // Initialize the master model tnmap and ntmap
    MLS.set(mtnmap, LuaNewTable);
    MLS.set(mntmap, LuaNewTable);
    MLS.rawset(LuaRegistry, "tnmap", mtnmap);
    MLS.rawset(LuaRegistry, "ntmap", mntmap);
    int s_ntables = SLS.rawlen(sntmap);
    for (int i = 1; i <= s_ntables; i++) {
        MLS.rawset(mntmap, i, 0);
    }

    // Keep track of which tables are already paired
    eng::vector<bool> paired;
    paired.assign(s_ntables + 1, false);

    // This records the top of the stack.
    int mtop = lua_gettop(master);

    for (int64_t id : basis) {
        MLS.rawget(mtab, mtangibles, id);
        SLS.rawget(stab, stangibles, id);
        assert(MLS.istable(mtab));
        assert(SLS.istable(stab));
        MLS.set(mkey, LuaNil);
        while (MLS.next(mtab, mkey, mval)) {
            if (!MLS.issortablekey(mkey)) continue;
            if (!MLS.istable(mval)) continue;
            MLS.movesortablekey(mkey, SLS, skey);
            SLS.rawget(sval, stab, skey);
            if (!SLS.istable(sval)) continue;
            lua_checkstack(master, 20);
            lua_checkstack(synch, 20);
            lua_pushvalue(master, mval.index());
            lua_pushvalue(synch, sval.index());
        }
    }

    while (lua_gettop(master) > mtop) {
        lua_replace(master, mtab.index());
        lua_replace(synch, stab.index());
        // If the master table is not a general table, skip.
        if (MLS.gettabletype(mtab) != LUA_TT_GENERAL) continue;
        // If the master table is already paired, skip.
        MLS.rawget(midx, mtnmap, mtab);
        if (MLS.isnumber(midx)) continue;
        // If the synch table doesn't have a number, skip.
        SLS.rawget(sidx, stnmap, stab);
        auto idx = SLS.trynumber(sidx);
        if (!idx) continue;
        assert((*idx >= 1) && (*idx <= s_ntables));
        // Pair the tables.
        MLS.rawset(mtnmap, mtab, *idx);
        MLS.rawset(mntmap, *idx, mtab);
        paired[*idx] = true;
        // Potentially pair the metatables.
        MLS.getmetatable(mval, mtab);
        if (MLS.istable(mval)) {
            SLS.getmetatable(sval, stab);
            if (SLS.istable(sval)) {
                lua_pushvalue(master, mval.index());
                lua_pushvalue(synch, sval.index());
            }
        }
        // Pair the subtables.
        MLS.set(mkey, LuaNil);
        while (MLS.next(mtab, mkey, mval)) {
            if (!MLS.issortablekey(mkey)) continue;
            if (!MLS.istable(mval)) continue;
            MLS.movesortablekey(mkey, SLS, skey);
            SLS.rawget(sval, stab, skey);
            if (!SLS.istable(sval)) continue;
            lua_checkstack(master, 20);
            lua_checkstack(synch, 20);
            lua_pushvalue(master, mval.index());
            lua_pushvalue(synch, sval.index());
        }
    }
}

int World::number_remaining_tables(const IdVector &basis, lua_State *master) {
    // This is conceptually recursive, but we're going to use an
    // explicit stack (the lua stack).
    eng::vector<bool> visited;
    int ntables;
    {
        lua_State *L = master;
        LuaVar tnmap, ntmap, tangibles, tab, key, val, xid;
        LuaExtStack LS(L, tnmap, ntmap, tangibles, tab, key, val, xid);
        LS.rawget(tnmap, LuaRegistry, "tnmap");
        LS.rawget(ntmap, LuaRegistry, "ntmap");
        LS.rawget(tangibles, LuaRegistry, "tangibles");
        ntables = LS.rawlen(ntmap);
        visited.assign(ntables + 1, false);
        int top = lua_gettop(L);

        // Push all subtables onto the stack.  Note that we may push
        // the same table twice, that's OK.
        for (int64_t id : basis) {
            LS.rawget(tab, tangibles, id);
            assert(LS.istable(tab));
            LS.set(key, LuaNil);
            while (LS.next(tab, key, val)) {
                if (LS.istable(val) && LS.gettabletype(val)==LUA_TT_GENERAL) {
                    lua_checkstack(L, 10);
                    lua_pushvalue(L, val.index());
                }
            }
        }

        // Pop tables from the stack one by one.  If the table is not
        // numbered, number it.  If it is not visited, visit it.
        while (lua_gettop(L) > top) {
            lua_replace(L, tab.index());
            int id = 0;
            LS.rawget(xid, tnmap, tab);
            if (!LS.isnumber(xid)) {
                id = visited.size();
                LS.rawset(tnmap, tab, id);
                LS.rawset(ntmap, id, tab);
                visited.push_back(false);
            } else {
                id = LS.cknumber(xid);
                assert((id >= 0) && (id < int(visited.size())));
            }
            if (!visited[id]) {
                visited[id] = true;
                // Traverse the metatable.
                LS.getmetatable(val, tab);
                if (LS.istable(val) && LS.gettabletype(val)==LUA_TT_GENERAL) {
                    lua_checkstack(L, 10);
                    lua_pushvalue(L, val.index());
                }
                // Traverse the subtables.
                LS.set(key, LuaNil);
                while (LS.next(tab, key, val)) {
                    if (LS.istable(val) && LS.gettabletype(val)==LUA_TT_GENERAL) {
                        lua_checkstack(L, 10);
                        lua_pushvalue(L, val.index());
                    }
                }
            }
        }
    }
    assert(stack_is_clear());
    return visited.size() - 1 - ntables;
}

void World::create_new_tables(int n) {
    {
        LuaVar tnmap, ntmap, tab;
        LuaExtStack LS(state(), tnmap, ntmap, tab);
        LS.rawget(tnmap, LuaRegistry, "tnmap");
        LS.rawget(ntmap, LuaRegistry, "ntmap");
        assert(LS.istable(tnmap));
        assert(LS.istable(ntmap));
        int ntables = LS.rawlen(ntmap);
        int nextid = ntables + 1;
        for (int i = 0; i < n; i++) {
            int id = nextid++;
            LS.newtable(tab);
            LS.rawset(ntmap, id, tab);
            LS.rawset(tnmap, tab, id);
        }
    }
    assert(stack_is_clear());
}

void World::unnumber_lua_tables() {
    // All we have to do is remove these tables from the registry.
    LuaExtStack LS(state());
    LS.rawset(LuaRegistry, "tnmap", LuaNil);
    LS.rawset(LuaRegistry, "ntmap", LuaNil);
}

