#pragma once
#include <vector>
#include <string>

class FormatVisitor : public BaseVisitor {
 private:
  std::vector<std::string> v;
  uint64_t dth = 0;
  std::string tb() {
    return std::string(dth * 4, ' ');
  }
  bool is_arg = false;
  bool is_lst = false;

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

  void Visit(const ClassDeclarationNode* node) override {
    int64_t Public_size = node->PublicFields().size();
    int64_t Protected_size = node->ProtectedFields().size();
    int64_t Private_size = node->PrivateFields().size();

    v.emplace_back(tb() + "class " + node->ClassName() + " {");

    if (Public_size) {
      v.emplace_back(tb() + "  public:");
      visiting(node->PublicFields());
      if (Protected_size || Private_size)
        v.emplace_back("");
    }

    if (Protected_size) {
      v.emplace_back(tb() + "  protected:");
      visiting(node->ProtectedFields());
      if (Private_size)
        v.emplace_back("");
    }

    if (Private_size) {
      v.emplace_back(tb() + "  private:");
      visiting(node->PrivateFields());
    }

    v.emplace_back(tb() + "};");
  }
  const std::vector<std::string>& GetFormattedCode() const {
    return v;
  }
  void Visit(const VarDeclarationNode* node) override {
    std::string arg = node->TypeName() + " " + node->VarName();
    if (is_arg) {
      v.back() += arg;
      if (!is_lst)
        v.back() += ", ";
    } else {
      v.emplace_back(tb() + arg + ";");
    }
  }

  void Visit(const MethodDeclarationNode* node) override {
    v.emplace_back(tb() + node->ReturnTypeName() + " "
                    + node->MethodName() + "(");
    is_arg = true;
    visiting(node->Arguments());
    v.back() += ");";
    is_arg = false;
  }

  void visiting(std::vector<BaseNode *> z) {
    dth++;
    for (uint64_t i = 0; i < z.size(); i++) {
      if (i == z.size() - 1)
        is_lst = true;
      z[i]->Visit(this);
    }
    is_lst = false;
    dth--;
  }
};