#include "luasnap.hpp"
#include "luastack.hpp"
#include <iostream>
#include <cassert>
#include <sstream>


LuaSnap::LuaSnap() {
    state_ = luaL_newstate();
    LuaStack LS(state_);

    // Create the persist table and the unpersist table.
    //
    // These tables need to contain all C functions.  Whenever
    // the source module inserts a C function into the lua environment,
    // it also needs to register the C function in these tables.
    //
    // I don't think anything else needs to go in these tables.
    //
    LS.rawset(LuaRegistry, "persist", LuaNewTable);
    LS.rawset(LuaRegistry, "unpersist", LuaNewTable);
}

LuaSnap::~LuaSnap() {
    std::cerr << "LuaSnap destructor not implemented yet" << std::endl;
}

bool LuaSnap::have_snapshot() const {
    return !snapshot_.empty();
}

static int oss_writer(lua_State *L, const void *p, size_t sz, void *ud) {
    std::ostringstream *oss = (std::ostringstream *)ud;
    oss->write((const char *)p, sz);
    return 0;
}

// Convert a C++ string into a lua_Reader.
class StringReader {
public:
    const std::string &str_;
    bool done_;
    StringReader(const std::string &s) : str_(s), done_(false) {}
    
    static const char *fn(lua_State *L, void *data, size_t *size) {
        StringReader *sr = (StringReader*)data;
        if (sr->done_) {
            *size = 0;
            return nullptr;
        } else {
            sr->done_ = true;
            *size = sr->str_.size();
            return sr->str_.c_str();
        }
    }
};

void LuaSnap::snapshot() {
    // Snapshot and the lua stack should both be empty.
    assert(snapshot_.empty());
    assert(lua_gettop(state_) == 0);

    // lua variables that we'll need.
    LuaVar key, value;
    LuaRet permstable, regcopy;
    LuaStack LS(state_, permstable, regcopy, key, value);

    // 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);

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

    // When we call 'LS.result', this should leave the
    // permstable and the regcopy on the stack.
    LS.result();
    assert(lua_gettop(state_) == 2);

    // Call eris to dump the state to an ostringstream,
    // then save the result.
    std::ostringstream oss;
    eris_dump(state_, oss_writer, (void *)&oss);
    snapshot_ = oss.str();
    std::cerr << "Eris dump is " << snapshot_.size() << " bytes." << std::endl;
    lua_settop(state_, 0);
}

void LuaSnap::rollback() {
    // Snapshot should have data, lua stack should be empty.
    assert(!snapshot_.empty());
    assert(lua_gettop(state_) == 0);

    // Call eris with the permanents table and passing the snapshot as a lua_Reader.
    lua_getfield(state_, LUA_REGISTRYINDEX, "unpersist");
    StringReader sr(snapshot_);
    eris_undump(state_, StringReader::fn, &sr);
    assert(lua_gettop(state_) == 2);

    // Set up a stack frame.
    LuaArg permstable, regcopy;
    LuaVar key, value;
    LuaStack LS(state_, permstable, regcopy, key, value);

    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);
    }
    LS.result();
    assert(lua_gettop(state_) == 0);
    snapshot_.clear();
}