#include <string>
#include <vector>
#include <map>
#include <set>
#include <algorithm>
#include <sstream>
#include <fstream>
#include <iostream>
#include "util.hpp"

#include "luastack.hpp"
#include "traceback.hpp"
#include "table.hpp"
#include "source.hpp"
#include "luasnap.hpp"

LuaDefine(source_makeclass, "f") {
    LuaArg classname;
    LuaRet classtab;
    LuaStack LS(L, classname, classtab);
    LS.makeclass(classtab, classname);
    return LS.result();
}

LuaDefine(source_classname, "f") {
    LuaArg table;
    LuaRet result;
    LuaStack LS(L, table, result);
    std::string rstr = LS.classname(table);
    if (rstr == "") {
        LS.set(result, LuaNil);
    } else {
        LS.set(result, rstr);
    }
    return LS.result();
}

LuaDefine(source_maketangible, "f") {
    LuaArg classname;
    LuaRet classtab;
    LuaVar subtab;
    LuaStack LS(L, classname, classtab, subtab);
    LS.makeclass(classtab, classname);
    LS.makesubtable(subtab, classtab, "action");
    return LS.result();
}


// the 'luaopen' function creates a new table.
//
// We don't want to create new tables during reload, so we call luaopen,
// then we copy the contents of the new table into the existing 'makeclass'
// table.
//
static void load_builtin_class(lua_State *L, const char *name, lua_CFunction func) {
    LuaVar sourcetab, classtab, key, value;
    LuaStack LS(L, sourcetab, classtab, key, value);
    LS.makeclass(classtab, name);
    func(L);
    lua_replace(L, sourcetab.index());
    LS.set(key, LuaNil);
    while (LS.next(sourcetab, key, value) != 0) {
        LS.rawset(classtab, key, value);
    }
    LS.result();
}

static void source_install_builtins(lua_State *L) {
    LuaVar nullstring, stringclass, globtab;
    LuaStack LS(L, nullstring, stringclass, globtab);
    luaopen_base(L);
    load_builtin_class(L, "table", luaopen_table);
    load_builtin_class(L, "string", luaopen_string);
    load_builtin_class(L, "math", luaopen_math);
    load_builtin_class(L, "debug", luaopen_debug);
    load_builtin_class(L, "coroutine", luaopen_coroutine);

    // Nuke a few of the builtin functions for sandboxing reasons.
    LS.getglobaltable(globtab);
    LS.rawset(globtab, "loadfile", LuaNil);
    
    // Set the metatable for strings.
    // Normally, this would be done by luaopen_string, but we're
    // messing with the tables so we have to redo it.
    LS.makeclass(stringclass, "string");
    LS.set(nullstring, "");
    LS.setmetatable(nullstring, stringclass);

    LS.result();
}

static void get_info_table(LuaStack &LS, LuaSlot db, LuaSlot info, const std::string &fn) {
    LS.rawget(info, db, fn);
    if (!LS.istable(info)) {
        LS.set(info, LuaNewTable);
        LS.rawset(db, fn, info);
    }
    LS.rawset(info, "name", fn);
}

static void calculate_loadresult(LuaStack &LS0, LuaSlot info, const std::string &fn, const std::string &code) {
    LuaVar loadresult;
    LuaStack LS(LS0.state(), loadresult);
    if (code == "") {
        LS.rawset(info, "loadresult", "missing or empty source file");
    } else {
        std::string chunk = "=" + fn;
        luaL_loadbuffer(LS.state(), code.c_str(), code.size(), chunk.c_str());
        lua_replace(LS.state(), loadresult.index());
        LS.rawset(info, "loadresult", loadresult);
    }
}


void SourceDB::diff(const SourceDB &auth, StreamBuffer *sb) {
    LuaVar sdb, sfn, sinfo, shash, sseq;
    LuaVar mdb, mfn, minfo, mhash, mseq, mcode;
    LuaStack SLS(lua_state_, sdb, sfn, sinfo, shash, sseq);
    LuaStack MLS(auth.lua_state_, mdb, mfn, minfo, mhash, mseq, mcode);
    sb->write_int32(0);
    int wc_after = sb->total_writes();
    int nupdates = 0;
    // Fetch the two source databases.
    SLS.rawget(sdb, LuaRegistry, "sourcedb");
    MLS.rawget(mdb, LuaRegistry, "sourcedb");
    // Loop over the master database.
    MLS.set(mfn, LuaNil);
    while (MLS.next(mdb, mfn, minfo) != 0) {
        std::string fn = MLS.ckstring(mfn);
        assert(MLS.istable(minfo));
        MLS.rawget(mseq, minfo, "sequence");
        MLS.rawget(mhash, minfo, "hash");
        bool diffhash = true;
        bool diffseq = true;
        SLS.set(sfn, fn);
        SLS.rawget(sinfo, sdb, sfn);
        if (SLS.istable(sinfo)) {
            SLS.rawget(shash, sinfo, "hash");
            SLS.rawget(sseq, sinfo, "sequence");
            diffhash = MLS.ckstring(mhash) != SLS.ckstring(shash);
            diffseq = MLS.ckinteger(mseq) != SLS.ckinteger(sseq);
        }
        if (diffhash || diffseq) {
            sb->write_string(MLS.ckstring(mfn));
            sb->write_int32(MLS.ckinteger(mseq));
            if (diffhash) {
                MLS.rawget(mcode, minfo, "code");
                sb->write_string(MLS.ckstring(mcode));
            } else {
                sb->write_string("\001");
            }
            nupdates += 1;
        }
    }
    // Loop over synch database.
    SLS.set(sfn, LuaNil);
    while (SLS.next(sdb, sfn, sinfo) != 0) {
        std::string fn = SLS.ckstring(sfn);
        assert(SLS.istable(sinfo));
        MLS.set(mfn, fn);
        MLS.rawget(minfo, mdb, mfn);
        if (!MLS.istable(minfo)) {
            sb->write_string(SLS.ckstring(sfn));
            sb->write_int32(-1);
            sb->write_string("");
            nupdates += 1;
        }
    }
    sb->overwrite_int32(wc_after, nupdates);
    MLS.result();
    SLS.result();
}

void SourceDB::patch(StreamBuffer *sb) {
    lua_State *L = lua_state_;
    LuaVar db, info;
    LuaStack LS(L, db, info);
    LS.rawget(db, LuaRegistry, "sourcedb");
    int nupdates = sb->read_int32();
    for (int i = 0; i < nupdates; i++) {
        std::string fn = sb->read_string();
        int sequence = sb->read_int32();
        std::string code = sb->read_string();
        if (sequence < 0) {
            LS.rawset(db, fn, LuaNil);
        } else {
            get_info_table(LS, db, info, fn);
            LS.rawset(info, "sequence", sequence);
            if (code != "\001") {
                LS.rawset(info, "code", code);
                LS.rawset(info, "hash", util::hash_to_hex(util::hash_string(code)));
                calculate_loadresult(LS, info, fn, code);
            }
        }
    }
    LS.result();
    if (nupdates > 0) rebuild(false);
}

void SourceDB::set(const std::string &fn, const std::string &code, int sequence) {
    lua_State *L = lua_state_;
    LuaVar db, info;
    LuaStack LS(L, db, info);
    LS.rawget(db, LuaRegistry, "sourcedb");
    get_info_table(LS, db, info, fn);
    LS.rawset(info, "sequence", sequence);
    LS.rawset(info, "code", code);
    LS.rawset(info, "fingerprint", "");
    LS.rawset(info, "hash", util::hash_to_hex(util::hash_string(code)));
    calculate_loadresult(LS, info, fn, code);
    LS.result();
}

std::string SourceDB::get(const std::string &fn) {
    lua_State *L = lua_state_;
    LuaVar db, info, code, sequence, loadresult;
    LuaStack LS(L, db, info, code, sequence, loadresult);
    LS.rawget(db, LuaRegistry, "sourcedb");
    LS.rawget(info, db, fn);
    if (!LS.istable(info)) {
        LS.result();
        return "<nonexistent>";
    }
    LS.rawget(code, info, "code");
    LS.rawget(sequence, info, "sequence");
    LS.rawget(loadresult, info, "loadresult");
    std::string ccode = LS.ckstring(code);
    int seqno = LS.ckint(sequence);
    std::string cloadresult;
    if (LS.isfunction(loadresult)) {
        cloadresult = "<function>";
    } else {
        cloadresult = LS.ckstring(loadresult);
    }
    std::ostringstream oss;
    oss << seqno << ":" << ccode << ":" << cloadresult;
    LS.result();
    return oss.str(); 
}

void SourceDB::update(const util::LuaSourceVec &source) {
    lua_State *L = lua_state_;
    LuaVar sourcedb, info;
    LuaStack LS(L, sourcedb, info);

    // Get and clear the source database.
    LS.rawget(sourcedb, LuaRegistry, "sourcedb");
    if (!LS.istable(sourcedb)) {
        LS.newtable(sourcedb);
        LS.rawset(LuaRegistry, "sourcedb", sourcedb);
    }
    LS.cleartable(sourcedb);

    for (int i = 0; i < int(source.size()); i++) {
        const std::string &file = source[i].first;
        const std::string &code = source[i].second;
        std::cerr << "Compiling " << file << std::endl;
        LS.newtable(info);
        LS.rawset(info, "name", file);
        LS.rawset(info, "code", code);
        LS.rawset(info, "hash", util::hash_to_hex(util::hash_string(code)));
        LS.rawset(info, "sequence", i + 1);
        calculate_loadresult(LS, info, file, code);
        LS.rawset(sourcedb, file, info);
    }
    LS.result();
}

// Delete everything from the global environment except
// the class tables.
//
static void source_clear_globals(lua_State *L) {
    LuaVar classname, classtab, key, globtab;
    LuaStack LS(L, classname, classtab, key, globtab);

    LS.getglobaltable(globtab);
    LS.rawset(globtab, "_G", LuaNil);
    LS.set(classname, LuaNil);
    while (LS.next(globtab, classname, classtab) != 0) {
        if (LS.istable(classtab)) {
            table_clear(LS, classtab);
        } else {
            LS.rawset(globtab, classname, LuaNil);
        }
    }
    LS.rawset(globtab, "_G", globtab);
    LS.result();
}


// Load all the 'LuaDefine' C functions into the lua state.
//
static void source_load_cfunctions(lua_State *L) {
    auto regs = LuaFunctionReg::all();
    for (const LuaFunctionReg *r : regs) {
        const std::string &name = r->get_name();
        size_t upos = name.find('_');
        std::string classname;
        std::string funcname;
        if (upos == std::string::npos) {
            continue;
        } else {
            funcname = name.substr(upos + 1);
            classname = name.substr(0, upos);
        }
        lua_CFunction func = r->get_func();
        std::string mode = r->get_mode();
        if (mode.find('c') != std::string::npos) { // Insert into class
            lua_pushlstring(L, classname.c_str(), classname.size());
            lfn_source_makeclass(L);
            lua_pushcfunction(L, func);
            lua_setfield(L, -2, funcname.c_str());
        }
        if (mode.find('f') != std::string::npos) { // Make global function
            lua_pushcfunction(L, func);
            lua_setglobal(L, funcname.c_str());
        }
    }
}

// Run all the closures from the source database.
//
static void source_load_lfunctions(lua_State *L, bool print_errors) {
    LuaRet errors;
    LuaVar sourcedb, key, info, seq, closure, err;
    LuaStack LS(L, sourcedb, errors, key, info, seq, closure, err);

    // Get the source database.
    LS.rawget(sourcedb, LuaRegistry, "sourcedb");
    if (LS.type(sourcedb) != LUA_TTABLE) {
        LS.newtable(sourcedb);
        LS.rawset(LuaRegistry, "sourcedb", sourcedb);
    }

    // Sort the keys by sequence number.
    std::map<int, std::string> indices;
    LS.set(key, LuaNil);
    while (LS.next(sourcedb, key, info) != 0) {
        LS.rawget(seq, info, "sequence");
        indices[LS.ckinteger(seq)] = LS.ckstring(key);
    }

    // Now call the closures in the proper order.
    std::stringstream errss;
    for (const auto &p : indices) {
        LS.rawget(info, sourcedb, p.second);
        LS.rawget(closure, info, "loadresult");

        // If there's already an error in the sourcedb, collect it.
        if (!LS.isfunction(closure)) {
            if (print_errors) {
                errss << LS.ckstring(closure) << "\n";
            }
            continue;
        }
        
        // Call the closure.  If there's an error, collect it.
        lua_pushvalue(L, closure.index());
        if (traceback_pcall(L, 0, 0) != 0) {
            lua_replace(L, err.index());
            errss << LS.ckstring(err);
        }
    }
    LS.set(errors, errss.str());
    LS.result();
}

void SourceDB::rebuild(bool print_errors) {
    lua_State *L = lua_state_;
    LuaVar errs;
    LuaStack LS(L, errs);
    source_clear_globals(L);
    source_install_builtins(L);
    source_load_cfunctions(L);
    source_load_lfunctions(L, print_errors);
    lua_replace(L, errs.index());
    std::string errstr = LS.ckstring(errs);
    std::cerr << errstr;
    LS.result();
}

void SourceDB::run_unittests() {
    lua_State *L = lua_state_;
    LuaVar unittests, name, func, err, globtab;
    LuaStack LS(L, unittests, name, func, err, globtab);

    LS.getglobaltable(globtab);
    LS.rawget(unittests, globtab, "unittests");

    // Sort the unit test names.
    std::set<std::string> names;
    LS.set(name, LuaNil);
    while (LS.next(unittests, name, func) != 0) {
        if (LS.isfunction(func) && LS.isstring(name)) {
            names.insert(LS.ckstring(name));
        }
    }

    // Run the functions in order
    bool any = false;
    for (const std::string &name : names) {
        std::cerr << "Running unittests." << name << std::endl;
        LS.rawget(func, unittests, name);

        lua_pushvalue(L, func.index());
        if (traceback_pcall(L, 0, 0) != 0) {
            lua_replace(L, err.index());
            std::cerr << LS.ckstring(err);
            any = true;
        }
    }

    if (any) {
        exit(1);
    }
    LS.result();
}

void SourceDB::init(lua_State *L) {
    lua_state_ = L;
    LuaVar globtab, persist, unpersist, classname, classtab, funcname, funcp, rawfunc;
    LuaStack LS(L, globtab, persist, unpersist, classname, classtab, funcname, funcp, rawfunc);

    // We need to register all C functions with the eris permanents tables.
    source_clear_globals(L);
    source_install_builtins(L);
    source_load_cfunctions(L);
    LS.getglobaltable(globtab);
    LS.rawget(persist, LuaRegistry, "persist");
    LS.rawget(unpersist, LuaRegistry, "unpersist");
    LS.rawset(LuaRegistry, "sourcedb", LuaNewTable);
    LS.set(classname, LuaNil);
    while (LS.next(globtab, classname, classtab) != 0) {
        if (LS.isstring(classname) && LS.istable(classtab)) {
            LS.set(funcname, LuaNil);
            while (LS.next(classtab, funcname, funcp) != 0) {
                if (LS.isstring(funcname) && LS.iscfunction(funcp)) {
                    std::string full = "cfunc:";
                    full += LS.ckstring(classname);
                    full += ".";
                    full += LS.ckstring(funcname);
                    lua_pushcfunction(L, lua_tocfunction(L, funcp.index()));
                    lua_replace(L, rawfunc.index());
                    LS.rawset(persist, rawfunc, full);
                    LS.rawset(unpersist, full, rawfunc);
                }
            }
        }
    }
}

// These should go away eventually.  They're for debugging.
LuaDefine(coroutine_setnextid, "c") {
    LuaArg co, lid;
    LuaStack LS(L, co, lid);
    lua_State *CO = LS.ckthread(co);
    lua_Number id = LS.ckinteger(lid);
    lua_setnextid(CO, id);
    return LS.result();
}

LuaDefine(coroutine_getnextid, "c") {
    LuaArg co;
    LuaRet lid;
    LuaStack LS(L, co, lid);
    lua_State *CO = LS.ckthread(co);
    LS.set(lid, lua_getnextid(CO));
    return LS.result();
}

LuaDefine(unittests_sourcedb, "c") {
    LuaSnap msnap;
    LuaSnap ssnap;
    SourceDB mdb;
    SourceDB sdb;
    mdb.init(msnap.state());
    sdb.init(ssnap.state());
    StreamBuffer sb;

    // Create a code database using 'set'.
    mdb.set("foo", "function foo() print('foo') end", 1);
    mdb.set("bar", "function bar() print('bar') end", 2);
    mdb.set("zoo", "funcjdshja mxooso yowza!", 3);
    
    // Verify that get works.
    LuaAssertStrEq(L, mdb.get("foo"), "1:function foo() print('foo') end:<function>");
    LuaAssertStrEq(L, mdb.get("bar"), "2:function bar() print('bar') end:<function>");
    LuaAssertStrEq(L, mdb.get("zoo"), "3:funcjdshja mxooso yowza!:zoo:1: syntax error near 'mxooso'");
    LuaAssertStrEq(L, mdb.get("baz"), "<nonexistent>");

    // Difference transmit.
    sdb.diff(mdb, &sb);

    // There should still be nothing in the sdb.
    LuaAssertStrEq(L, sdb.get("foo"), "<nonexistent>");
    LuaAssertStrEq(L, sdb.get("bar"), "<nonexistent>");
    LuaAssertStrEq(L, sdb.get("zoo"), "<nonexistent>");
    LuaAssertStrEq(L, sdb.get("baz"), "<nonexistent>");

    // Apply the diffs.
    sdb.patch(&sb);

    // Everything should now be copied to sdb.
    LuaAssertStrEq(L, sdb.get("foo"), "1:function foo() print('foo') end:<function>");
    LuaAssertStrEq(L, sdb.get("bar"), "2:function bar() print('bar') end:<function>");
    LuaAssertStrEq(L, sdb.get("zoo"), "3:funcjdshja mxooso yowza!:zoo:1: syntax error near 'mxooso'");
    LuaAssertStrEq(L, sdb.get("baz"), "<nonexistent>");

    // Make a single change to the code.
    mdb.set("bar", "function bar1() print('bar1') end", 2);

    // Diff and patch
    sdb.diff(mdb, &sb);
    sdb.patch(&sb);

    // Verify that it's been updated.
    LuaAssertStrEq(L, sdb.get("foo"), "1:function foo() print('foo') end:<function>");
    LuaAssertStrEq(L, sdb.get("bar"), "2:function bar1() print('bar1') end:<function>");
    LuaAssertStrEq(L, sdb.get("zoo"), "3:funcjdshja mxooso yowza!:zoo:1: syntax error near 'mxooso'");
    LuaAssertStrEq(L, sdb.get("baz"), "<nonexistent>");

    // Make a single change to just a sequence number.
    mdb.set("foo", "function foo() print('foo') end", 6);

    // Diff
    sdb.diff(mdb, &sb);

    // The only thing we've updated is a single sequence number.  Make
    // sure the number of bytes transmitted is reasonable.
    //    4 bytes for the count of changes
    //    4 bytes for "foo"
    //    4 bytes for the sequence number
    //    2 bytes for unmodified code
    // This verifies that the hash comparisons are working.
    LuaAssert(L, sb.fill() == 14);

    // Patch.
    sdb.patch(&sb);

    // Verify that it's been updated.
    LuaAssertStrEq(L, sdb.get("foo"), "6:function foo() print('foo') end:<function>");
    LuaAssertStrEq(L, sdb.get("bar"), "2:function bar1() print('bar1') end:<function>");
    LuaAssertStrEq(L, sdb.get("zoo"), "3:funcjdshja mxooso yowza!:zoo:1: syntax error near 'mxooso'");
    LuaAssertStrEq(L, sdb.get("baz"), "<nonexistent>");

    return 0;
}