#include "compile/compilation.h"

#include "command/command.h"
#include "compile/diagnostic.h"
#include "compile/implement.h"
#include "semantic/ast_utility.h"
#include "support/logging.h"

#include "llvm/Support/Error.h"
#include "clang/Frontend/MultiplexConsumer.h"
#include "clang/Frontend/TextDiagnosticPrinter.h"
#include "clang/Lex/PreprocessorOptions.h"

namespace clice {

CompilationUnitRef::Self::~Self() {
    if(action) {
        // We already notified the pp of end-of-file earlier, so detach it first.
        // We must keep it alive until after EndSourceFile(), Sema relies on this.
        std::shared_ptr<clang::Preprocessor> pp = instance->getPreprocessorPtr();
        // Detach so we don't send EOF again
        instance->setPreprocessor(nullptr);
        action->EndSourceFile();
    }
}

std::unique_ptr<clang::CompilerInvocation>
    CompilationUnitRef::Self::create_invocation(this Self& self,
                                                CompilationParams& params,
                                                clang::DiagnosticConsumer* consumer) {
    if(params.arguments.empty()) {
        LOG_ERROR_RET(nullptr, "Fail to create invocation: empty argument list from database");
    }

    /// Temporary diagnostic engine, only used for command line parsing.
    /// For compilation, we need to create a new diagnostic engine. See also
    /// https://github.com/llvm/llvm-project/pull/139584#issuecomment-2920704282.
    clang::DiagnosticOptions options;
    llvm::IntrusiveRefCntPtr diagnostic_engine =
        clang::CompilerInstance::createDiagnostics(*params.vfs, options, consumer, false);
    if(!diagnostic_engine) {
        LOG_ERROR_RET(nullptr, "Fail to create diagnostics engine");
    }

    std::unique_ptr<clang::CompilerInvocation> invocation;

    /// If the second argument is "-cc1", the arguments are already expanded
    /// (e.g. from compilation database + query_toolchain). Skip driver and "-cc1"
    /// and create invocation directly from the cc1 args.
    bool is_cc1 = params.arguments.size() >= 2 && llvm::StringRef(params.arguments[1]) == "-cc1";
    if(is_cc1) {
        invocation = std::make_unique<clang::CompilerInvocation>();
        if(!clang::CompilerInvocation::CreateFromArgs(
               *invocation,
               llvm::ArrayRef(params.arguments).drop_front(2),
               *diagnostic_engine,
               params.arguments[0])) {
            LOG_ERROR_RET(nullptr,
                          " Fail to create invocation, arguments list is: {}",
                          print_argv(params.arguments));
        }
    } else {
        /// Create clang invocation.
        clang::CreateInvocationOptions options = {
            .Diags = diagnostic_engine,
            .VFS = params.vfs,

            /// Avoid replacing -include with -include-pch, also
            /// see https://github.com/clangd/clangd/issues/856.
            .ProbePrecompiled = false,
        };

        invocation = clang::createInvocation(params.arguments, options);
        if(!invocation) {
            LOG_ERROR_RET(nullptr,
                          " Fail to create invocation, arguments list is: {}",
                          print_argv(params.arguments));
        }
    }

    auto& pp_opts = invocation->getPreprocessorOpts();

    // CompilerInstance does not deterministically clear RetainRemappedFileBuffers,
    // especially if compilation aborts early, so we keep them alive and clean up
    // in CompilationUnit's destructor instead.
    pp_opts.RetainRemappedFileBuffers = true;

    for(auto& [file, buffer]: params.buffers) {
        pp_opts.addRemappedFile(file, buffer.get());
    }
    self.remapped_buffers = std::move(params.buffers);

    auto [pch, bound] = params.pch;
    pp_opts.ImplicitPCHInclude = std::move(pch);
    if(bound != 0) {
        pp_opts.PrecompiledPreambleBytes = {bound, false};
    }

    // We don't want to write comment locations into PCM. They are racy and slow
    // to read back. We rely on dynamic index for the comments instead.
    pp_opts.WriteCommentListToPCH = false;

    auto& header_search_opts = invocation->getHeaderSearchOpts();
    header_search_opts.Verbose = false;
    for(auto& [name, path]: params.pcms) {
        header_search_opts.PrebuiltModuleFiles.try_emplace(name.str(), std::move(path));
    }

    auto& front_opts = invocation->getFrontendOpts();
    front_opts.DisableFree = false;
    front_opts.ShowHelp = false;
    front_opts.ShowStats = false;
    front_opts.ShowVersion = false;
    front_opts.StatsFile = "";
    front_opts.TimeTracePath = "";
    front_opts.TimeTraceVerbose = false;
    front_opts.TimeTraceGranularity = 0;
    front_opts.PrintSupportedCPUs = false;
    front_opts.PrintEnabledExtensions = false;
    front_opts.PrintSupportedExtensions = false;

    /// Compiler flags (like gcc/clang's -M, -MD, -MMD, -H, or msvc's /showIncludes)
    /// can generate dependency files or print included headers to stdout/stderr.
    ///
    /// This output can interfere with or corrupt the Language Server Protocol (LSP)
    /// communication if the server is configured to use stdio for its JSON-RPC transport.
    /// We explicitly disables all related options to ensure no side-effect output is
    /// generated during parsing.
    auto& deps_opts = invocation->getDependencyOutputOpts();
    deps_opts.IncludeSystemHeaders = false;
    deps_opts.ShowSkippedHeaderIncludes = false;
    deps_opts.UsePhonyTargets = false;
    deps_opts.AddMissingHeaderDeps = false;
    deps_opts.IncludeModuleFiles = false;
    deps_opts.ShowIncludesDest = clang::ShowIncludesDestination::None;
    deps_opts.OutputFile.clear();
    deps_opts.HeaderIncludeOutputFile.clear();
    deps_opts.Targets.clear();
    deps_opts.ExtraDeps.clear();
    deps_opts.DOTOutputFile.clear();
    deps_opts.ModuleDependencyOutputDir.clear();

    auto& lang_opts = invocation->getLangOpts();
    lang_opts.CommentOpts.ParseAllComments = true;
    lang_opts.RetainCommentsFromSystemHeaders = true;

    return invocation;
}

void CompilationUnitRef::Self::configure_tidy(tidy::TidyParams tidy_params) {
    checker = tidy::configure(*instance, tidy_params);
}

void CompilationUnitRef::Self::run_tidy() {
    if(checker) {
        // AST traversals should exclude the preamble, to avoid performance cliffs.
        // TODO: is it okay to affect the unit-level traversal scope here?
        auto& Ctx = instance->getASTContext();
        Ctx.setTraversalScope(top_level_decls);
        checker->finder.matchAST(Ctx);

        /// XXX: This is messy: clang-tidy checks flush some diagnostics at EOF.
        /// However Action->EndSourceFile() would destroy the ASTContext!
        /// So just inform the preprocessor of EOF, while keeping everything alive.
        instance->getPreprocessor().EndSourceFile();
    }
}

namespace {

/// A wrapper ast consumer, so that we can cancel the ast parse
class ProxyASTConsumer final : public clang::MultiplexConsumer {
public:
    ProxyASTConsumer(std::unique_ptr<clang::ASTConsumer> consumer, CompilationUnitRef unit) :
        clang::MultiplexConsumer(std::move(consumer)), unit(unit) {}

    void collect_decl(clang::Decl* decl) {
        if(unit.file_id(unit.expansion_location(decl->getLocation())) != unit.interested_file()) {
            return;
        }

        if(const clang::NamedDecl* named_decl = dyn_cast<clang::NamedDecl>(decl)) {
            if(ast::is_implicit_template_instantiation(named_decl)) {
                return;
            }
        }

        unit->top_level_decls.push_back(decl);
    }

    auto HandleTopLevelDecl(clang::DeclGroupRef group) -> bool final {
        if(unit->kind == CompilationKind::Content) {
            if(group.isDeclGroup()) {
                for(auto decl: group) {
                    collect_decl(decl);
                }
            } else {
                collect_decl(group.getSingleDecl());
            }
        }

        /// TODO: check atomic variable after the parse of each declaration
        /// may result in performance issue, benchmark in the future.
        if(unit->stop && unit->stop->load()) {
            return false;
        }

        return clang::MultiplexConsumer::HandleTopLevelDecl(group);
    }

private:
    CompilationUnitRef unit;
};

class ProxyAction final : public clang::WrapperFrontendAction {
public:
    ProxyAction(std::unique_ptr<clang::FrontendAction> action, CompilationUnitRef unit) :
        clang::WrapperFrontendAction(std::move(action)), unit(unit) {}

    auto CreateASTConsumer(clang::CompilerInstance& instance, llvm::StringRef file)
        -> std::unique_ptr<clang::ASTConsumer> final {
        return std::make_unique<ProxyASTConsumer>(
            WrapperFrontendAction::CreateASTConsumer(instance, file),
            unit);
    }

    /// Make this public.
    using clang::WrapperFrontendAction::EndSourceFile;

private:
    CompilationUnitRef unit;
};

}  // namespace

CompilationStatus CompilationUnitRef::Self::run_clang(
    this Self& self,
    CompilationParams& params,
    std::unique_ptr<clang::FrontendAction> action,
    llvm::function_ref<void(clang::CompilerInstance&)> before_execute) {
    std::unique_ptr diagnostic_consumer = self.create_diagnostic();
    std::unique_ptr invocation = self.create_invocation(params, diagnostic_consumer.get());
    if(!invocation) {
        return CompilationStatus::SetupFail;
    }

    self.instance = std::make_unique<clang::CompilerInstance>(std::move(invocation));
    auto& instance = *self.instance;
    instance.createDiagnostics(*params.vfs, diagnostic_consumer.release(), true);

    if(auto remapping = clang::createVFSFromCompilerInvocation(instance.getInvocation(),
                                                               instance.getDiagnostics(),
                                                               params.vfs)) {
        instance.createFileManager(std::move(remapping));
    }

    if(!instance.createTarget()) {
        return CompilationStatus::SetupFail;
    }

    if(before_execute) {
        before_execute(instance);
    }

    self.action = std::make_unique<ProxyAction>(std::move(action), &self);

    if(!self.action->BeginSourceFile(instance, instance.getFrontendOpts().Inputs[0])) {
        /// If the action is not empty, we will call `EndSourceFile` at the destructor of `Self`.
        /// But if we fail to `BeginSourceFile` we don't need to call `EndSourceFile`. So just
        /// reset it.
        self.action.reset();
        return CompilationStatus::SetupFail;
    }

    /// FIXME: include-fixer, etc?

    /// Add PPCallbacks to collect preprocessing information.
    self.collect_directives();

    if(params.clang_tidy) {
        self.configure_tidy({});
    }

    std::optional<clang::syntax::TokenCollector> token_collector;
    if(!instance.hasCodeCompletionConsumer()) {
        /// It is not necessary to collect tokens if we are running code completion.
        /// And in fact will cause assertion failure.
        token_collector.emplace(instance.getPreprocessor());
    }

    if(auto error = self.action->Execute()) {
        // Upstream FrontendAction::Execute() always returns success (errors go through
        // diagnostics); log here only as a guard in case a custom action ever returns
        // an unexpected llvm::Error.
        LOG_ERROR("FrontendAction::Execute failed: {}", error);
        return CompilationStatus::FatalError;
    }

    /// If the output file is not empty, it represents that we are
    /// generating a PCH or PCM. If error occurs, the AST must be
    /// invalid to some extent, serialization of such AST may result
    /// in crash frequently. So forbidden it here and return as error.
    if(!instance.getFrontendOpts().OutputFile.empty() &&
       instance.getDiagnostics().hasErrorOccurred()) {
        return CompilationStatus::FatalError;
    }

    /// Check whether the compilation is canceled, if so we think
    /// it is an error.
    if(self.stop && self.stop->load()) {
        return CompilationStatus::Cancelled;
    }

    if(token_collector) {
        self.buffer = std::move(*token_collector).consume();
    }

    self.run_tidy();

    if(instance.hasSema()) {
        self.resolver.emplace(instance.getSema());
    }

    return CompilationStatus::Completed;
}

CompilationUnit run_clang(CompilationParams& params,
                          std::unique_ptr<clang::FrontendAction> action,
                          llvm::function_ref<void(clang::CompilerInstance&)> before_execute = {},
                          llvm::function_ref<void(CompilationUnitRef)> after_execute = {}) {
    LOG_DEBUG("Compile command: {}", print_argv(params.arguments));

    auto self = new CompilationUnitRef::Self();
    self->kind = params.kind;
    self->stop = std::move(params.stop);

    using namespace std::chrono;
    self->build_at = duration_cast<milliseconds>(system_clock::now().time_since_epoch());
    auto build_start = steady_clock::now().time_since_epoch();

    self->status = self->run_clang(params, std::move(action), before_execute);

    auto build_end = steady_clock::now().time_since_epoch();
    self->build_duration = duration_cast<milliseconds>(build_end - build_start);

    if(self->status == CompilationStatus::Completed && after_execute) {
        after_execute(self);
    }

    return CompilationUnit(self);
}

CompilationUnit preprocess(CompilationParams& params) {
    return run_clang(params, std::make_unique<clang::PreprocessOnlyAction>());
}

CompilationUnit compile(CompilationParams& params) {
    return run_clang(params,
                     std::make_unique<clang::SyntaxOnlyAction>(),
                     [](clang::CompilerInstance& instance) {
                         /// Make sure the output file is empty.
                         instance.getFrontendOpts().OutputFile.clear();
                     });
}

CompilationUnit compile(CompilationParams& params, PCHInfo& out) {
    assert(!params.output_file.empty() && "PCH file path cannot be empty");

    /// Record the begin time of PCH building.
    auto now = std::chrono::system_clock::now().time_since_epoch();
    out.mtime = std::chrono::duration_cast<std::chrono::milliseconds>(now).count();

    return run_clang(
        params,
        std::make_unique<clang::GeneratePCHAction>(),
        [&](clang::CompilerInstance& instance) {
            /// Set options to generate PCH.
            instance.getFrontendOpts().OutputFile = params.output_file.str();
            instance.getFrontendOpts().ProgramAction = clang::frontend::GeneratePCH;
            instance.getPreprocessorOpts().GeneratePreamble = true;

            // We don't want to write comment locations into PCH. They are racy and slow
            // to read back. We rely on dynamic index for the comments instead.
            instance.getPreprocessorOpts().WriteCommentListToPCH = false;

            instance.getLangOpts().CompilingPCH = true;
        },
        [&](CompilationUnitRef unit) {
            out.path = params.output_file.str();
            out.preamble = unit.interested_content();
            out.deps = unit.deps();
            out.arguments = params.arguments;
        });
}

CompilationUnit compile(CompilationParams& params, PCMInfo& out) {
    assert(!params.output_file.empty() && "PCM file path cannot be empty");

    return run_clang(
        params,
        std::make_unique<clang::GenerateReducedModuleInterfaceAction>(),
        [&](clang::CompilerInstance& instance) {
            /// Set options to generate PCH.
            instance.getFrontendOpts().OutputFile = params.output_file.str();
            instance.getFrontendOpts().ProgramAction =
                clang::frontend::GenerateReducedModuleInterface;

            out.srcPath = instance.getFrontendOpts().Inputs[0].getFile();
        },
        [&](CompilationUnitRef unit) {
            out.path = params.output_file.str();

            for(auto& [name, path]: params.pcms) {
                out.mods.emplace_back(name);
            }
        });
}

CompilationUnit complete(CompilationParams& params, clang::CodeCompleteConsumer* consumer) {
    auto& [file, offset] = params.completion;

    /// The location of clang is 1-1 based.
    std::uint32_t line = 1;
    std::uint32_t column = 1;

    /// FIXME:
    assert(params.buffers.size() == 1);
    llvm::StringRef content = params.buffers.begin()->second->getBuffer();

    for(auto c: content.substr(0, offset)) {
        if(c == '\n') {
            line += 1;
            column = 1;
            continue;
        }
        column += 1;
    }

    return run_clang(params,
                     std::make_unique<clang::SyntaxOnlyAction>(),
                     [&](clang::CompilerInstance& instance) {
                         /// Set options to run code completion.
                         instance.getFrontendOpts().CodeCompletionAt.FileName = std::move(file);
                         instance.getFrontendOpts().CodeCompletionAt.Line = line;
                         instance.getFrontendOpts().CodeCompletionAt.Column = column;
                         instance.setCodeCompletionConsumer(consumer);
                     });
}

}  // namespace clice
