#pragma once
#include <vector>
#include <string>
 
class FormatVisitor : public BaseVisitor {
 public:
    void Visit(const BaseNode* node) override {
        node->Visit(this);
    }
 
    void Visit(const ClassDeclarationNode* node) override {
        int64_t pub_size = node->PublicFields().size();
        int64_t prot_size = node->ProtectedFields().size();
        int64_t priv_size = node->PrivateFields().size();
 
        v.push_back(tubs() + "class " + node->ClassName() + " {");
 
        if (pub_size) {
            v.push_back(tubs() + "  public:");
            visiting(node->PublicFields());
            if (prot_size || priv_size)
                v.push_back("");
        }
 
        if (prot_size) {
            v.push_back(tubs() + "  protected:");
            visiting(node->ProtectedFields());
            if (priv_size)
                v.push_back("");
        }
 
        if (priv_size) {
            v.push_back(tubs() + "  private:");
            visiting(node->PrivateFields());
        }
 
        v.push_back(tubs() + "};");
    }
 
    void Visit(const VarDeclarationNode* node) override {
        std::string arg = node->TypeName() + " " + node->VarName();
        if (var_is_arg) {
            v.back() += arg;
            if (!is_last)
                v.back() += ", ";
        } else {
            v.push_back(tubs() + arg + ";");
        }
    }
 
    void Visit(const MethodDeclarationNode* node) override {
        v.push_back(tubs() + node->ReturnTypeName() + " "
        + node->MethodName() + "(");
        var_is_arg = 1;
        visiting(node->Arguments());
        v.back() += ");";
        var_is_arg = 0;
    }
 
    const std::vector<std::string>& GetFormattedCode() const {
        return v;
    }
 
 private:
    std::vector<std::string> v;
 
    uint64_t depth = 0;
 
    std::string tubs() {
        return std::string(depth * 4, ' ');
    }
 
    bool var_is_arg = 0;
 
    bool is_last = 0;
 
    void visiting(std::vector<BaseNode *> z) {
        depth++;
        for (uint64_t i = 0; i < z.size(); i++) {
            if (i == z.size() - 1)
                is_last = 1;
            z[i]->Visit(this);
        }
        is_last = 0;
        depth--;
    }
};