#pragma once

#include <string>
#include <vector>

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

    void Visit(const ClassDeclarationNode* node) override {
        bool pub_f = node->PublicFields().size() > 0;
        bool prot_f = node->ProtectedFields().size() > 0;
        bool priv_f = node->PrivateFields().size() > 0;
        code_.push_back("class " + node->ClassName() + " {");
        if (pub_f) code_.push_back("  public:");
        std::uint64_t s = code_.size();
        for (std::uint64_t i = 0; i < node->PublicFields().size(); i++) {
            node->PublicFields()[i]->Visit(this);
        }
        for (std::uint64_t i = s; i < code_.size(); i++) {
            code_[i] = "    " + code_[i];
        }
        if (pub_f && prot_f) code_.push_back("");
        if (prot_f) code_.push_back("  protected:");
        s = code_.size();
        for (std::uint64_t i = 0; i < node->ProtectedFields().size(); i++) {
            node->ProtectedFields()[i]->Visit(this);
        }
        for (std::uint64_t i = s; i < code_.size(); i++) {
            code_[i] = "    " + code_[i];
        }
        if (priv_f && (pub_f || prot_f)) code_.push_back("");
        if (priv_f) code_.push_back("  private:");
        s = code_.size();
        for (std::uint64_t i = 0; i < node->PrivateFields().size(); i++) {
            node->PrivateFields()[i]->Visit(this);
        }
        for (std::uint64_t i = s; i < code_.size(); i++) {
            code_[i] = "    " + code_[i];
        }
        code_.push_back("};");
    }
    void Visit(const VarDeclarationNode* node) override {
        code_.push_back(node->TypeName() + " " + node->VarName() + ";");
    }
    void Visit(const MethodDeclarationNode* node) override {
        std::string rt = node->ReturnTypeName();
        std::string mn = node->MethodName();
        code_.push_back(rt + " " + mn + "(");
        for (std::uint64_t i = 0; i < node->Arguments().size(); i++) {
            node->Arguments()[i]->Visit(this);
            code_[code_.size() - 2] += code_[code_.size() - 1];
            code_[code_.size() - 2].pop_back();
            code_.pop_back();
            if (i != node->Arguments().size() - 1)
                code_[code_.size() - 1] += ", ";
        }
        code_[code_.size() - 1] += ");";
    }

    const std::vector<std::string>& GetFormattedCode() const {
        return code_;
    }
 private:
     std::vector<std::string> code_;
     std::uint64_t strings_ = 0;
};