#include "pprint.hpp"
#include "util.hpp"
#include "table.hpp"
#include <ostream>
#include <iostream>


void atomic_print(LuaStack &LS, LuaSlot val, std::ostream *os) {
    switch (LS.type(val)) {
    case LUA_TNIL:
        (*os) << "nil";
        return;
    case LUA_TSTRING:
        // TODO: this could be more efficient.
        (*os) << LS.ckstring(val);
        return;
    case LUA_TNUMBER: {
        double value = LS.cknumber(val);
        int64_t ivalue = int64_t(value);
        if (double(ivalue) == value) {
            (*os) << ivalue;
        } else {
            (*os) << value;
        }
        return;
    }
    case LUA_TBOOLEAN:
        (*os) << (LS.ckboolean(val) ? "true" : "false");
        return;
    case LUA_TTABLE:
        (*os) << "table";
        return;
    default:
        (*os) << "unknown";
        return;
    }
}

static bool string_quote(LuaStack &LS, LuaSlot val, std::ostream *os) {
    switch (LS.type(val)) {
    case LUA_TNIL:
        (*os) << "nil";
        return true;
    case LUA_TSTRING:
        util::quote_string(LS.ckstring(val), os);
        return true;
    case LUA_TNUMBER: {
        double value = LS.cknumber(val);
        int64_t ivalue = int64_t(value);
        if (double(ivalue) == value) {
            (*os) << ivalue;
        } else {
            (*os) << value;
        }
        return true;
    }
    case LUA_TBOOLEAN:
        (*os) << (LS.ckboolean(val) ? "true" : "false");
        return true;
    default:
        return false;
    }
}


// Find tables recursively.
//
// Builds a table (tabcount) whose keys are tables.  If a table
// is visited exactly once, then tabcount[t]=0.  If a table is
// visited more than once, then tabcount[t]=-1.
//
static void findtables(LuaStack &LS0, LuaSlot root, LuaSlot tabcount) {
    lua_State *L = LS0.state();
    LuaVar key, val, tab, count;
    LuaStack LS(L, key, val, tab, count);

    LS.newtable(tabcount);
    int top = lua_gettop(L);
    if (LS.istable(root)) {
        lua_pushvalue(L, root.index());
    }
    while (lua_gettop(L) > top) {
        lua_checkstack(L, 10);
        lua_replace(L, tab.index());
        LS.rawget(count, tabcount, tab);
        if (LS.isnil(count)) {
            LS.rawset(tabcount, tab, 0);
            LS.set(key, LuaNil);
            while (LS.next(tab, key, val)) {
                if (LS.istable(key)) {
                    lua_pushvalue(L, key.index());
                }
                if (LS.istable(val)) {
                    lua_pushvalue(L, val.index());
                }
            }
            LS.getmetatable(val, tab);
            if (LS.istable(val)) {
                lua_pushvalue(L, val.index());
            }
        } else {
            LS.rawset(tabcount, tab, -1);
        }
    }
    LS.result();
}

LuaDefine(table_findtables, "c") {
    LuaArg root;
    LuaRet tabcount;
    LuaStack LS(L, root, tabcount);
    findtables(LS, root, tabcount);
    return LS.result();
}

struct Inspector {
    lua_State *L;
    LuaVar ids;
    int nextid;
    bool indent;
    int maxlen;
    bool anyindent;
    std::ostream *stream;
};

static void tabify(Inspector &insp, int level) {
    if (insp.indent) {
        (*insp.stream) << std::endl;
        for (int i = 0; i < level; i++) {
            (*insp.stream) << "  ";
        }
        insp.anyindent = true;
    } else {
        (*insp.stream) << " ";
    }
}

static void pprint_r(Inspector &insp, int level, LuaSlot root) {
    lua_checkstack(insp.L, 20);
    LuaVar idv, pairs, key, val;
    LuaStack LS(insp.L, idv, pairs, key, val);

    // If it's a simple type, print it quoted.
    if (string_quote(LS, root, insp.stream)) {
        LS.result();
        return;
    }

    // If it's not a table, just print the typename.
    int t = LS.type(root);
    if (t != LUA_TTABLE) {
        (*insp.stream) << "<" << lua_typename(insp.L, t) << ">";
        LS.result();
        return;
    }

    // Determine the table's ID, allocating an ID if necessary.
    LS.rawget(idv, insp.ids, root);
    int id = LS.ckint(idv);
    bool new_id = false;
    if (id < 0) {
        new_id = true;
        id = insp.nextid++;
        LS.rawset(insp.ids, root, id);
    }
    
    // Print the table's name, if any.
    bool is_class = false;
    bool is_tangible = false;
    std::string cname = LS.classname(root);
    if (cname != "") {
        is_class = true;
        (*insp.stream) << "<class " << cname << ">";
    } else {
        int64_t tid = LS.tanid(root);
        if (tid > 0) {
            is_tangible = true;
            (*insp.stream) << "<tangible " << tid << ">";
        } else if (id > 0) {
            (*insp.stream) << "<table " << id << ">";
        }
    }

    // If this is a class, and we're not at the top level, truncate.
    if ((is_class || is_tangible) && (level > 0)) {
        LS.result();
        return;
    }

    // If this is a table we've already printed, truncate it.
    if ((id > 0) && (!new_id)) {
        if (lua_nkeys(insp.L, root.index())==0) {
            (*insp.stream) << "{}";
        } else {
            (*insp.stream) << "{...}";
        }
        LS.result();
        return;
    }

    // State variables.
    int nextseq = 1;
    bool needcomma = false;
    bool multiline = false;

    // Open the brackets.
    (*insp.stream) << "{";

    // Output the table keys.
    table_getpairs(LS, root, pairs, true);
    for (int i = 2; ; i+=2) {
        LS.rawget(key, pairs, i);
        if (LS.isnil(key)) break;
        LS.rawget(val, pairs, i+1);
        if (needcomma) (*insp.stream) << ",";
        needcomma = true;
        if (LS.isnumber(key) && (LS.ckint(key) == nextseq)) {
            (*insp.stream) << " ";
            pprint_r(insp, level + 1, val);
            nextseq = nextseq + 1;
        } else {
            multiline = true;
            tabify(insp, level + 1);
            if (LS.isstring(key) && util::is_identifier(LS.ckstring(key))) {
                (*insp.stream) << LS.ckstring(key);
            } else {
                (*insp.stream) << "[";
                pprint_r(insp, level + 1, key);
                (*insp.stream) << "]";
            }
            if (insp.indent) {
                (*insp.stream) << " = ";
            } else {
                (*insp.stream) << "=";
            }
            pprint_r(insp, level + 1, val);
        }
    }

    // Output the metatable.
    LS.getmetatable(val, root);
    if (LS.istable(val)) {
        multiline = true;
        if (needcomma) (*insp.stream) << ",";
        needcomma = true;
        tabify(insp, level + 1);
        (*insp.stream) << "<meta> = ";
        pprint_r(insp, level + 1, val);
    }
    
    // Close the brackets.
    if (multiline) {
        tabify(insp, level);
    } else if (nextseq > 1) {
        (*insp.stream) << " ";
    }
    (*insp.stream) << "}";
    
    LS.result();
}

void pprint(LuaStack &LS0, LuaSlot root, bool indent, std::ostream *os) {
    Inspector insp;
    LuaStack LS(LS0.state(), insp.ids);
    findtables(LS, root, insp.ids);
    insp.L = LS0.state();
    insp.nextid = 1;
    insp.indent = indent;
    insp.anyindent = false;
    insp.stream = os;
    pprint_r(insp, 0, root);
    LS.result();
}

LuaDefine(string_isidentifier, "c") {
    LuaArg str;
    LuaRet result;
    LuaStack LS(L, str, result);
    if (LS.isstring(str)) {
        std::string s = LS.ckstring(str);
        LS.set(result, util::is_identifier(s));
    } else {
        LS.set(result, false);
    }
    return LS.result();
}

LuaDefine(string_print, "c") {
    LuaArg val;
    LuaRet result;
    LuaStack LS(L, val, result);
    std::ostringstream oss;
    atomic_print(LS, val, &oss);
    LS.set(result, oss.str());
    return LS.result();
}

LuaDefine(string_pprint, "c") {
    LuaArg root, indent;
    LuaRet result;
    LuaStack LS(L, root, indent, result);
    bool ind = LS.ckboolean(indent);
    std::ostringstream oss;
    pprint(LS, root, ind, &oss);
    LS.set(result, oss.str());
    return LS.result();
}

LuaDefine(string_tostring, "f") {
    LuaArg val;
    LuaRet result;
    LuaStack LS(L, val, result);
    std::ostringstream oss;
    atomic_print(LS, val, &oss);
    LS.set(result, oss.str());
    return LS.result();
}