#include "clang/AST/AST.h"
#include "clang/AST/ASTConsumer.h"
#include "clang/AST/RecursiveASTVisitor.h"
#include "clang/Frontend/CompilerInstance.h"
#include "clang/Frontend/FrontendPluginRegistry.h"
#include "clang/Sema/Sema.h"
#include "llvm/Support/raw_ostream.h"
using namespace clang;
using namespace std;

namespace {

	class VarDeclVisitor : public RecursiveASTVisitor<VarDeclVisitor> {
	public:
		explicit VarDeclVisitor(DiagnosticsEngine& d) : m_diag(d) {}

		bool VisitVarDecl(VarDecl* varDecl) {
			// Checks for all initialised variable declarations
			if (varDecl->getInit()) {

				// Checks for all strings
				if (varDecl->getType().getAsString().find("char") != string::npos) {

					// Create a new string literal
					StringLiteral* stringLiteral =
						StringLiteral::Create(varDecl->getASTContext(), StringRef("New Text"),
							StringLiteral::StringKind::Ascii, false,
							varDecl->getType(), varDecl->getLocation());

					// Change the initialised variable declaration
					varDecl->setInit(stringLiteral);

					// Print the edited AST node
					varDecl->dump();
				}
			}

	private:
		DiagnosticsEngine& m_diag;
	};

	class VarDeclConsumer : public ASTConsumer {
	public:
		explicit VarDeclConsumer(DiagnosticsEngine& d)
			: m_visitor(VarDeclVisitor(d)) {}

		virtual bool HandleTopLevelDecl(DeclGroupRef dg) override {
			for (Decl* decl : dg) {
				m_visitor.TraverseDecl(decl);
			}

			return true;
		}

	private:
		VarDeclVisitor m_visitor;
	};

	class CaesarStrings : public PluginASTAction {
	protected:
		std::unique_ptr<ASTConsumer> CreateASTConsumer(CompilerInstance& ci,
			llvm::StringRef) override {
			return llvm::make_unique<VarDeclConsumer>(ci.getDiagnostics());
		}

		bool ParseArgs(const CompilerInstance&,
			const std::vector<std::string>&) override {
			return true;
		}
	};
} // namespace

static FrontendPluginRegistry::Add<CaesarStrings>
X("my-plugin", "what my plugin does");