/* Copyright 2017-2018 ccls Authors

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

  http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "clang_complete.hh"

#include "clang_utils.h"
#include "filesystem.hh"
#include "log.hh"
#include "platform.h"

#include <clang/Frontend/CompilerInstance.h>
#include <clang/Lex/PreprocessorOptions.h>
#include <clang/Sema/CodeCompleteConsumer.h>
#include <clang/Sema/Sema.h>
#include <llvm/ADT/Twine.h>
#include <llvm/Config/llvm-config.h>
#include <llvm/Support/CrashRecoveryContext.h>
#include <llvm/Support/Threading.h>
using namespace clang;
using namespace llvm;

#include <algorithm>
#include <thread>

namespace ccls {
namespace {

std::string StripFileType(const std::string &path) {
  SmallString<128> Ret;
  sys::path::append(Ret, sys::path::parent_path(path), sys::path::stem(path));
  return sys::path::convert_to_slash(Ret);
}

bool LocationInRange(SourceLocation L, CharSourceRange R,
                     const SourceManager &M) {
  assert(R.isCharRange());
  if (!R.isValid() || M.getFileID(R.getBegin()) != M.getFileID(R.getEnd()) ||
      M.getFileID(R.getBegin()) != M.getFileID(L))
    return false;
  return L != R.getEnd() && M.isPointWithin(L, R.getBegin(), R.getEnd());
}

CharSourceRange DiagnosticRange(const clang::Diagnostic &D, const LangOptions &L) {
  auto &M = D.getSourceManager();
  auto Loc = M.getFileLoc(D.getLocation());
  // Accept the first range that contains the location.
  for (const auto &CR : D.getRanges()) {
    auto R = Lexer::makeFileCharRange(CR, M, L);
    if (LocationInRange(Loc, R, M))
      return R;
  }
  // The range may be given as a fixit hint instead.
  for (const auto &F : D.getFixItHints()) {
    auto R = Lexer::makeFileCharRange(F.RemoveRange, M, L);
    if (LocationInRange(Loc, R, M))
      return R;
  }
  // If no suitable range is found, just use the token at the location.
  auto R = Lexer::makeFileCharRange(CharSourceRange::getTokenRange(Loc), M, L);
  if (!R.isValid()) // Fall back to location only, let the editor deal with it.
    R = CharSourceRange::getCharRange(Loc);
  return R;
}



class StoreDiags : public DiagnosticConsumer {
  const LangOptions *LangOpts;
  std::optional<Diag> last;
  std::vector<Diag> output;
  void Flush() {
    if (!last)
      return;
    bool mentions = last->inside_main || last->edits.size();
    if (!mentions)
      for (auto &N : last->notes)
        if (N.inside_main)
          mentions = true;
    if (mentions)
      output.push_back(std::move(*last));
    last.reset();
  }
public:
  std::vector<Diag> Take() {
    return std::move(output);
  }
  void BeginSourceFile(const LangOptions &Opts, const Preprocessor *) override {
    LangOpts = &Opts;
  }
  void EndSourceFile() override {
    Flush();
  }
  void HandleDiagnostic(DiagnosticsEngine::Level Level,
                        const Diagnostic &Info) override {
    DiagnosticConsumer::HandleDiagnostic(Level, Info);
    SourceLocation L = Info.getLocation();
    if (!L.isValid()) return;
    const SourceManager &SM = Info.getSourceManager();
    bool inside_main = SM.isInMainFile(L);
    auto fillDiagBase = [&](DiagBase &d) {
      llvm::SmallString<64> Message;
      Info.FormatDiagnostic(Message);
      d.range =
          FromCharSourceRange(SM, *LangOpts, DiagnosticRange(Info, *LangOpts));
      d.message = Message.str();
      d.inside_main = inside_main;
      d.file = SM.getFilename(Info.getLocation());
      d.level = Level;
      d.category = DiagnosticIDs::getCategoryNumberForDiag(Info.getID());
    };

    auto addFix = [&](bool SyntheticMessage) -> bool {
      if (!inside_main)
        return false;
      for (const FixItHint &FixIt : Info.getFixItHints()) {
        if (!SM.isInMainFile(FixIt.RemoveRange.getBegin()))
          return false;
        lsTextEdit edit;
        edit.newText = FixIt.CodeToInsert;
        auto r = FromCharSourceRange(SM, *LangOpts, FixIt.RemoveRange);
        edit.range =
            lsRange{{r.start.line, r.start.column}, {r.end.line, r.end.column}};
        last->edits.push_back(std::move(edit));
      }
      return true;
    };

    if (Level == DiagnosticsEngine::Note || Level == DiagnosticsEngine::Remark) {
      if (Info.getFixItHints().size()) {
        addFix(false);
      } else {
        Note &n = last->notes.emplace_back();
        fillDiagBase(n);
      }
    } else {
      Flush();
      last = Diag();
      fillDiagBase(*last);
      if (!Info.getFixItHints().empty())
        addFix(true);
    }
  }
};

std::unique_ptr<CompilerInstance> BuildCompilerInstance(
    CompletionSession &session, std::unique_ptr<CompilerInvocation> CI,
    DiagnosticConsumer &DC, const WorkingFiles::Snapshot &snapshot,
    std::vector<std::unique_ptr<llvm::MemoryBuffer>> &Bufs) {
  for (auto &file : snapshot.files) {
    Bufs.push_back(llvm::MemoryBuffer::getMemBuffer(file.content));
    if (file.filename == session.file.filename) {
      if (auto Preamble = session.GetPreamble()) {
#if LLVM_VERSION_MAJOR >= 7
        Preamble->Preamble.OverridePreamble(*CI, session.FS,
                                            Bufs.back().get());
#else
        Preamble->Preamble.AddImplicitPreamble(*CI, session.FS,
                                               Bufs.back().get());
#endif
      } else {
        CI->getPreprocessorOpts().addRemappedFile(
            CI->getFrontendOpts().Inputs[0].getFile(), Bufs.back().get());
      }
    } else {
      CI->getPreprocessorOpts().addRemappedFile(file.filename,
                                                Bufs.back().get());
    }
  }

  auto Clang = std::make_unique<CompilerInstance>(session.PCH);
  Clang->setInvocation(std::move(CI));
  Clang->setVirtualFileSystem(session.FS);
  Clang->createDiagnostics(&DC, false);
  Clang->setTarget(TargetInfo::CreateTargetInfo(
      Clang->getDiagnostics(), Clang->getInvocation().TargetOpts));
  if (!Clang->hasTarget())
    return nullptr;
  return Clang;
}

bool Parse(CompilerInstance &Clang) {
  SyntaxOnlyAction Action;
  if (!Action.BeginSourceFile(Clang, Clang.getFrontendOpts().Inputs[0]))
    return false;
  if (!Action.Execute())
    return false;
  Action.EndSourceFile();
  return true;
}

void CompletionPreloadMain(CompletionManager *manager) {
  while (true) {
    // Fetching the completion request blocks until we have a request.
    auto request = manager->preload_requests_.Dequeue();

    // If we don't get a session then that means we don't care about the file
    // anymore - abandon the request.
    std::shared_ptr<CompletionSession> session = manager->TryGetSession(
        request.path, false /*mark_as_completion*/, false /*create_if_needed*/);
    if (!session)
      continue;

    const auto &args = session->file.args;
    WorkingFiles::Snapshot snapshot = session->wfiles->AsSnapshot(
      {StripFileType(session->file.filename)});

    LOG_S(INFO) << "create completion session for " << session->file.filename;
    if (std::unique_ptr<CompilerInvocation> CI =
            BuildCompilerInvocation(args, session->FS))
      session->BuildPreamble(*CI);
    if (g_config->diagnostics.onSave) {
      lsTextDocumentIdentifier document;
      document.uri = lsDocumentUri::FromPath(request.path);
      manager->diagnostic_request_.PushBack({document}, true);
    }
  }
}

void CompletionMain(CompletionManager *completion_manager) {
  while (true) {
    // Fetching the completion request blocks until we have a request.
    std::unique_ptr<CompletionManager::CompletionRequest> request =
        completion_manager->completion_request_.Dequeue();

    // Drop older requests if we're not buffering.
    while (g_config->completion.dropOldRequests &&
           !completion_manager->completion_request_.IsEmpty()) {
      completion_manager->on_dropped_(request->id);
      request->Consumer.reset();
      request->on_complete(nullptr);
      request = completion_manager->completion_request_.Dequeue();
    }

    std::string path = request->document.uri.GetPath();

    std::shared_ptr<CompletionSession> session =
        completion_manager->TryGetSession(path, true /*mark_as_completion*/,
                                          true /*create_if_needed*/);

    std::unique_ptr<CompilerInvocation> CI =
        BuildCompilerInvocation(session->file.args, session->FS);
    if (!CI)
      continue;
    auto &FOpts = CI->getFrontendOpts();
    FOpts.CodeCompleteOpts = request->CCOpts;
    FOpts.CodeCompletionAt.FileName = path;
    FOpts.CodeCompletionAt.Line = request->position.line + 1;
    FOpts.CodeCompletionAt.Column = request->position.character + 1;
    FOpts.SkipFunctionBodies = true;
    CI->getLangOpts()->CommentOpts.ParseAllComments = true;

    DiagnosticConsumer DC;
    WorkingFiles::Snapshot snapshot =
      completion_manager->working_files_->AsSnapshot({StripFileType(path)});
    std::vector<std::unique_ptr<llvm::MemoryBuffer>> Bufs;
    auto Clang = BuildCompilerInstance(*session, std::move(CI), DC, snapshot, Bufs);
    if (!Clang)
      continue;

    Clang->setCodeCompletionConsumer(request->Consumer.release());
    if (!Parse(*Clang))
      continue;
    for (auto &Buf : Bufs)
      Buf.release();

    request->on_complete(&Clang->getCodeCompletionConsumer());
  }
}

llvm::StringRef diagLeveltoString(DiagnosticsEngine::Level Lvl) {
  switch (Lvl) {
  case DiagnosticsEngine::Ignored:
    return "ignored";
  case DiagnosticsEngine::Note:
    return "note";
  case DiagnosticsEngine::Remark:
    return "remark";
  case DiagnosticsEngine::Warning:
    return "warning";
  case DiagnosticsEngine::Error:
    return "error";
  case DiagnosticsEngine::Fatal:
    return "fatal error";
  }
}

void printDiag(llvm::raw_string_ostream &OS, const DiagBase &d) {
  if (d.inside_main)
    OS << llvm::sys::path::filename(d.file);
  else
    OS << d.file;
  auto pos = d.range.start;
  OS << ":" << (pos.line + 1) << ":" << (pos.column + 1) << ":"
     << (d.inside_main ? " " : "\n");
  OS << diagLeveltoString(d.level) << ": " << d.message;
}

void DiagnosticMain(CompletionManager *manager) {
  while (true) {
    // Fetching the completion request blocks until we have a request.
    CompletionManager::DiagnosticRequest request =
        manager->diagnostic_request_.Dequeue();
    std::string path = request.document.uri.GetPath();

    std::shared_ptr<CompletionSession> session = manager->TryGetSession(
        path, true /*mark_as_completion*/, true /*create_if_needed*/);

    std::unique_ptr<CompilerInvocation> CI =
        BuildCompilerInvocation(session->file.args, session->FS);
    if (!CI)
      continue;
    CI->getDiagnosticOpts().IgnoreWarnings = false;
    CI->getLangOpts()->SpellChecking = g_config->diagnostics.spellChecking;
    StoreDiags DC;
    WorkingFiles::Snapshot snapshot =
        manager->working_files_->AsSnapshot({StripFileType(path)});
    std::vector<std::unique_ptr<llvm::MemoryBuffer>> Bufs;
    auto Clang = BuildCompilerInstance(*session, std::move(CI), DC, snapshot, Bufs);
    if (!Clang)
      continue;
    if (!Parse(*Clang))
      continue;
    for (auto &Buf : Bufs)
      Buf.release();

    auto Fill = [](const DiagBase &d, lsDiagnostic &ret) {
      ret.range = lsRange{{d.range.start.line, d.range.start.column},
                          {d.range.end.line, d.range.end.column}};
      switch (d.level) {
      case DiagnosticsEngine::Ignored:
        // llvm_unreachable
      case DiagnosticsEngine::Remark:
        ret.severity = lsDiagnosticSeverity::Hint;
        break;
      case DiagnosticsEngine::Note:
        ret.severity = lsDiagnosticSeverity::Information;
        break;
      case DiagnosticsEngine::Warning:
        ret.severity = lsDiagnosticSeverity::Warning;
        break;
      case DiagnosticsEngine::Error:
      case DiagnosticsEngine::Fatal:
        ret.severity = lsDiagnosticSeverity::Error;
        break;
      }
      ret.code = d.category;
      return ret;
    };

    std::vector<lsDiagnostic> ls_diags;
    for (auto &d : DC.Take()) {
      if (!d.inside_main)
        continue;
      std::string buf;
      llvm::raw_string_ostream OS(buf);
      lsDiagnostic &ls_diag = ls_diags.emplace_back();
      Fill(d, ls_diag);
      ls_diag.fixits_ = d.edits;
      OS << d.message;
      for (auto &n : d.notes) {
        OS << "\n\n";
        printDiag(OS, n);
      }
      OS.flush();
      ls_diag.message = std::move(buf);
      for (auto &n : d.notes) {
        if (!n.inside_main)
          continue;
        lsDiagnostic &ls_diag1 = ls_diags.emplace_back();
        Fill(n, ls_diag1);
        OS << n.message << "\n\n";
        printDiag(OS, d);
        OS.flush();
        ls_diag1.message = std::move(buf);
      }
    }
    manager->on_diagnostic_(path, ls_diags);
  }
}

} // namespace

std::shared_ptr<PreambleData> CompletionSession::GetPreamble() {
  std::lock_guard<std::mutex> lock(mutex);
  return preamble;
}

void CompletionSession::BuildPreamble(CompilerInvocation &CI) {
  std::shared_ptr<PreambleData> OldP = GetPreamble();
  std::string content = wfiles->GetContent(file.filename);
  std::unique_ptr<llvm::MemoryBuffer> Buf =
      llvm::MemoryBuffer::getMemBuffer(content);
  auto Bounds = ComputePreambleBounds(*CI.getLangOpts(), Buf.get(), 0);
  if (OldP && OldP->Preamble.CanReuse(CI, Buf.get(), Bounds, FS.get()))
    return;
  CI.getDiagnosticOpts().IgnoreWarnings = false;
  CI.getFrontendOpts().SkipFunctionBodies = true;
  CI.getLangOpts()->CommentOpts.ParseAllComments = true;
#if LLVM_VERSION_MAJOR >= 7
  CI.getPreprocessorOpts().WriteCommentListToPCH = false;
#endif

  StoreDiags DC;
  IntrusiveRefCntPtr<DiagnosticsEngine> DE =
      CompilerInstance::createDiagnostics(&CI.getDiagnosticOpts(), &DC, false);
  PreambleCallbacks PP;
  if (auto NewPreamble = PrecompiledPreamble::Build(CI, Buf.get(), Bounds,
      *DE, FS, PCH, true, PP)) {
    std::lock_guard<std::mutex> lock(mutex);
    preamble =
        std::make_shared<PreambleData>(std::move(*NewPreamble), DC.Take());
  }
}

} // namespace ccls

CompletionManager::CompletionManager(Project *project,
                                     WorkingFiles *working_files,
                                     OnDiagnostic on_diagnostic,
                                     OnDropped on_dropped)
    : project_(project), working_files_(working_files),
      on_diagnostic_(on_diagnostic), on_dropped_(on_dropped),
      preloaded_sessions_(kMaxPreloadedSessions),
      completion_sessions_(kMaxCompletionSessions),
      PCH(std::make_shared<PCHContainerOperations>()) {
  std::thread([&]() {
    set_thread_name("comp");
    ccls::CompletionMain(this);
  })
      .detach();
  std::thread([&]() {
    set_thread_name("comp-preload");
    ccls::CompletionPreloadMain(this);
  })
      .detach();
  std::thread([&]() {
    set_thread_name("diag");
    ccls::DiagnosticMain(this);
  })
      .detach();
}

void CompletionManager::CodeComplete(
    const lsRequestId &id,
    const lsTextDocumentPositionParams &completion_location,
    const OnComplete &on_complete) {
}

void CompletionManager::DiagnosticsUpdate(
    const lsTextDocumentIdentifier &document) {
  bool has = false;
  diagnostic_request_.Iterate([&](const DiagnosticRequest &request) {
    if (request.document.uri == document.uri)
      has = true;
  });
  if (!has)
    diagnostic_request_.PushBack(DiagnosticRequest{document},
                                 true /*priority*/);
}

void CompletionManager::NotifyView(const std::string &path) {
  // Only reparse the file if we create a new CompletionSession.
  if (EnsureCompletionOrCreatePreloadSession(path))
    preload_requests_.PushBack(PreloadRequest{path}, true);
}

void CompletionManager::NotifySave(const std::string &filename) {
  //
  // On save, always reparse.
  //

  EnsureCompletionOrCreatePreloadSession(filename);
  preload_requests_.PushBack(PreloadRequest{filename}, true);
}

void CompletionManager::NotifyClose(const std::string &filename) {
  //
  // On close, we clear any existing CompletionSession instance.
  //

  std::lock_guard<std::mutex> lock(sessions_lock_);

  // Take and drop. It's okay if we don't actually drop the file, it'll
  // eventually get pushed out of the caches as the user opens other files.
  auto preloaded_ptr = preloaded_sessions_.TryTake(filename);
  LOG_IF_S(INFO, !!preloaded_ptr)
      << "Dropped preloaded-based code completion session for " << filename;
  auto completion_ptr = completion_sessions_.TryTake(filename);
  LOG_IF_S(INFO, !!completion_ptr)
      << "Dropped completion-based code completion session for " << filename;

  // We should never have both a preloaded and completion session.
  assert((preloaded_ptr && completion_ptr) == false);
}

bool CompletionManager::EnsureCompletionOrCreatePreloadSession(
    const std::string &path) {
  std::lock_guard<std::mutex> lock(sessions_lock_);

  // Check for an existing CompletionSession.
  if (preloaded_sessions_.TryGet(path) ||
      completion_sessions_.TryGet(path)) {
    return false;
  }

  // No CompletionSession, create new one.
  auto session = std::make_shared<ccls::CompletionSession>(
      project_->FindCompilationEntryForFile(path), working_files_, PCH);
  preloaded_sessions_.Insert(session->file.filename, session);
  return true;
}

std::shared_ptr<ccls::CompletionSession>
CompletionManager::TryGetSession(const std::string &path,
                                 bool mark_as_completion,
                                 bool create_if_needed) {
  std::lock_guard<std::mutex> lock(sessions_lock_);

  // Try to find a preloaded session.
  std::shared_ptr<ccls::CompletionSession> preloaded =
      preloaded_sessions_.TryGet(path);

  if (preloaded) {
    // If this request is for a completion, we should move it to
    // |completion_sessions|.
    if (mark_as_completion) {
      preloaded_sessions_.TryTake(path);
      completion_sessions_.Insert(path, preloaded);
    }
    return preloaded;
  }

  // Try to find a completion session. If none create one.
  std::shared_ptr<ccls::CompletionSession> session =
      completion_sessions_.TryGet(path);
  if (!session && create_if_needed) {
    session = std::make_shared<ccls::CompletionSession>(
        project_->FindCompilationEntryForFile(path), working_files_, PCH);
    completion_sessions_.Insert(path, session);
  }

  return session;
}

void CompletionManager::FlushSession(const std::string &path) {
  std::lock_guard<std::mutex> lock(sessions_lock_);

  preloaded_sessions_.TryTake(path);
  completion_sessions_.TryTake(path);
}

void CompletionManager::FlushAllSessions() {
  LOG_S(INFO) << "flush all clang complete sessions";
  std::lock_guard<std::mutex> lock(sessions_lock_);

  preloaded_sessions_.Clear();
  completion_sessions_.Clear();
}