#pragma once

#include <vector>
#include <string>

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

    void Visit(const ClassDeclarationNode *node) override {
        if (formatted_code_.empty())
            formatted_code_.emplace_back(tab);

        formatted_code_.back() += "class " + node->ClassName() + " {";

        ClassField(node->PublicFields(), "public");

        if (!node->PublicFields().empty() &&
            (!node->ProtectedFields().empty() ||
            !node->PrivateFields().empty()))
            formatted_code_.emplace_back("");

        ClassField(node->ProtectedFields(), "protected");

        if (!node->PrivateFields().empty() &&
        (!node->PublicFields().empty() || !node->PrivateFields().empty()))
            formatted_code_.emplace_back("");

        ClassField(node->PrivateFields(), "private");

        formatted_code_.emplace_back(tab);
        formatted_code_.back() += "};";
    }

    void Visit(const VarDeclarationNode *node) override {
        formatted_code_.back() += node->TypeName() + " " + node->VarName();
    }

    void Visit(const MethodDeclarationNode *node) override {
        formatted_code_.back() += node->ReturnTypeName() +
        " " + node->MethodName() + "(";

        for (std::size_t i = 0; i < node->Arguments().size(); i++) {
            Visit(node->Arguments()[i]);

            if (i != node->Arguments().size() - 1)
                formatted_code_.back() += ", ";
        }

        formatted_code_.back() += ")";
    }

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

 private:
    std::vector<std::string> formatted_code_;

    std::string tab;

    void ClassField(std::vector<BaseNode*> field, std::string field_name) {
        if (!field.empty()) {
            formatted_code_.emplace_back(tab);
            formatted_code_.back() += "  " + field_name + ":";
        }

        tab += "    ";

        for (const auto & item : field) {
            formatted_code_.emplace_back(tab);
            Visit(item);

            if (formatted_code_.back().back() != ';')
                formatted_code_.back() += ";";
        }

        for (int i = 0; i < 4; i++)
            tab.pop_back();
    }
};