#include <string>
#include <vector>

using std::vector;
using std::string;

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

  void Visit(const ClassDeclarationNode* node) override {
    code.push_back(_addTabs() + "class " + node->ClassName() + " {");
    _tabLevel++;

    if (node->PublicFields().size() > 0) {
      code.push_back(_addTabs() + "public:");
      _tabLevel++;
      for (auto m : node->PublicFields())
        this->Visit(m);
      _tabLevel--;
    }

    if (node->ProtectedFields().size() > 0) {
      code.push_back("");
      code.push_back(_addTabs() + "protected:");
      _tabLevel++;
      for (auto m : node->ProtectedFields())
        this->Visit(m);
      _tabLevel--;
    }

    if (node->PrivateFields().size() > 0) {
      code.push_back("");
      code.push_back(_addTabs() + "private:");
      _tabLevel++;
      for (auto m : node->PrivateFields())
        this->Visit(m);
      _tabLevel--;
    }

    _tabLevel--;
    code.push_back(_addTabs() + "};");
  }

  void Visit(const VarDeclarationNode* node) override {
    auto type = node->TypeName();
    auto name = node->VarName();

    if (_inMethod) {
      _buffMeth += type + " " + name;
      return;
    }

    code.push_back(_addTabs() + type + " " + name + ";");
  }

  void Visit(const MethodDeclarationNode* node) override {
    _inMethod = true;
    _buffMeth = _addTabs() + node->ReturnTypeName() + " "
      + node->MethodName() + "(";

    for (size_t i = 0; i < node->Arguments().size(); i++) {
      Visit(node->Arguments()[i]);
      _buffMeth += (node->Arguments().size() - 1 > i) ? ", " : "";
    }

    _buffMeth += ");";
    code.push_back(_buffMeth);
    _inMethod = false;
  }

  const std::vector<std::string>& GetFormattedCode() const {
    return code;
  }

 private:
  bool _inMethod = false;
  string _buffMeth = "";
  vector<string> code;
  size_t _tabLevel = 0;
  string _addTabs() {
    string buff;
    for (size_t i = 0; i < _tabLevel; i++)
      buff += "  ";
    return buff;
  }
};