#include "ld.h"
#include "GnuLdLexer.h"
#include "GnuLdParser.h"
#include "GnuLdParserBaseListener.h"
#include "antlr4-runtime.h"

#include <vector>
#include <set>
#include <map>
#include <locale>
#include <string>
#include <algorithm>
#include <iostream>

using antlr4::ANTLRFileStream;
using namespace std;

enum class MemoryAttribute {
    R, W, X
};

class MemoryArea {
public:
    string name;
    uint64_t origin;
    uint64_t length;
    set<MemoryAttribute> attributes;
};

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

using ParseTree = antlr4::tree::ParseTree;

template<typename V>
class ParseTreeProperty {
public:
    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) {
        _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;
//    map<ParserRuleContext*, uint64_t> expr;

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

        int factor = 1;
        if (endsWith(str, "k")) {
            factor = 1024;
            str = str.substr(0, str.length() - 1);
        } else if (endsWith(str, "k")) {
            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;
    }

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

    void exitExpInt(GnuLdParser::ExpIntContext *ctx) override {
//        System.out.println("ElfinfoGnuLdListener.exitExpInt: ctx->INT()->getText() = " + ctx->INT()->getText());
        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).get());
        uint64_t b = expr.get(ctx->exp(1).get());
        uint64_t x = a - b;
        expr.put(ctx, x);
    }

    void exitExpAdd(GnuLdParser::ExpAddContext *ctx) override {
        uint64_t a = expr.get(ctx->exp(0).get());
        uint64_t b = expr.get(ctx->exp(1).get());
        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());
//        System.out.println("ma.length = " + ma.length);
        expr.put(ctx, ma.length);
    }

    void exitExpOrigin(GnuLdParser::ExpOriginContext *ctx) override {
//        System.out.println("ElfinfoGnuLdListener.exitExpOrigin: " + ctx->getText());
        MemoryArea
                ma = getMemoryArea(ctx->NAME()->getText());
//        System.out.println("ma.origin = " + ma.origin);
        expr.put(ctx, ma.origin);
    }

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

    void enterMustbe_exp(GnuLdParser::Mustbe_expContext *ctx) override {
//        System.out.println("ElfinfoGnuLdListener.enterMustbe_exp");
    }

    void exitMustbe_exp(GnuLdParser::Mustbe_expContext *ctx) override {
//        System.out.println("ElfinfoGnuLdListener.exitMustbe_exp");

        expr.put(ctx, expr.get(ctx->exp().get()));
    }


    void enterOrigin_spec(GnuLdParser::Origin_specContext *ctx) override {
//        System.out.println("ElfinfoGnuLdListener.enterOrigin_spec");
    }

    void exitOrigin_spec(GnuLdParser::Origin_specContext *ctx) override {
//        System.out.println("ElfinfoGnuLdListener.exitOrigin_spec");
    }

    void enterLength_spec(GnuLdParser::Length_specContext *ctx) override {
//        System.out.println("ElfinfoGnuLdListener.enterLength_spec");
    }

    void exitLength_spec(GnuLdParser::Length_specContext *ctx) override {
//        System.out.println("ElfinfoGnuLdListener.exitLength_spec");
    }

    void enterMemory_spec(GnuLdParser::Memory_specContext *ctx) override {
//        System.out.println("ElfinfoGnuLdListener.enterMemory_spec");
    }

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

    MemoryAttribute attribute;
    bool attributesInverted;
    set<MemoryAttribute> attributes;

    void exitAttributes_opt(GnuLdParser::Attributes_optContext *ctx) override {
//        System.out.println("ElfinfoGnuLdListener.exitAttributes_opt");
        attributes.clear();
    }

    void enterAttributeInverted(GnuLdParser::AttributeInvertedContext *ctx) override {
//        System.out.println("ElfinfoGnuLdListener.enterAttributeInverted");

        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();
        const string &name = ctx->NAME()->getText();
//        System.out.println("ctx->ATTRIBUTE()->getText() = " + name);

        attributesInverted = true;
    }

    void enterAttributeNormal(GnuLdParser::AttributeNormalContext *ctx) override {
//        System.out.println("ElfinfoGnuLdListener.enterAttributeNormal");

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

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

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

    GnuLdParser parser(&tokens);
    ElfinfoGnuLdBaseListener listener;
    parser.addParseListener(&listener);
    auto file = parser.file();
    std::cout << file->toStringTree(&parser) << std::endl << std::endl;
    return {};
}