#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"

// Read control.lst
//
// - trim all lines
// - remove blank lines
// - remove comment lines
//
util::stringvec read_control_lst(const std::string &path) {
    util::stringvec lines = util::get_file_lines(path);
    util::stringvec result;
    for (int i = 0; i < int(lines.size()); i++) {
        std::string trimmed = util::trim(lines[i]);
        if ((trimmed.size() > 0) && (trimmed[0] != '#')) {
            result.push_back(trimmed);
        }
    }
    return result;
}

LuaDefine(source_makeclass, "f") {
    LuaArg classname;
    LuaRet classtab;
    LuaStack LS(L, classname, classtab);
    LS.makeclass(classtab, classname);
    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);
    }
}

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

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

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

    // 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.rawget(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("lua/" + 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("lua/" + cfn);
        LS.rawset(info, "name", fn);
        LS.rawset(info, "fingerprint", new_fingerprint);
        LS.rawset(info, "code", ccode);
        if ((new_fingerprint == "")||(ccode == "")) {
            LS.rawset(info, "loadresult", "cannot read source file");
        } else {
            std::string chunk = "=" + cfn;
            luaL_loadbuffer(L, ccode.c_str(), ccode.size(), chunk.c_str());
            lua_replace(L, loadresult.index());
            LS.rawset(info, "loadresult", 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.rawget(sourcedb, LuaRegistry, "sourcedb");
    if (!LS.istable(sourcedb)) {
        LS.newtable(sourcedb);
    }

    // Read the list of filenames.
    util::stringvec filenames = read_control_lst("lua/control.lst");
    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.rawset(info, "sequence", seq);
        LS.rawset(newdb, fn, info);
    }

    // Store the new source db.
    LS.rawset(LuaRegistry, "sourcedb", newdb);
    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)) {
            LS.call(table_clear, 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());
            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.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)) {
            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_install_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, 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();
}