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

#include <cmath>
#include <cinttypes>

class PrintMachine {
public:
    LuaVar tabchpos_;
    LuaExtStack LS_;
    int next_id_;
    bool indent_;
    std::ostream *output_;

    // When a general table appears in the output, we use the character
    // position where the table first appeared as a unique ID for the
    // general table.
    eng::map<int, int> chpos_to_table_number_;
    eng::map<int, eng::string> chpos_to_header_;

    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: {
            eng::string classname = LS_.classname(val);
            if (classname.empty()) {
                (*output_) << "<table>";
            } else {
                (*output_) << "<" << classname << ">";
            }
            return;
        }
        case LUA_TT_TANGIBLE: {
            eng::string classname = LS_.classname(val);
            if (classname.empty()) {
                (*output_) << "<tan " << LS_.tanid(val) << ">";
            } else {
                (*output_) << "<tan " << classname << " " << 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;
        LuaExtStack LS(L, loffset, pairs, key, val, lchpos);

        // 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);
        eng::string classname = LS.classname(value);

        // Print the atomic portion or table header.
        // Make a decision about whether to print key-value pairs,
        // if not, then return.
        if (xtype < LUA_TT_GENERAL) {
            atomic_print(xtype, value, true);
            return;
        } else if (xtype > LUA_TT_GENERAL) {
            atomic_print(xtype, value, true);
            if (!expand) return;
        } else {
            LS.rawget(lchpos, tabchpos_, value);
            int chpos = LS.tryint(lchpos).value_or(-1);
            if (chpos < 0) {
                // First time.
                // * The character position where the table first appears
                //   serves as a unique ID for the table.  Record it in the tabchpos
                //   map.  Then, generate an initial tentative header for the table,
                //   and insert it into the header map.  Do not output the header to
                //   output stream: it will be post-injected into the output stream.
                chpos = int((*output_).tellp());
                LS.rawset(tabchpos_, value, chpos);
                if (!classname.empty()) {
                    chpos_to_header_[chpos] = util::ss("<", classname, ">");
                }
            } else {
                // Second time.
                // See if the table already has a table number assigned.  If not,
                // assign it one.  Update the header for the table.
                // Do output the header to the output stream, it is only
                // post-injected at the first table appearance.
                int tabnum = chpos_to_table_number_[chpos];
                if (tabnum == 0) {
                    tabnum = next_id_++;
                    chpos_to_table_number_[chpos] = tabnum;
                    if (classname.empty()) {
                        chpos_to_header_[chpos] = util::ss("<table ", tabnum, ">");
                    } else {
                        chpos_to_header_[chpos] = util::ss("<", classname, " ", tabnum, ">");
                    }
                }
                (*output_) << chpos_to_header_[chpos];
                return; // Do not output key-value pairs the second time.
            }
        }

        // How many keys in the table?
        int nkeys = LS.nkeys(value);

        // Decide whether we're going to print a line for the metatable.
        // If the table has a classname, we don't need, to, because the
        // classname (which is in the header) tell the reader what metatable
        // is being used.  Also, tangible metatables are not shown, because
        // tangibles keep secret stuff in the metatable.   (this may change).
        bool print_meta = false;
        if (classname.empty()) {
            LS.getmetatable(val, value);
            if (LS.istable(val) && (LS.gettabletype(val) != LUA_TT_TANGIBLEMETA)) {
                print_meta = true;
            }
        }

        // Count the number of array-style keys.
        // Also, check if there are any tables in the array keys.
        int narray_keys = 0;
        bool array_contains_table = false;
        for (int i = 1; i <= nkeys; i++) {
            LS.rawget(val, value, i);
            if (LS.isnil(val)) break;
            narray_keys = i;
            if (LS.type(val) == LUA_TTABLE) array_contains_table = true;
        }
        
        // Maybe print it array-style.  This code is simple because
        // we don't do indentation, and we don't handle any values
        // that aren't atomic.
        bool array_style = (narray_keys == nkeys) && (!array_contains_table);
        if (array_style) {
            (*output_) << "{";
            for (int i = 1; i <= nkeys; i++) {
                if (i > 1) (*output_) << ", ";
                LS.rawget(val, value, i);
                atomic_print(LS.xtype(val), val, true);
            }
            (*output_) << "}";
        }

        // Print it table-style.
        if (!array_style) {
            table_getpairs(LS, value, pairs, true);
            (*output_) << "{";
            bool needcomma = false;
            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;
                tabify(level + 1);
                if (!LS.rawequal(key, (i/2))) {
                    if (LS.isstring(key) && sv::is_lua_id(LS.ckstring(key))) {
                        (*output_) << LS.ckstring(key);
                    } else {
                        (*output_) << "[";
                        pprint_r(level + 1, false, key);
                        (*output_) << "]";
                    }
                    if (indent_) {
                        (*output_) << " = ";
                    } else {
                        (*output_) << "=";
                    }
                }
                pprint_r(level + 1, false, val);
            }
            if (print_meta) {
                LS.getmetatable(val, value);
                if (needcomma) (*output_) << ",";
                tabify(level + 1);
                (*output_) << "<meta> = ";
                pprint_r(level + 1, false, val);
            }
            tabify(level);
            (*output_) << "}";
        }
    }

    PrintMachine(LuaCoreStack &LS0, LuaSlot root, const PrettyPrint &opts, std::ostream *os) :
            LS_(LS0.state(), tabchpos_) {
        if (opts.atomic_) {
            output_ = os;
            atomic_print(LS_.xtype(root), root, opts.quote_);
            return;
        }
        next_id_ = 1;
        indent_ = opts.indent_;
        LS_.newtable(tabchpos_);
        util::ostringstream preoutput;
        output_ = &preoutput;
        pprint_r(opts.level_, opts.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 headers wherever the chpos_to_header map says
        // so.
        chpos_to_header_.emplace(0x7FFFFFFF, "");
        auto iter = chpos_to_header_.begin();
        for (int i = 0; i < int(pre.size()); i++) {
            if (i == iter->first) {
                (*os) << iter->second;
                iter++;
            }
            (*os).put(pre[i]);
        }
    }
};


void PrettyPrint::print(LuaCoreStack &LS, LuaSlot val, std::ostream *os) const {
    PrintMachine pm(LS, val, *this, os);
}


//////////////////////////////////////////////////////////////////////////////////
//
// Format 
//
// Printf-style formatting that consumes arguments from LuaExtraArgs.
//
//////////////////////////////////////////////////////////////////////////////////

class FormatDirective
{
public:
    std::string_view precedingliteral;
    std::string_view parameters;    // Raw string_view of '%', flags, width, precision.
    std::string_view modifiers;
    char directive;
    char rebuilt[100];

    // Parsed fields from parameters.
    bool flag_minus;
    bool flag_plus;
    bool flag_space;
    bool flag_hash;
    bool flag_zero;
    int width;              // -1 if unspecified.
    int precision;          // -1 if unspecified.

    const int PARAMETERS_TOO_LONG = 50;

    // Return an error message declaring this format specifier to be invalid.
    //
    eng::string invalid()
    {
        return util::ss("Invalid format specifier: '", parameters, modifiers, directive, "'");
    }

    // Rebuild the format directive, using the specified suffix
    // instead of the stored modifiers and directive.
    //
    const char *rebuild(std::string_view suffix)
    {
        std::string_view params = parameters;
        if (int(params.size()) > PARAMETERS_TOO_LONG) params = "%";
        memcpy(rebuilt, params.data(), params.size());
        memcpy(rebuilt + params.size(), suffix.data(), suffix.size());
        rebuilt[params.size() + suffix.size()] = 0;
        return rebuilt;
    }

    // Read one directive from fmt, advancing fmt past everything consumed.
    //
    // On return:
    //   directive != 0, parameters non-empty  — normal format directive
    //   directive != 0, parameters empty       — not possible
    //   directive == 0, parameters empty       — end of string (may have precedingliteral)
    //   directive == 0, parameters non-empty   — truncated format (missing conversion char)
    //
    void read(std::string_view &fmt)
    {
        // Find the preceding literal (everything before the first '%').
        size_t pct = fmt.find('%');
        precedingliteral = fmt.substr(0, pct);
        if (pct == std::string_view::npos) {
            parameters = {};
            modifiers = {};
            directive = 0;
            fmt = {};
            return;
        }

        // Parse format parameters: '%', flags, width, precision.
        // We parse into individual fields AND store the raw string_view.
        std::string_view rest = fmt.substr(pct + 1);

        // Flags.
        flag_minus = false;
        flag_plus = false;
        flag_space = false;
        flag_hash = false;
        flag_zero = false;
        for (;;) {
            switch (sv::zfront(rest)) {
            case '-': flag_minus = true; rest.remove_prefix(1); continue;
            case '+': flag_plus  = true; rest.remove_prefix(1); continue;
            case ' ': flag_space = true; rest.remove_prefix(1); continue;
            case '#': flag_hash  = true; rest.remove_prefix(1); continue;
            case '0': flag_zero  = true; rest.remove_prefix(1); continue;
            default: break;
            }
            break;
        }

        // Width.
        std::string_view width_str = sv::read_number(rest, false, false, false, false);
        width = width_str.empty() ? -1 : (int)sv::to_int64(width_str);

        // Precision.
        precision = -1;
        if (sv::zfront(rest) == '.') {
            rest.remove_prefix(1);
            std::string_view prec_str = sv::read_number(rest, false, false, false, false);
            precision = prec_str.empty() ? 0 : (int)sv::to_int64(prec_str);
        }

        parameters = fmt.substr(pct, fmt.size() - pct - rest.size());
        fmt = rest;

        // Read 'l' modifiers.
        std::string_view mod_start = fmt;
        while (sv::zfront(fmt) == 'l')
            fmt.remove_prefix(1);
        modifiers = mod_start.substr(0, mod_start.size() - fmt.size());

        // Read conversion character.
        directive = sv::zfront(fmt);
        if (directive) fmt.remove_prefix(1);
    }
};


static void format_signed(LuaCoreStack &LS, LuaSlot arg, const char *format, std::ostream *os) {
    auto ni = LS.tryinteger(arg);
    if (ni) {
        char buf[64];
        snprintf(buf, sizeof(buf), format, *ni);
        (*os) << buf;
    }
}

static void format_unsigned(LuaCoreStack &LS, LuaSlot arg, const char *format, std::ostream *os) {
    auto ni = LS.tryinteger(arg);
    if (ni) {
        char buf[64];
        snprintf(buf, sizeof(buf), format, (uint64_t)(*ni));
        (*os) << buf;
    }
}

static void format_value(LuaCoreStack &LS, LuaSlot arg, FormatDirective &fd, const PrettyPrint &opts, std::ostream *os) {
    if (fd.width <= 0) {
        PrintMachine pm(LS, arg, opts, os);
        return;
    }
    eng::ostringstream tmp;
    PrintMachine pm(LS, arg, opts, &tmp);
    std::string_view sv = tmp.view();
    int pad = fd.width - (int)sv.size();
    if (!fd.flag_minus)
        for (int i = 0; i < pad; i++) os->put(' ');
    os->write(sv.data(), sv.size());
    if (fd.flag_minus)
        for (int i = 0; i < pad; i++) os->put(' ');
}

static void format_double(LuaCoreStack &LS, LuaSlot arg, const char *format, std::ostream *os) {
    auto ni = LS.trynumber(arg);
    if (ni) {
        char buf[64];
        snprintf(buf, sizeof(buf), format, *ni);
        (*os) << buf;
    }
}


eng::string PrettyPrint::format(LuaCoreStack &LS, std::string_view fmt, LuaExtraArgs args, std::ostream *os) {
    FormatDirective fd;

    // First pass: validate the format string and the argument types.
    std::string_view fmtcopy = fmt;
    int nargs = 0;
    while (!fmtcopy.empty()) {
        fd.read(fmtcopy);

        if (int(fd.parameters.size()) > fd.PARAMETERS_TOO_LONG) return fd.invalid();

        if (fd.directive == 0 && !fd.parameters.empty()) return fd.invalid();

        if (fd.directive == '%') {
            if ((fd.parameters.size() != 1)||(fd.modifiers.size() != 0))
                return fd.invalid();
        } else if (fd.directive != 0) {
            if (nargs >= args.size())
                return util::ss("expected more than ", args.size(), " arguments");
            LuaSpecial arg = args[nargs++];
            switch (fd.directive) {
            case 'd': case 'i': case 'o': case 'u': case 'x': case 'X':
            case 'c':
                if (!LS.isinteger(arg))
                    return util::ss("bad argument #", nargs, " (integer expected, got ", LS.typestr(arg), ")");
                break;
            case 'e': case 'E': case 'f': case 'g': case 'G':
                if (!LS.isnumber(arg))
                    return util::ss("bad argument #", nargs, " (number expected, got ", LS.typestr(arg), ")");
                break;
            case 's': case 'q': case 'p': case 'P':
                break;
            default:
                return fd.invalid();
            }
        }
    }

    if (nargs != args.size())
        return util::ss("expected ", nargs, " arguments, got ", args.size());

    // Second pass: produce output.
    int argidx = 0;
    while (!fmt.empty()) {
        fd.read(fmt);

        if (!fd.precedingliteral.empty())
            os->write(fd.precedingliteral.data(), fd.precedingliteral.size());

        if (fd.directive == 0) break;
        if (fd.directive == '%') { os->put('%'); continue; }

        LuaSpecial arg = args[argidx++];

        switch (fd.directive) {
        case 'd': format_signed(LS, arg, fd.rebuild(PRId64), os); break;
        case 'i': format_signed(LS, arg, fd.rebuild(PRId64), os); break;
        case 'o': format_unsigned(LS, arg, fd.rebuild(PRIo64), os); break;
        case 'u': format_unsigned(LS, arg, fd.rebuild(PRIu64), os); break;
        case 'x': format_unsigned(LS, arg, fd.rebuild(PRIx64), os); break;
        case 'X': format_unsigned(LS, arg, fd.rebuild(PRIX64), os); break;
        case 'e': format_double(LS, arg, fd.rebuild("e"), os); break;
        case 'E': format_double(LS, arg, fd.rebuild("E"), os); break;
        case 'f': format_double(LS, arg, fd.rebuild("f"), os); break;
        case 'g': format_double(LS, arg, fd.rebuild("g"), os); break;
        case 'G': format_double(LS, arg, fd.rebuild("G"), os); break;
        case 'c': format_signed(LS, arg, fd.rebuild("c"), os); break;
        case 's': format_value(LS, arg, fd, PrettyPrint::Atomic(false), os); break;
        case 'q': format_value(LS, arg, fd, PrettyPrint::Atomic(true), os); break;
        case 'p': format_value(LS, arg, fd, PrettyPrint::Pretty(fd.modifiers.empty(), false), os); break;
        case 'P': format_value(LS, arg, fd, PrettyPrint::Pretty(fd.modifiers.empty(), true), os); break;
        default: break;
        }
    }

    return {};
}

//////////////////////////////////////////////////////////////////////////////////
//
// Lua Interfaces to the Various Printing Routines
//
// Note: there are more functions like this in world-accessor.
//
//////////////////////////////////////////////////////////////////////////////////

LuaDefine(tostring, "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;
    PrettyPrint::Atomic(false).print(LS, val, &oss);
    LS.set(result, oss.str());
    return LS.result();
}


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

