#include <string>
#include <vector>

class FormatVisitor : public BaseVisitor {
 public:
    void Visit(const BaseNode *node) override {
        node->Visit(this);
    }

    void Visit(const ClassDeclarationNode *node) override {
        this->node_class = "class";
        this->formattedCode.push_back("class " + node->ClassName() + " {");
        if (!node->PublicFields().empty())
            this->formattedCode.push_back("  public:");
        for (auto i : node->PublicFields()) {
            auto Fv = new FormatVisitor();
            Fv->Visit(i);
            for (auto j : Fv->GetFormattedCode()) {
                this->formattedCode.push_back("    " + j + (j[j.size() - 1] ==
                ';' || j[j.size() - 1] == ':' || j[j.size() - 1] ==
                '{' ? "" : ";"));
            }
            delete Fv;
        }
        if (!node->ProtectedFields().empty()) {
            this->formattedCode.push_back("");
            this->formattedCode.push_back("  protected:");
        }
        for (auto i : node->ProtectedFields()) {
            auto Fv = new FormatVisitor();
            Fv->Visit(i);
            for (auto j : Fv->GetFormattedCode()) {
                this->formattedCode.push_back("    " + j + (j[j.size() - 1] ==
                ';' || j[j.size() - 1] == ':' || j[j.size() - 1] ==
                '{' ? "" : ";"));
            }
            delete Fv;
        }

        if (!node->PrivateFields().empty()) {
            this->formattedCode.push_back("");
            this->formattedCode.push_back("  private:");
        }
        for (auto i : node->PrivateFields()) {
            auto Fv = new FormatVisitor();
            Fv->Visit(i);
            for (auto j : Fv->GetFormattedCode()) {
                this->formattedCode.push_back("    " + j + (j[j.size() - 1] ==
                ';' || j[j.size() - 1] == ':' || j[j.size() - 1] ==
                '{' ? "" : ";"));
            }
            delete Fv;
        }
        this->formattedCode.push_back("};");
    };

    void Visit(const VarDeclarationNode *node) override {
        this->node_class = "var";
        this->formattedCode.push_back(node->TypeName() + " " + node->VarName());
    };

    void Visit(const MethodDeclarationNode *node) override {
        this->node_class = "method";
        std::string tmp = node->ReturnTypeName() + " " + node->MethodName() +
                          "(";
        auto args = node->Arguments();
        for (size_t i = 0; i < args.size(); ++i) {
            auto Fv = new FormatVisitor();
            Fv->Visit(args[i]);
            tmp += Fv->GetFormattedCode()[0] + (args.size() - i == 1 ? "" :
                   ", ");
            delete Fv;
        }
        tmp += ");";
        this->formattedCode.push_back(tmp);
    };

    const std::vector<std::string> &GetFormattedCode() const {
        return this->formattedCode;
    }

 private:
    std::string node_class;
    std::vector<std::string> formattedCode{};
};