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

#include <cmath>

class PrintMachine {
public:
    LuaVar tabchpos_;
    LuaExtStack LS_;
    int next_id_;
    bool indent_;
    std::ostream *output_;
    eng::map<int, int> chpos_to_tabnum_;

    void atomic_print(int xtype, LuaSlot val, bool quote) {
        switch (xtype) {
        case LUA_TNIL: {
            (*output_) << "nil";
            return;
        }
        case LUA_TSTRING: {
            if (quote) {
                util::quote_string(LS_.ckstring(val), output_);
            } else {
                // TODO: this could be more efficient.
                (*output_) << LS_.ckstring(val);
            }   
            return;
        }
        case LUA_TNUMBER: {
            double value = LS_.cknumber(val);
            if (std::isnan(value)) {
                (*output_) << "nan";
            } else {
                int64_t ivalue = int64_t(value);
                if (double(ivalue) == value) {
                    (*output_) << ivalue;
                } else {
                    (*output_) << value;
                }
            }
            return;
        }
        case LUA_TBOOLEAN: {
            (*output_) << (LS_.ckboolean(val) ? "true" : "false");
            return;
        }
        case LUA_TFUNCTION: {
            (*output_) << "<function>";
            return;
        }
        case LUA_TTHREAD: {
            (*output_) << "<thread>";
            return;
        }
        case LUA_TLIGHTUSERDATA: {
            LuaToken token = LS_.cktoken(val);
            (*output_) << "[" << token.str() << "]";
            return;
        }
        case LUA_TT_GENERAL: {
            (*output_) << "<table>";
            return;
        }
        case LUA_TT_TANGIBLE: {
            (*output_) << "<tangible " << LS_.tanid(val) << ">";
            return;
        }
        case LUA_TT_CLASS: {
            (*output_) << "<class " << LS_.classname(val) << ">";
            return;
        }
        case LUA_TT_GLOBALENV: {
            (*output_) << "<global-env>";
            return;
        }
        case LUA_TT_TANGIBLEMETA: {
            (*output_) << "<tangible-metatable>";
            return;
        }
        default: {
            (*output_) << "<unknown type #" << xtype << ">";
            return;
        }}
    }

    void tabify(int level) {
        if (indent_) {
            (*output_) << std::endl;
            for (int i = 0; i < level; i++) {
                (*output_) << "  ";
            }
        } else {
            (*output_) << " ";
        }
    }

    void pprint_r(int level, bool expand, LuaSlot value) {
        lua_State *L = LS_.state();
        lua_checkstack(L, 20);
        LuaVar loffset, pairs, key, val, lchpos, nextseq;
        LuaExtStack LS(L, loffset, pairs, key, val, lchpos, nextseq);

        // Determine the extended type of the object. If the
        // expand flag is true, try to coerce it to a general table.
        int xtype = LS.xtype(value);

        // Print the atomic portion.
        if (xtype != LUA_TT_GENERAL) {
            atomic_print(xtype, value, true);
            if ((xtype < LUA_TT_GENERAL) || (!expand)) {
                return;
            }
        }

        // Find out whether it has a table number.  If necessary,
        // assign one.
        int tabnum = 0;
        LS.rawget(lchpos, tabchpos_, value);
        auto chpos = LS.tryint(lchpos);
        if (!chpos) {
            // First time. Record the character position where the
            // table first appears in the output stream.
            LS.rawset(tabchpos_, value, int((*output_).tellp()));
        } else {
            tabnum = chpos_to_tabnum_[*chpos];
            if (tabnum == 0) {
                // Second time. The table is already in the output,
                // but it hasn't been assigned a number.  Assign one.
                tabnum = next_id_++;
                chpos_to_tabnum_[*chpos] = tabnum;
            }
        }
        
        // If the table has a number, that means we're visiting
        // it for the second time.  Just print an abbreviated version
        // and return.
        if (tabnum > 0) {
            (*output_) << "<table " << tabnum << ">";
            if (lua_nkeys(L, value.index())==0) {
                (*output_) << "{}";
            } else {
                (*output_) << "{...}";
            }
            return;
        }

        // State variables.
        bool needcomma = false;
        bool multiline = false;
        LS.set(nextseq, 1);

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

        // Output the table keys.
        table_getpairs(LS, value, 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) (*output_) << ",";
            needcomma = true;
            if (LS.rawequal(key, nextseq)) {
                (*output_) << " ";
                pprint_r(level + 1, false, val);
                LS.set(nextseq, LS.ckinteger(nextseq) + 1);
            } else {
                multiline = true;
                tabify(level + 1);
                if (LS.isstring(key) && sv::is_lua_id(LS.ckstring(key))) {
                    (*output_) << LS.ckstring(key);
                } else if (LS.istoken(key)) {
                    atomic_print(LUA_TLIGHTUSERDATA, key, false);
                } else {
                    (*output_) << "[";
                    pprint_r(level + 1, false, key);
                    (*output_) << "]";
                }
                if (indent_) {
                    (*output_) << " = ";
                } else {
                    (*output_) << "=";
                }
                pprint_r(level + 1, false, val);
            }
        }

        // Output the metatable.
        LS.getmetatable(val, value);
        if (LS.istable(val) && (LS.gettabletype(val) != LUA_TT_TANGIBLEMETA)) {
            multiline = true;
            if (needcomma) (*output_) << ",";
            needcomma = true;
            tabify(level + 1);
            (*output_) << "<meta> = ";
            pprint_r(level + 1, false, val);
        }
        
        // Close the brackets.
        if (multiline) {
            tabify(level);
        } else if (LS.ckinteger(nextseq) > 1) {
            (*output_) << " ";
        }
        (*output_) << "}";
    }

    // Atomic print interface.
    PrintMachine(LuaCoreStack &LS0, LuaSlot root, bool quote, std::ostream *os) :
            LS_(LS0.state(), tabchpos_) {
        output_ = os;
        atomic_print(LS_.xtype(root), root, quote);
    }

    // Pretty print interface.
    PrintMachine(LuaCoreStack &LS0, LuaSlot root, bool indent, int level, bool expand, std::ostream *os) : 
            LS_(LS0.state(), tabchpos_) {
        next_id_ = 1;
        indent_ = indent;
        LS_.newtable(tabchpos_);
        util::ostringstream preoutput;
        output_ = &preoutput;
        pprint_r(level, expand, root);
        std::string_view pre = preoutput.view();
        
        // Output the results.  We would just copy the characters
        // one by one to the target stream, except that we have to
        // insert <table XX> in front of all tables that got referenced.
        chpos_to_tabnum_.emplace(0x7FFFFFFF, 0);
        auto iter = chpos_to_tabnum_.begin();
        for (int i = 0; i < int(pre.size()); i++) {
            if (i == iter->first) {
                (*os) << "<table " << iter->second << ">";
                iter++;
            }
            (*os).put(pre[i]);
        }
    }
};

void PrettyPrintOptions::parse(LuaKeywordParser &kp) {
    LuaVar option;
    LuaExtStack LS(kp.state(), option);
    kp.check_throw();
    if (kp.optional(option, "indent")) {
        indent = LS.ckboolean(option);
    }
    if (kp.optional(option, "level")) {
        level = LS.ckint(option);
    }
    if (kp.optional(option, "expand")) {
        expand = LS.ckboolean(option);
    }
}

void atomic_print(LuaCoreStack &LS, LuaSlot val, bool quote, std::ostream *os) {
    PrintMachine pm(LS, val, quote, os);
}

void pprint(LuaCoreStack &LS, LuaSlot val, const PrettyPrintOptions &opts, std::ostream *os) {
    PrintMachine pm(LS, val, opts.indent, opts.level, opts.expand, os);
}

LuaDefine(string_pprint, "obj1, obj2, ...",
    "|Pretty-print the specified objects into a string."
    "|"
    "|See also: string.pprintx, which has a lot more options."
    "|This function uses the default options: pretty print indented,"
    "|start at indentation level zero, and always expand the"
    "|top-level table."
    "|") {
    LuaRet result;
    LuaExtraArgs extra;
    LuaDefStack LS(L, result, extra);

    util::ostringstream oss;
    for (int i = 0; i < extra.size(); i++) {
        pprint(LS, extra[i], PrettyPrintOptions(), &oss);
        oss << "\n";
    }
    LS.set(result, oss.view());
    return LS.result();
}

LuaDefine(string_pprintx, "options",
    "|Pretty-print the specified object into a string, with options"
    "|"
    "|Options is a table with these fields:"
    "|"
    "|   value    - the object to pretty-print"
    "|   indent   - if false, suppress newlines and indentation (default: true)"
    "|   level    - base level of indentation (default: zero)"
    "|   expand   - if true, force expansion of top-level table (default: false)"
    "|"
    "|About the expand flag: normally, when you print a class, it just "
    "|prints '<class name>', and when you print a tangible, it just"
    "|prints '<tangible id>'.  But sometimes, you want to see the details."
    "|The expand flag forces it to expand the top-level table, even if the"
    "|top-level table is a tangible or class."
    "|") {
    LuaArg loptions;
    LuaRet result;
    LuaVar value;
    LuaDefStack LS(L, loptions, result, value);
    PrettyPrintOptions options;
    LuaKeywordParser kp(LS, loptions);
    options.parse(kp);
    if (!kp.optional(value, "value")) {
        LS.set(value, LuaNil);
    }
    kp.final_check_throw();
    util::ostringstream oss;
    pprint(LS, value, options, &oss);
    LS.set(result, oss.view());
    return LS.result();
}

LuaDefine(string_print, "obj",
    "|Concise print the specified object into a string"
    "|"
    "|This prints a concise representation of obj into a string.  Tables"
    "|are not expanded: for that, use string.pprint or string.pprintx."
    "|"
    "|The functions string.print and tostring are identical."
    "|") {
    LuaArg val;
    LuaRet result;
    LuaDefStack LS(L, val, result);
    eng::ostringstream oss;
    atomic_print(LS, val, false, &oss);
    LS.set(result, oss.str());
    return LS.result();
}

LuaDefineAlias(tostring, string_print);

LuaDefine(string_isidentifier, "str", "return true if the string is a valid lua identifier") {
    LuaArg str;
    LuaRet result;
    LuaDefStack LS(L, str, result);
    if (LS.isstring(str)) {
        eng::string s = LS.ckstring(str);
        LS.set(result, sv::is_lua_id(s));
    } else {
        LS.set(result, false);
    }
    return LS.result();
}