#include "trygvis/elfinfo/Ld.h"
#include "trygvis/antlr.h"
#include "GnuLdLexer.h"
#include "GnuLdParser.h"
#include "GnuLdParserBaseListener.h"

namespace trygvis {
namespace elfinfo {

using antlr4::ANTLRFileStream;
using antlr4::tree::ParseTree;
using namespace std;
using namespace trygvis::antlr;

static MemoryAttribute valueOf(char c) {
    switch (c) {
        case 'r':
        case 'R':
            return MemoryAttribute::R;
        case 'w':
        case 'W':
            return MemoryAttribute::W;
        case 'x':
        case 'X':
            return MemoryAttribute::X;
        default:
            throw domain_error("Invalid memory attribute: " + c);
    }
}

static bool endsWith(const string &a, const string &b) {
    return b.length() <= a.length() && a.compare(a.length() - b.length(), b.length(), b) == 0;
}

class ElfinfoGnuLdBaseListener : public GnuLdParserBaseListener {
private:
public:
    vector<MemoryArea> memoryAreas;
    ParseTreeProperty<uint64_t> expr;
    vector<Section> sections;
    MemoryAttribute attribute;
    bool attributesInverted;
    set<MemoryAttribute> attributes;

    MemoryArea getMemoryArea(const string &name) {
        for (MemoryArea &ma : memoryAreas) {
            if (ma.name == name) {
                return ma;
            }
        }
        throw out_of_range("No such memory area: " + name);
    }

    Section &getSection(const string &name) {
        for (Section &s : sections) {
            if (s.name == name) {
                return s;
            }
        }
        sections.emplace_back(Section{.name = name});
        return sections[sections.size() - 1];
    }

    static uint64_t parseInt(const string &s) {
        string str = s;
        transform(begin(str), end(str), begin(str), ::tolower);
        int base = 10;
        if (str.compare(0, 2, "0x") == 0) {
            base = 16;
            str = str.substr(2);
        }

        int factor = 1;
        if (endsWith(str, "k")) {
            factor = 1024;
            str = str.substr(0, str.length() - 1);
        } else if (endsWith(str, "m")) {
            factor = 1024 * 1024;
            str = str.substr(0, str.length() - 1);
        }

        unsigned long long i = strtoull(str.c_str(), NULL, base);

        if (factor > 1) {
            i = i * factor;
        }
        return i;
    }

    void exitSection(GnuLdParser::SectionContext *ctx) override {
        auto name = ctx->NAME()->getText();
        sections.push_back(Section{.name = name});
    }

    void exitExpAlign(GnuLdParser::ExpAlignContext *ctx) override {
        expr.put(ctx, 0);
    }

    virtual void exitExpDefined(GnuLdParser::ExpDefinedContext *ctx) override {
        expr.put(ctx, 0);
    }

    void exitExpInt(GnuLdParser::ExpIntContext *ctx) override {
        uint64_t i = parseInt(ctx->INT()->getText());
        expr.put(ctx, i);
    }

    void exitExpSub(GnuLdParser::ExpSubContext *ctx) override {
        uint64_t a = expr.get(ctx->exp(0));
        uint64_t b = expr.get(ctx->exp(1));
        uint64_t x = a - b;
        expr.put(ctx, x);
    }

    void exitExpAdd(GnuLdParser::ExpAddContext *ctx) override {
        uint64_t a = expr.get(ctx->exp(0));
        uint64_t b = expr.get(ctx->exp(1));
        uint64_t x = a + b;
        expr.put(ctx, x);
    }

    void exitExpMul(GnuLdParser::ExpMulContext *ctx) override {
        uint64_t a = expr.get(ctx->exp(0));
        uint64_t b = expr.get(ctx->exp(1));
        uint64_t x = a * b;
        expr.put(ctx, x);
    }

    void exitExpDiv(GnuLdParser::ExpDivContext *ctx) override {
        uint64_t a = expr.get(ctx->exp(0));
        uint64_t b = expr.get(ctx->exp(1));
        uint64_t x = b > 0 ? a / b : 0;
        expr.put(ctx, x);
    }

    void exitExpName(GnuLdParser::ExpNameContext *ctx) override {
        expr.put(ctx, 0);
    }

    void exitExpAddr(GnuLdParser::ExpAddrContext *ctx) override {
        expr.put(ctx, 0);
    }

    void exitExpSizeof(GnuLdParser::ExpSizeofContext *ctx) override {
        expr.put(ctx, 0);
    }

    void exitExpLengthExp(GnuLdParser::ExpLengthExpContext *ctx) override {
        MemoryArea ma = getMemoryArea(ctx->NAME()->getText());
        expr.put(ctx, ma.length);
    }

    void exitExpOrigin(GnuLdParser::ExpOriginContext *ctx) override {
        MemoryArea ma = getMemoryArea(ctx->NAME()->getText());
        expr.put(ctx, ma.origin);
    }

    void exitMustbe_exp(GnuLdParser::Mustbe_expContext *ctx) override {
        expr.put(ctx, expr.get(ctx->exp()));
    }

    void exitExpParen(GnuLdParser::ExpParenContext *ctx) override {
        expr.put(ctx, expr.get(ctx->exp()));
    }

    void exitExpLoadaddr(GnuLdParser::ExpLoadaddrContext *ctx) override {
        auto &section = getSection(ctx->NAME()->getText());
        expr.put(ctx, 0);
    }

    void exitExpTernary(GnuLdParser::ExpTernaryContext *ctx) override {
        uint64_t a = expr.get(ctx->exp(0));
        uint64_t b = expr.get(ctx->exp(1));
        uint64_t c = expr.get(ctx->exp(2));
        uint64_t x = a ? b : c;
        expr.put(ctx, x);
    }

    void exitExpConstant(GnuLdParser::ExpConstantContext *ctx) override {
        expr.put(ctx, expr.get(ctx->NAME()));
    }

    void exitMemory_spec(GnuLdParser::Memory_specContext *ctx) override {
        MemoryArea ma;
        ma.name = ctx->NAME()->getText();
        ma.attributes = attributes;
        ma.origin = expr.get(ctx->origin_spec()->mustbe_exp());
        ma.length = expr.get(ctx->length_spec()->mustbe_exp());
        memoryAreas.push_back(ma);
    }

    void exitAttributes_opt(GnuLdParser::Attributes_optContext *ctx) override {
        attributes.clear();
    }

    void enterAttributeInverted(GnuLdParser::AttributeInvertedContext *ctx) override {
        if (!attributes.empty()) {
            throw new RuntimeException(
                "Attributes for memory areas can only be attributesInverted (with '!') as the first character in a specification; foo(!rw), not foo(x!rw).");
        }

        const string &name = ctx->NAME()->getText();

        attributesInverted = true;
    }

    void enterAttributeNormal(GnuLdParser::AttributeNormalContext *ctx) override {

        const string &name = ctx->NAME()->getText();
        for (int i = 0; i < name.length(); i++) {
            attribute = valueOf(name[i]);
            attributes.insert(attribute);
        }
        attributesInverted = false;
    }
};

LdScriptLoader::LdScriptLoader() : debug_(false) {
}

class LdErrorListener : public BaseErrorListener {
public:
    vector<string> messages;

    void syntaxError(IRecognizer *recognizer, Token *offendingSymbol, size_t line, int charPositionInLine,
                     const string &msg, exception_ptr e) override {
        messages.push_back("line " + to_string(line) + ":" + to_string(charPositionInLine) + ": " + msg);
    }
};

LdScript LdScriptLoader::load(string path) {
    ifstream stream(path, ios::binary);

    if (!stream.good() || stream.eof()) {
        return {};
    }

    string data;
    while (!stream.eof()) {
        char buf[1000];
        stream.read(buf, sizeof(buf));
        auto count = stream.gcount();
        for (int i = 0; i < count; i++) {
            data.push_back((char) (buf[i] & 0x7f));
        }
    }

    ANTLRInputStream input(data);
    GnuLdLexer lexer(&input);
    CommonTokenStream tokens(&lexer);
    tokens.fill();

    if (debug_) {
        for (auto token : tokens.getTokens()) {
            cout << token->toString() << endl;
        }
    }

    GnuLdParser parser(&tokens);
    ElfinfoGnuLdBaseListener listener;
    parser.addParseListener(&listener);

    LdErrorListener ldErrorListener;
    parser.removeErrorListeners();
    parser.addErrorListener(&ldErrorListener);

    auto file = parser.file();

    if (!ldErrorListener.messages.empty()) {
        throw LdParseException(ldErrorListener.messages);
    }

    if (debug_ && parser.getNumberOfSyntaxErrors() == 0) {
        cout << file->toStringTree(&parser) << endl;
    }

    return {
        .memoryAreas = listener.memoryAreas
    };
}

void LdScriptLoader::setDebug(bool debug) {
    debug_ = debug;
}

} // namespace elfinfo
} // namespace trygvis