#include "Ld.h"
#include "GnuLdLexer.h"
#include "GnuLdParser.h"
#include "GnuLdParserBaseListener.h"

namespace trygvis {
namespace elfinfo {

using antlr4::ANTLRFileStream;
using namespace std;

using ParseTree = antlr4::tree::ParseTree;

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 std::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;
}

template<typename V>
class ParseTreeProperty {
public:
    virtual V get(Ref<ParseTree> node) {
        return get(node.get());
    }

    virtual V get(ParseTree *const node) {
//        try {
//             cout << "node = " << node->getText() << endl;
            return _annotations.at(node);
//        } catch (std::out_of_range &e) {
//             cout << "out of range: " << node->getText() << endl;
//            throw e;
//        }
    }

    virtual void put(ParseTree *const node, V value) {
        // cout << "put(" << node << ", " << value << ")" << endl;
        _annotations[node] = value;
    }

    virtual V removeFrom(ParseTree *const node) {
        return _annotations.erase(node);
    }

protected:
    std::map<ParseTree *, V> _annotations;

private:
};

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;
            }
        }
        throw out_of_range("No such section: " + name);
    }

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

    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 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 exitExpLoadaddr(GnuLdParser::ExpLoadaddrContext *ctx) override {
        auto &section = getSection(ctx->NAME()->getText());
        expr.put(ctx, 0);
    }

    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) {
}

LdScript LdScriptLoader::load(std::string path) {
    ANTLRFileStream input(path);
    GnuLdLexer lexer(&input);
    CommonTokenStream tokens(&lexer);
    tokens.fill();

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

    GnuLdParser parser(&tokens);
    ElfinfoGnuLdBaseListener listener;
    parser.addParseListener(&listener);
    auto file = parser.file();

    if (debug_) {
        std::cout << file->toStringTree(&parser) << std::endl << std::endl;
    }

    return {
        .memoryAreas = listener.memoryAreas
    };
}

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

} // namespace elfinfo
} // namespace trygvis