//
// This file is distributed under the MIT License. See LICENSE.md for details.
//

#include "mlir/IR/Attributes.h"
#include "mlir/IR/BuiltinOps.h"

#include "revng/CliftEmitC/CEmitter.h"
#include "revng/CliftEmitC/Headers.h"
#include "revng/CliftEmitC/TypeDefinitionEmitter.h"
#include "revng/CliftEmitC/TypeDependencyGraph.h"
#include "revng/PTML/CTokenEmitter.h"
#include "revng/Pipeline/Location.h"
#include "revng/Pipes/Ranks.h"

class CHeaderEmitterImpl : TypeDefinitionEmitter {

public:
  explicit CHeaderEmitterImpl(ptml::CTokenEmitter &Tokens,
                              const CDataModel &DataModel,
                              TypeEmitterConfiguration Configuration) :
    TypeDefinitionEmitter(Tokens, DataModel, Configuration) {}

public:
  void emitHeaderPrologue() {
    // TODO: emit header location *definition*, so that ctrl+click on includes
    //       (references) leads to the correct file.

    Tokens.emitPragmaOnceDirective();
    Tokens.emitNewline();

    emitCategoryComment("This header has been generated using rev.ng.");

    // TODO: emit the license information, revng version information, etc.
  }

  void emitCommonIncludes() {
    // TODO: attach proper header *reference* locations.

    Tokens.emitIncludeDirective("attributes.h",
                                "",
                                ptml::CTokenEmitter::IncludeMode::Quote);
    Tokens.emitIncludeDirective("primitive-types.h",
                                "",
                                ptml::CTokenEmitter::IncludeMode::Quote);
    Tokens.emitNewline();
  }

private:
  void emitTypeGraph(const TypeDependencyGraph &Graph) {
    // In order to improve the printing order, do the visit it in two parts:
    // first only start from nodes without any successors (real roots),
    // only then, resolve potential loops by starting from arbitrary nodes.
    std::unordered_set<const TypeDependencyNode *> Emitted;
    for (const auto *Root : Graph.nodes())
      if (not Root->predecessorCount())
        emitTypeTree(*Root, Emitted);
    for (const auto *Root : Graph.nodes())
      emitTypeTree(*Root, Emitted);

    revng_assert(Graph.size() == Emitted.size());
  }

public:
  void emitTypes(mlir::ModuleOp Module) {
    auto Graph = TypeDependencyGraph::makeModelGraph(Module);
    if (not Graph.empty()) {
      emitCategoryComment("Types");
      emitTypeGraph(Graph);
      Tokens.emitNewline();
    }
  }

private:
  template<pipeline::RankSpecialization RankFilter>
  void emitFunctionsImpl(mlir::ModuleOp Module,
                         const RankFilter &Rank,
                         llvm::StringRef CategoryComment) {
    bool CommentEmitted = false;
    Module->walk([&](clift::FunctionOp Function) {
      if (pipeline::locationFromString(Rank, Function.getHandle())) {
        if (not CommentEmitted) {
          emitCategoryComment(CategoryComment);
          CommentEmitted = true;
        }

        emitFunctionPrototype(Function);
        Tokens.emitPunctuator(ptml::CTokenEmitter::Punctuator::Semicolon);
        Tokens.emitNewline();
        Tokens.emitNewline();
      }
    });
  }

public:
  void emitFunctions(mlir::ModuleOp Module) {
    emitFunctionsImpl(Module, revng::ranks::Function, "Functions");
  }

  void emitDynamicFunctions(mlir::ModuleOp Module) {
    emitFunctionsImpl(Module,
                      revng::ranks::DynamicFunction,
                      "Imported Dynamic Functions");
  }

public:
  void emitSegments(mlir::ModuleOp Module) {
    bool CommentEmitted = false;
    Module->walk([this, &CommentEmitted](clift::GlobalVariableOp Segment) {
      auto MaybeLocation = pipeline::locationFromString(revng::ranks::Segment,
                                                        Segment.getHandle());
      revng_assert(MaybeLocation.has_value());

      if (not CommentEmitted) {
        emitCategoryComment("Segments");
        CommentEmitted = true;
      }

      emitGlobalDoxygenComment(Segment);

      using EntityKind = ptml::CTokenEmitter::EntityKind;
      emitDeclaration(Segment.getType(),
                      CEmitter::DeclaratorInfo{
                        .Identifier = Segment.getName(),
                        .Location = Segment.getHandle(),
                        .CAttributes = {},
                        .Kind = EntityKind::GlobalVariable,
                        .Parameters = {} });

      Tokens.emitPunctuator(ptml::CTokenEmitter::Punctuator::Semicolon);
      Tokens.emitNewline();
      Tokens.emitNewline();
    });
  }

public:
  void emitHelpers(llvm::ArrayRef<mlir::ModuleOp> Modules) {
    // TODO: emit `#include`s

    auto Graph = TypeDependencyGraph::makeHelperGraph(Modules);
    if (not Graph.empty()) {
      emitCategoryComment("Types");

      revng_assert(!Modules.empty());
      emitTypeGraph(Graph);
    }

    bool CommentEmitted = false;
    std::unordered_set<std::string_view> EmittedFunctions;
    for (mlir::ModuleOp Module : Modules) {
      Module->walk([this,
                    &CommentEmitted,
                    &EmittedFunctions](clift::FunctionOp Function) {
        if (EmittedFunctions.contains(Function.getHandle()))
          return;

        if (pipeline::locationFromString(revng::ranks::HelperFunction,
                                         Function.getHandle())) {
          if (not CommentEmitted) {
            emitCategoryComment("Functions");
            CommentEmitted = true;
          }

          emitFunctionPrototype(Function);
          Tokens.emitPunctuator(ptml::CTokenEmitter::Punctuator::Semicolon);
          Tokens.emitNewline();
          Tokens.emitNewline();

          auto [_, Success] = EmittedFunctions.emplace(Function.getHandle());
          revng_assert(Success);
        }
      });
    }

    Tokens.emitNewline();
  }
};

void emitTypeAndGlobalHeader(ptml::CTokenEmitter &Tokens,
                             mlir::ModuleOp Module,
                             TypeEmitterConfiguration Configuration) {
  CHeaderEmitterImpl Emitter(Tokens,
                             clift::getDataModel(Module),
                             Configuration);

  Emitter.emitHeaderPrologue();
  Emitter.emitCommonIncludes();

  // TODO: split the following into separate headers.
  Emitter.emitTypes(Module);
  Emitter.emitFunctions(Module);
  Emitter.emitDynamicFunctions(Module);
  Emitter.emitSegments(Module);
}

void emitHelperHeader(ptml::CTokenEmitter &Tokens,
                      llvm::ArrayRef<mlir::ModuleOp> Modules) {
  revng_check(not Modules.empty());

  const CDataModel &DataModel = clift::getDataModel(Modules.front());
  revng_assert(llvm::all_of(Modules.drop_front(),
                            [&DataModel](mlir::ModuleOp Module) {
                              return clift::getDataModel(Module) == DataModel;
                            }));

  TypeEmitterConfiguration Configuration = {
    .ExplicitPadding = true,
  };

  CHeaderEmitterImpl Emitter(Tokens, DataModel, Configuration);

  Emitter.emitHeaderPrologue();
  Emitter.emitHelpers(Modules);
}