#include "clang/AST/ASTConsumer.h"
#include "clang/AST/RecursiveASTVisitor.h"
#include "clang/Frontend/CompilerInstance.h"
#include "clang/Frontend/FrontendAction.h"
#include "clang/Tooling/CommonOptionsParser.h"
#include "clang/Tooling/Tooling.h"
#include "llvm/Support/CommandLine.h"
#include <spdlog/spdlog.h>

using namespace clang;
using namespace clang::tooling;
using namespace llvm;

// Apply a custom category to all command-line options so that they are the
// only ones displayed.
static cl::OptionCategory MyToolCategory("my-tool options");

// CommonOptionsParser declares HelpMessage with a description of the common
// command-line options related to the compilation database and input files.
// It's nice to have this help message in all tools.
static cl::extrahelp CommonHelp(CommonOptionsParser::HelpMessage);

// A help message for this specific tool can be added afterwards.
static cl::extrahelp MoreHelp("\nMore help text...\n");

class FunctionVisitor : public RecursiveASTVisitor<FunctionVisitor> {
private:
    ASTContext *Context;

public:
    explicit FunctionVisitor(ASTContext *Context) : Context(Context) {}

    bool VisitFunctionDecl(FunctionDecl *Declaration) {
        if (Declaration->hasBody()) {
            SourceManager &SM = Context->getSourceManager();
            SourceLocation Loc = Declaration->getLocation();
            
            if (SM.isInMainFile(Loc)) {
                std::string FuncName = Declaration->getNameInfo().getName().getAsString();
                unsigned Line = SM.getExpansionLineNumber(Loc);
                
                spdlog::info("Found function '{}' at line {}", FuncName, Line);
                
                // Count parameters
                unsigned ParamCount = Declaration->getNumParams();
                spdlog::debug("  - Function '{}' has {} parameters", FuncName, ParamCount);
                
                // Check if it's a main function
                if (FuncName == "main") {
                    spdlog::warn("Found main function at line {}", Line);
                }
            }
        }
        return true;
    }

    bool VisitVarDecl(VarDecl *Declaration) {
        SourceManager &SM = Context->getSourceManager();
        SourceLocation Loc = Declaration->getLocation();
        
        if (SM.isInMainFile(Loc)) {
            std::string VarName = Declaration->getNameAsString();
            unsigned Line = SM.getExpansionLineNumber(Loc);
            
            spdlog::debug("Found variable '{}' at line {}", VarName, Line);
        }
        return true;
    }
};

class FunctionConsumer : public ASTConsumer {
private:
    FunctionVisitor Visitor;

public:
    explicit FunctionConsumer(ASTContext *Context) : Visitor(Context) {}

    void HandleTranslationUnit(ASTContext &Context) override {
        spdlog::info("Starting AST traversal...");
        Visitor.TraverseDecl(Context.getTranslationUnitDecl());
        spdlog::info("AST traversal completed");
    }
};

class FunctionAction : public ASTFrontendAction {
public:
    std::unique_ptr<ASTConsumer> CreateASTConsumer(CompilerInstance &Compiler,
                                                   StringRef InFile) override {
        spdlog::info("Processing file: {}", InFile.str());
        return std::make_unique<FunctionConsumer>(&Compiler.getASTContext());
    }
};

int main(int argc, const char **argv) {
    // Set up logging
    spdlog::set_level(spdlog::level::debug);
    spdlog::info("Starting clang tooling example");

    auto ExpectedParser = CommonOptionsParser::create(argc, argv, MyToolCategory);
    if (!ExpectedParser) {
        // Fail gracefully for unsupported options.
        llvm::errs() << ExpectedParser.takeError();
        return 1;
    }
    CommonOptionsParser &OptionsParser = ExpectedParser.get();
    ClangTool Tool(OptionsParser.getCompilations(),
                   OptionsParser.getSourcePathList());

    spdlog::info("Running tool on {} source files", 
                 OptionsParser.getSourcePathList().size());

    int Result = Tool.run(newFrontendActionFactory<FunctionAction>().get());
    
    if (Result == 0) {
        spdlog::info("Tool completed successfully");
    } else {
        spdlog::error("Tool failed with exit code {}", Result);
    }
    
    return Result;
}