
#include "wrap-sstream.hpp"
#include "wrap-string.hpp"

#include "luasnap.hpp"
#include "luastack.hpp"
#include "streambuffer.hpp"

#include <cassert>

class LuaByteReader {
private:
    const char *data_;
    int64_t size_;
public:
    LuaByteReader(const char *d, int64_t s) : data_(d), size_(s) {}
    static const char *lua_reader(lua_State *L, void *ud, size_t *size) {
        LuaByteReader *reader = (LuaByteReader*)ud;
        *size = reader->size_;
        const char *data = reader->data_;
        reader->data_ = 0;
        reader->size_ = 0;
        return data;
    }
};

LuaSnap::LuaSnap() {
    state_ = LuaCoreStack::newstate(eng::l_alloc);
    LuaExtStack LS(state_);
}

LuaSnap::~LuaSnap() {
    lua_close(state_);
}

void LuaSnap::serialize(StreamBuffer *sb) {
    // Lua stack should be empty.
    assert(lua_gettop(state_) == 0);

    // We'll allocate stack slots manually, to ensure
    // that the thread serializes as we expect it to.
    lua_pushnil(state_); // permstable
    lua_pushnil(state_); // regcopy
    lua_pushnil(state_); // key
    lua_pushnil(state_); // value
    LuaSpecial permstable(1);
    LuaSpecial regcopy(2);
    LuaSpecial key(3);
    LuaSpecial value(4);
    LuaCoreStack LS(state_);

    // Construct a copy of the registry table.
    LS.set(regcopy, LuaNewTable);
    LS.set(key, LuaNil);
    while (LS.next(LuaRegistry, key, value) != 0) {
        LS.rawset(regcopy, key, value);
    }

    // Remove certain things from the copy.  These items
    // don't get included in the snapshot.
    //
    // From a modularity perspective, this is bad.
    //
    LS.rawset(regcopy, LUA_RIDX_MAINTHREAD, LuaNil);
    LS.rawset(regcopy, "persist", LuaNil);
    LS.rawset(regcopy, "unpersist", LuaNil);
    LS.rawset(regcopy, "world", LuaNil);
    LS.rawset(regcopy, "gui", LuaNil);
    LS.rawset(regcopy, "tnmap", LuaNil);
    LS.rawset(regcopy, "ntmap", LuaNil);

    // Get the eris permanents table from the registry.
    LS.rawget(permstable, LuaRegistry, "persist");
    assert(LS.istable(permstable));

    // Remove key and value from stack, leaving permstable and regcopy.
    lua_settop(state_, 2);

    // Write dummy length, use eris to write data, then overwrite length.
    sb->write_int64(0);
    int64_t pos1 = sb->total_writes();

    eris_dump(state_, lua_writer_into_streambuffer, sb);
    int64_t pos2 = sb->total_writes();
    sb->overwrite_int64(pos1, pos2 - pos1);
    lua_settop(state_, 0);
}

void LuaSnap::deserialize(StreamBuffer *sb) {
    // Lua stack should be empty.
    assert(lua_gettop(state_) == 0);

    // Get the eris data and convert it to a byte reader.
    int64_t eris_len = sb->read_int64();
    const char *eris_bytes = sb->read_bytes(eris_len);
    LuaByteReader bytereader(eris_bytes, eris_len);
    
    // Call eris with the permanents table and passing the snapshot as a lua_Reader.
    lua_pushstring(state_, "unpersist");
    lua_rawget(state_, LUA_REGISTRYINDEX);
    eris_undump(state_, bytereader.lua_reader, &bytereader);

    // Set up a stack frame manually.  Eris deserialization
    // should have left permstable and regcopy on the stack.
    assert(lua_gettop(state_) == 2);
    [[maybe_unused]] LuaSpecial permstable(1);
    LuaSpecial regcopy(2);
    lua_pushnil(state_);
    lua_pushnil(state_);
    LuaSpecial key(3);
    LuaSpecial value(4);
    LuaCoreStack LS(state_);
    assert(LS.istable(regcopy));

    // Copy the contents of the snapshot over to the registry.
    LS.set(key, LuaNil);
    while (LS.next(regcopy, key, value) != 0) {
        LS.rawset(LuaRegistry, key, value);
    }
    lua_settop(state_, 0);
}

// Snapshot and rollback can trivially be implemented on top of serialize and
// deserialize.  However, it's also possible to implement snapshot and rollback
// using an alternative technique:
//
// 1. When constructing the lua interpreter, use a custom memory allocator that
//    keeps track of all the memory blocks used by lua.
//
// 2. Snapshot simply copies all the memory blocks used by lua into a buffer.
//
// 3. Rollback restores lua's memory blocks back to their previous state.  This
//    has the effect of restoring lua's state.
//
// A proof-of-concept implementation of the memory-snapshotting design was
// created, and it worked. It is probably faster than using serialize and
// deserialize.
//
// Note: even if we implement this alternative design, we still need to keep
// serialize and deserialize around in order to implement the save-game
// functionality. So for now, we're sticking with this design, which doesn't
// require us to maintain any additional code.


