#pragma once

#include "antlr4-runtime.h"

namespace trygvis {
namespace antlr {

// This namespace is shared copied code

using ParseTree = antlr4::tree::ParseTree;

class MissingParseTreeProperty : public std::out_of_range {
public:
    explicit MissingParseTreeProperty(const std::string &what) : out_of_range(what) { }
};

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

    virtual V get(ParseTree *const node) {
        if (!debug) {
            return _annotations.at(node);
        }

        try {
//            cerr << "node = " << node->getText() << endl;
            return _annotations.at(node);
        } catch (std::out_of_range &e) {
            std::cerr << "get(" << node << "), text=" << node->getText() << std::endl;
            std::stringstream buf;
            buf << "out of range: " << node << ", text=" << node->getText();
            auto msg = buf.str();
            std::cerr << msg << std::endl;
            throw MissingParseTreeProperty(msg);
        }
    }

    virtual void put(ParseTree *const node, V value) {
        if (debug) {
            std::cerr << "put(" << node << ", " << value << "), text: " << node->getText() << std::endl;
        }
        _annotations[node] = value;
    }

    virtual V removeFrom(ParseTree *const node) {
        auto it = _annotations.find(node);

        if (it == _annotations.end()) {
            throw MissingParseTreeProperty(node->getText());
        }

        return it->second;
    }

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

private:
};

} // namespace antlr

}