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

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


LuaDefine(source_makeclass, "f") {
    LuaArg classname;
    LuaVar action, gname;
    LuaRet classtab;
    LuaStack LS(L, classname, classtab, action, gname);

    LS.checkstring(classname);

    // Special case: if the classname is _G, return global env.
    LS.set(gname, "_G");
    if (LS.equal(classname, gname)) {
        LS.set(classtab, LuaGlobals);
        return LS.result();
    }
    
    // Get the classtab from the global environment.
    // Create it if it doesn't exist.
    LS.rawget(classtab, LuaGlobals, classname);
    if (LS.isnil(classtab)) {
        LS.newtable(classtab);
        LS.rawset(LuaGlobals, classname, classtab);
    }

    // If the name isn't bound to a table, abort.
    if (!LS.istable(classtab)) {
        luaL_error(L, "%s is not a class", LS.ckstring(classname).c_str());
    }

    // Repair the special fields.
    LS.setfield(classtab, "__index", classtab);
    LS.setfield(classtab, "__class", classname);

    // Repair the action table.
    LS.getfield(action, classtab, "action");
    if (!LS.istable(action)) {
        LS.setfield(classtab, "action", LuaNewTable);
    }

    return LS.result();
}

// Load the builtins.

static void load_builtin(lua_State *L, const char *name, lua_CFunction func) {
    lua_pushcfunction(L, func);
    lua_pushstring(L, name);
    lua_call(L, 1, 0);
}

static void source_install_builtins(lua_State *L) {
    LuaStack LS(L);
    load_builtin(L, "base", luaopen_base);
    load_builtin(L, "table", luaopen_table);
    load_builtin(L, "string", luaopen_string);
    load_builtin(L, "math", luaopen_math);
    load_builtin(L, "bit", luaopen_math);
    load_builtin(L, "debug", luaopen_debug);
}

void SourceDB::initialize(lua_State *L) {
    lua_state_ = L;
    LuaVar key, value, skey, svalue, snapshot, ssnapshot;
    LuaStack LS(L, snapshot, key, value, skey, svalue, ssnapshot);

    // Note: the global environment contains _G.  This routine
    // perceives this as a subtable, so it backs up the top level
    // as well.
    source_install_builtins(L);
    LS.newtable(snapshot);
    LS.set(key, LuaNil);
    while (LS.next(LuaGlobals, key, value) != 0) {
        if (LS.istable(value)) {
            LS.newtable(ssnapshot);
            LS.rawset(snapshot, key, ssnapshot);
            LS.set(skey, LuaNil);
            while (LS.next(value, skey, svalue) != 0) {
                if (LS.isfunction(svalue) || LS.isstring(svalue) || LS.isnumber(svalue)) {
                    LS.rawset(ssnapshot, skey, svalue);
                }
            }
        }
    }
    LS.setfield(LuaRegistry, "source_snapshot_builtins", snapshot);
    LS.result();
}

static int source_updatefile(lua_State *L) {
    LuaArg source, fn;
    LuaRet info;
    LuaVar fingerprint, null;
    LuaStack LS(L, source, fn, info, fingerprint, null);

    // Get the existing info table from the source DB.
    if (LS.istable(source)) {
        LS.rawget(info, source, fn);
        if (!LS.istable(info)) {
            LS.newtable(info);
        }
    } else {
        LS.newtable(info);
    }

    // If the file modification is wrong, update
    // these fields: code, fingerprint, closure, error
    // Otherwise, update nothing.
    std::string cfn = LS.ckstring(fn);
    LS.getfield(fingerprint, info, "fingerprint");
    std::string old_fingerprint;
    if (LS.isstring(fingerprint)) {
        old_fingerprint = LS.ckstring(fingerprint);
    }
    // std::cerr << "Probing " << cfn << std::endl;
    std::string new_fingerprint = util::get_file_fingerprint("syslua/" + cfn);
    LS.set(null, LuaNil);
    if ((old_fingerprint == "") || (old_fingerprint != new_fingerprint)) {
        std::cerr << "Rereading " << cfn << std::endl;
        std::string ccode = util::get_file_contents("syslua/" + cfn);
        LS.setfield(info, "name", fn);
        LS.setfield(info, "fingerprint", new_fingerprint);
        LS.setfield(info, "code", ccode);
        if ((new_fingerprint == "")||(ccode == "")) {
            LS.setfield(info, "loadresult", "cannot read source file");
        } else {
            std::string chunk = "=" + cfn;
            luaL_loadbuffer(L, ccode.c_str(), ccode.size(), chunk.c_str());
            lua_setfield(L, info.index(), "loadresult");
        }
    }
    return LS.result();
}

void SourceDB::update() {
    lua_State *L = lua_state_;
    LuaVar sourcedb, newdb, info, fn, seq;
    LuaStack LS(L, newdb, sourcedb, info, fn, seq);

    // Get the (old) source database.
    LS.getfield(sourcedb, LuaRegistry, "sourcedb");
    if (!LS.istable(sourcedb)) {
        LS.newtable(sourcedb);
    }

    // Read the list of filenames.
    std::string ctrl = "syslua/control.lst";
    util::stringvec filenames = util::trim_and_uncomment(util::read_lines(ctrl));
    if (filenames.empty()) {
        luaL_error(L, "cannot read source database control.lst");
    }

    // Process the files one by one.
    LS.newtable(newdb);
    for (int i = 0; i < int(filenames.size()); i++) {
        LS.set(fn, filenames[i]);

        // Call source_updatefile to get the updated info for one file.
        LS.call(info, source_updatefile, sourcedb, fn);

        // Insert the sequence number and put finalized info into the new database.
        LS.set(seq, i + 1);
        LS.setfield(info, "sequence", seq);
        LS.rawset(newdb, fn, info);
    }

    // Store the new source db.
    LS.setfield(LuaRegistry, "sourcedb", newdb);
    LS.result();
}

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

    LS.setfield(LuaGlobals, "_G", LuaNil);
    LS.set(classname, LuaNil);
    while (LS.next(LuaGlobals, classname, classtab) != 0) {
        if (LS.istable(classtab)) {
            bool keep_action = false;
            LS.getfield(action, classtab, "action");
            if (LS.istable(action)) {
                LS.call(table_clear, action);
                keep_action = true;
            }
            LS.call(table_clear, classtab);
            if (keep_action) {
                LS.setfield(classtab, "action", action);
            }
        } else {
            LS.rawset(LuaGlobals, classname, LuaNil);
        }
    }
    LS.setfield(LuaGlobals, "_G", LuaGlobals);
    LS.result();
}

// Restore the lua builtins from the backup snapshot.
//
static void source_restore_builtins(lua_State *L) {
    LuaVar snapshot, key, value, skey, svalue, subglobal;
    LuaStack LS(L, snapshot, key, value, skey, svalue, subglobal);

    LS.getfield(snapshot, LuaRegistry, "source_snapshot_builtins");
    LS.setfield(LuaGlobals, "_G", LuaGlobals);
    LS.set(key, LuaNil);
    while (LS.next(snapshot, key, value) != 0) {
        LS.checktable(value);
        LS.call(subglobal, source_makeclass, key);
        LS.set(skey, LuaNil);
        while (LS.next(value, skey, svalue) != 0) {
            LS.rawset(subglobal, skey, svalue);
        }
    }
    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());
            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) {
    LuaRet errors;
    LuaVar sourcedb, key, info, seq, closure, err;
    LuaStack LS(L, sourcedb, errors, key, info, seq, closure, err);

    // Get the source database.
    LS.getfield(sourcedb, LuaRegistry, "sourcedb");
    if (LS.type(sourcedb) != LUA_TTABLE) {
        LS.newtable(sourcedb);
        LS.setfield(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.getfield(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.getfield(closure, info, "loadresult");

        // If there's already an error in the sourcedb, collect it.
        if (!LS.isfunction(closure)) {
            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() {
    lua_State *L = lua_state_;
    LuaVar errs;
    LuaStack LS(L, errs);
    source_clear_globals(L);
    source_restore_builtins(L);
    source_load_cfunctions(L);
    source_load_lfunctions(L);
    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;
    LuaStack LS(L, unittests, name, func, err);

    LS.getfield(unittests, LuaGlobals, "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();
}