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

class FormatVisitor : public BaseVisitor {
 private:
  std::vector<std::string> v;

  uint64_t depth = 0;

  std::string tubs() {
    return std::string(depth * 4, ' ');
  }

 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 pr_size = node->ProtectedFields().size();
    int64_t prv_size = node->PrivateFields().size();
    v.push_back(tubs() + "class " + node->ClassName() + " {");
    if (pub_size) {
      v.push_back(tubs() + "  public:");
      visiting(node->PublicFields());
      if (pr_size || prv_size)
        v.emplace_back("");
    }
    if (pr_size) {
      v.push_back(tubs() + "  protected:");
      visiting(node->ProtectedFields());
      if (prv_size)
        v.emplace_back("");
    }
    if (prv_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 = true;
    visiting(node->Arguments());
    v.back() += ");";
    var_is_arg = false;
  }
  const std::vector<std::string> &GetFormattedCode() const {
    return v;
  }
  bool var_is_arg = false;
  bool is_last = false;

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