#include "IndexAction.h"
#include "clang/Frontend/CompilerInstance.h"
#include "clang/Index/IndexDataConsumer.h"
#include "clang/Index/IndexingAction.h"
#include "clang/Tooling/Tooling.h"

using namespace llvm;
namespace clang {
namespace clangd {
namespace {

llvm::Optional<std::string> toURI(const FileEntry *File) {
  if (!File)
    return llvm::None;
  auto AbsolutePath = File->tryGetRealPathName();
  if (AbsolutePath.empty())
    return llvm::None;
  return URI::create(AbsolutePath).toString();
}

// Collects the nodes and edges of include graph during indexing action.
// Important: The graph generated by those callbacks might contain cycles and
// self edges.
struct IncludeGraphCollector : public PPCallbacks {
public:
  IncludeGraphCollector(const SourceManager &SM, IncludeGraph &IG)
      : SM(SM), IG(IG) {}

  // Populates everything except direct includes for a node, which represents
  // edges in the include graph and populated in inclusion directive.
  // We cannot populate the fields in InclusionDirective because it does not
  // have access to the contents of the included file.
  void FileChanged(SourceLocation Loc, FileChangeReason Reason,
                   SrcMgr::CharacteristicKind FileType,
                   FileID PrevFID) override {
    // We only need to process each file once. So we don't care about anything
    // but entries.
    if (Reason != FileChangeReason::EnterFile)
      return;

    const auto FileID = SM.getFileID(Loc);
    const auto File = SM.getFileEntryForID(FileID);
    auto URI = toURI(File);
    if (!URI)
      return;
    auto I = IG.try_emplace(*URI).first;

    auto &Node = I->getValue();
    // Node has already been populated.
    if (Node.URI.data() == I->getKeyData()) {
#ifndef NDEBUG
      auto Digest = digestFile(SM, FileID);
      assert(Digest && Node.Digest == *Digest &&
             "Same file, different digest?");
#endif
      return;
    }
    if (auto Digest = digestFile(SM, FileID))
      Node.Digest = std::move(*Digest);
    Node.IsTU = FileID == SM.getMainFileID();
    Node.URI = I->getKey();
  }

  // Add edges from including files to includes.
  void InclusionDirective(SourceLocation HashLoc, const Token &IncludeTok,
                          StringRef FileName, bool IsAngled,
                          CharSourceRange FilenameRange, const FileEntry *File,
                          StringRef SearchPath, StringRef RelativePath,
                          const Module *Imported,
                          SrcMgr::CharacteristicKind FileType) override {
    auto IncludeURI = toURI(File);
    if (!IncludeURI)
      return;

    auto IncludingURI = toURI(SM.getFileEntryForID(SM.getFileID(HashLoc)));
    if (!IncludingURI)
      return;

    auto NodeForInclude = IG.try_emplace(*IncludeURI).first->getKey();
    auto NodeForIncluding = IG.try_emplace(*IncludingURI);

    NodeForIncluding.first->getValue().DirectIncludes.push_back(NodeForInclude);
  }

  // Sanity check to ensure we have already populated a skipped file.
  void FileSkipped(const FileEntry &SkippedFile, const Token &FilenameTok,
                   SrcMgr::CharacteristicKind FileType) override {
#ifndef NDEBUG
    auto URI = toURI(&SkippedFile);
    if (!URI)
      return;
    auto I = IG.try_emplace(*URI);
    assert(!I.second && "File inserted for the first time on skip.");
    assert(I.first->getKeyData() == I.first->getValue().URI.data() &&
           "Node have not been populated yet");
#endif
  }

private:
  const SourceManager &SM;
  IncludeGraph &IG;
};

// Wraps the index action and reports index data after each translation unit.
class IndexAction : public WrapperFrontendAction {
public:
  IndexAction(std::shared_ptr<SymbolCollector> C,
              std::unique_ptr<CanonicalIncludes> Includes,
              const index::IndexingOptions &Opts,
              std::function<void(SymbolSlab)> SymbolsCallback,
              std::function<void(RefSlab)> RefsCallback,
              std::function<void(IncludeGraph)> IncludeGraphCallback)
      : WrapperFrontendAction(index::createIndexingAction(C, Opts, nullptr)),
        SymbolsCallback(SymbolsCallback), RefsCallback(RefsCallback),
        IncludeGraphCallback(IncludeGraphCallback), Collector(C),
        Includes(std::move(Includes)),
        PragmaHandler(collectIWYUHeaderMaps(this->Includes.get())) {}

  std::unique_ptr<ASTConsumer> CreateASTConsumer(CompilerInstance &CI,
                                                 StringRef InFile) override {
    CI.getPreprocessor().addCommentHandler(PragmaHandler.get());
    if (IncludeGraphCallback != nullptr)
      CI.getPreprocessor().addPPCallbacks(
          llvm::make_unique<IncludeGraphCollector>(CI.getSourceManager(), IG));
    return WrapperFrontendAction::CreateASTConsumer(CI, InFile);
  }

  bool BeginInvocation(CompilerInstance &CI) override {
    // We want all comments, not just the doxygen ones.
    CI.getLangOpts().CommentOpts.ParseAllComments = true;
    return WrapperFrontendAction::BeginInvocation(CI);
  }

  void EndSourceFileAction() override {
    WrapperFrontendAction::EndSourceFileAction();

    const auto &CI = getCompilerInstance();
    if (CI.hasDiagnostics() &&
        CI.getDiagnostics().hasUncompilableErrorOccurred()) {
      errs() << "Skipping TU due to uncompilable errors\n";
      return;
    }
    SymbolsCallback(Collector->takeSymbols());
    if (RefsCallback != nullptr)
      RefsCallback(Collector->takeRefs());
    if (IncludeGraphCallback != nullptr) {
#ifndef NDEBUG
      // This checks if all nodes are initialized.
      for (const auto &Node : IG)
        assert(Node.getKeyData() == Node.getValue().URI.data());
#endif
      IncludeGraphCallback(std::move(IG));
    }
  }

private:
  std::function<void(SymbolSlab)> SymbolsCallback;
  std::function<void(RefSlab)> RefsCallback;
  std::function<void(IncludeGraph)> IncludeGraphCallback;
  std::shared_ptr<SymbolCollector> Collector;
  std::unique_ptr<CanonicalIncludes> Includes;
  std::unique_ptr<CommentHandler> PragmaHandler;
  IncludeGraph IG;
};

} // namespace

std::unique_ptr<FrontendAction> createStaticIndexingAction(
    SymbolCollector::Options Opts,
    std::function<void(SymbolSlab)> SymbolsCallback,
    std::function<void(RefSlab)> RefsCallback,
    std::function<void(IncludeGraph)> IncludeGraphCallback) {
  index::IndexingOptions IndexOpts;
  IndexOpts.SystemSymbolFilter =
      index::IndexingOptions::SystemSymbolFilterKind::All;
  Opts.CollectIncludePath = true;
  Opts.CountReferences = true;
  Opts.Origin = SymbolOrigin::Static;
  if (RefsCallback != nullptr) {
    Opts.RefFilter = RefKind::All;
    Opts.RefsInHeaders = true;
  }
  auto Includes = llvm::make_unique<CanonicalIncludes>();
  addSystemHeadersMapping(Includes.get());
  Opts.Includes = Includes.get();
  return llvm::make_unique<IndexAction>(
      std::make_shared<SymbolCollector>(std::move(Opts)), std::move(Includes),
      IndexOpts, SymbolsCallback, RefsCallback, IncludeGraphCallback);
}

} // namespace clangd
} // namespace clang