#include <string>
#include <vector>
#include <functional>
#include <algorithm>

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

  void Visit(const ClassDeclarationNode *node) override {
    formated.push_back(tab + "class " + node->ClassName() + " {");
    if (!node->PublicFields().empty()) {
      tab += "  ";
      formated.push_back(tab + "public:");
      tab += "  ";
      for (auto it : node->PublicFields()) {
        Visit(it);
      }
      tab.erase(0, 4);
    }
    if (!node->ProtectedFields().empty()) {
      if (formated.back().back() != '{') {
        formated.push_back("");
      }
      tab += "  ";
      formated.push_back(tab + "protected:");
      tab += "  ";
      for (auto it : node->ProtectedFields()) {
        Visit(it);
      }
      tab.erase(0, 4);
    }
    if (!node->PrivateFields().empty()) {
      if (formated.back().back() != '{') {
        formated.push_back("");
      }
      tab += "  ";
      formated.push_back(tab + "private:");
      tab += "  ";
      for (auto it : node->PrivateFields()) {
        Visit(it);
      }
      tab.erase(0, 4);
    }
    formated.push_back(tab + "};");
  }

  void Visit(const VarDeclarationNode *node) override {
    std::string tmp = node->TypeName() + " " + node->VarName();
    if (!args.empty()) {
      args += tmp;
    } else {
      formated.push_back(tab + tmp + ";");
    }
  }
  void Visit(const MethodDeclarationNode *node) override {
    formated.push_back(tab + node->ReturnTypeName() + " " + node->MethodName());
    args = "(";
    if (!node->Arguments().empty()) {
      for (auto it : node->Arguments()) {
        Visit(it);
        args += ", ";
      }
      args.pop_back();
      args.pop_back();
    }
    formated.back() += args + ");";
    args = "";
  }

  [[nodiscard]] const std::vector<std::string> &GetFormattedCode() const {
    return formated;
  }

 private:
  std::vector<std::string> formated;
  std::string args;
  std::string tab = "";
};