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

#include "revng/EarlyFunctionAnalysis/ControlFlowGraph.h"
#include "revng/EarlyFunctionAnalysis/ControlFlowGraphCache.h"
#include "revng/Lift/LoadBinaryPass.h"
#include "revng/Model/Binary.h"
#include "revng/Model/NameBuilder.h"
#include "revng/PTML/Constants.h"
#include "revng/PTML/Doxygen.h"
#include "revng/PTML/Tag.h"
#include "revng/Pipeline/AllRegistries.h"
#include "revng/Pipeline/Pipe.h"
#include "revng/Pipeline/RegisterPipe.h"
#include "revng/Pipes/Kinds.h"
#include "revng/Pipes/ModelGlobal.h"
#include "revng/Support/YAMLTraits.h"
#include "revng/Yield/Assembly/DisassemblyHelper.h"
#include "revng/Yield/Function.h"
#include "revng/Yield/PTML.h"
#include "revng/Yield/Pipes/ProcessAssembly.h"
#include "revng/Yield/Pipes/YieldAssembly.h"

namespace revng::pipes {

void ProcessAssembly::run(pipeline::ExecutionContext &Context,
                          const BinaryFileContainer &SourceBinary,
                          const CFGMap &CFGMap,
                          FunctionAssemblyStringMap &Output) {
  if (not SourceBinary.exists())
    return;

  // Access the model
  const auto &Model = getModelFromContext(Context);
  model::AssemblyNameBuilder NameBuilder = *Model;

  // Access the binary
  revng_assert(SourceBinary.path().has_value());
  auto BufferOrError = llvm::MemoryBuffer::getFileOrSTDIN(*SourceBinary.path());
  auto Buffer = cantFail(errorOrToExpected(std::move(BufferOrError)));
  RawBinaryView BinaryView(*Model, Buffer->getBuffer());

  // Define the helper object to store the disassembly pipeline.
  // This allows it to only be created once.
  DissassemblyHelper Helper;

  ControlFlowGraphCache Cache(CFGMap);

  for (const model::Function &Function :
       getFunctionsAndCommit(Context, Output.name())) {

    const auto &Metadata = Cache.getControlFlowGraph(Function.Entry());

    auto Disassembled = Helper.disassemble(Function,
                                           Metadata,
                                           BinaryView,
                                           *Model,
                                           NameBuilder);
    revng_assert(Disassembled.verify());
    Output.insert_or_assign(Function.Entry(), toString(Disassembled));
  }
}

} // namespace revng::pipes

namespace revng::pypeline::piperuns {

ProcessAssembly::ProcessAssembly(const class Model &Model,
                                 llvm::StringRef Config,
                                 llvm::StringRef DynamicConfig,
                                 const BinariesContainer &BinariesContainer,
                                 const CFGMap &CFG,
                                 AssemblyInternalContainer &Output) :
  Binary(*Model.get().get()), CFG(CFG), Output(Output), NameBuilder(Binary) {
  Helper = std::make_unique<DissassemblyHelper>();

  auto BinaryBuffer = BinariesContainer.getFile(0);
  BinaryView = std::make_unique<RawBinaryView>(Binary,
                                               llvm::StringRef{
                                                 BinaryBuffer.data(),
                                                 BinaryBuffer.size() });
};

ProcessAssembly::~ProcessAssembly() = default;

void ProcessAssembly::runOnFunction(const model::Function &TheFunction) {
  ObjectID Object(TheFunction.Entry());
  const auto &Metadata = CFG.getElement(Object);

  TupleTree<yield::Function> &OutputFunction = Output.getElement(Object);
  Helper->disassemble(TheFunction,
                      *Metadata,
                      *BinaryView,
                      Binary,
                      NameBuilder,
                      *OutputFunction);
}

} // namespace revng::pypeline::piperuns

namespace revng::pipes {

void YieldAssembly::run(pipeline::ExecutionContext &Context,
                        const FunctionAssemblyStringMap &Input,
                        FunctionAssemblyPTMLStringMap &Output) {
  // Access the model
  const auto &Model = getModelFromContext(Context);
  model::CNameBuilder NameBuilder = *Model;

  ptml::MarkupBuilder B;
  for (const model::Function &Function :
       getFunctionsAndCommit(Context, Output.name())) {
    MetaAddress Address = Function.Entry();
    llvm::StringRef YamlText = Input.at(Address);
    auto MaybeFunction = TupleTree<yield::Function>::fromString(YamlText);

    revng_assert(MaybeFunction && MaybeFunction->verify());
    revng_assert((*MaybeFunction)->verify());
    revng_assert((*MaybeFunction)->Entry() == Address);

    const model::Architecture::Values A = Model->Architecture();
    auto CommentIndicator = model::Architecture::getAssemblyCommentIndicator(A);

    const model::Configuration &Configuration = Model->Configuration();
    uint64_t LineWidth = Configuration.CommentLineWidth();

    std::string R = ptml::functionComment(B,
                                          Function,
                                          *Model,
                                          CommentIndicator,
                                          0,
                                          LineWidth,
                                          NameBuilder);
    R += yield::ptml::functionAssembly(B, **MaybeFunction, *Model);
    R = B.getTag(ptml::tags::Div, std::move(R)).toString();
    Output.insert_or_assign((*MaybeFunction)->Entry(), std::move(R));
  }
}

} // end namespace revng::pipes

using namespace revng::pipes;
using namespace pipeline;
static RegisterDefaultConstructibleContainer<FunctionAssemblyStringMap> X1;
static RegisterDefaultConstructibleContainer<FunctionAssemblyPTMLStringMap> X2;

static pipeline::RegisterPipe<revng::pipes::ProcessAssembly> ProcessPipe;
static pipeline::RegisterPipe<revng::pipes::YieldAssembly> YieldPipe;

namespace revng::pypeline::piperuns {

void YieldAssembly::runOnFunction(const model::Function &TheFunction) {
  MetaAddress Address = TheFunction.Entry();
  ObjectID Object(Address);
  const TupleTree<yield::Function> &Function = Input.getElement(Object);

  revng_assert(Function.verify());
  revng_assert(Function->verify());
  revng_assert(Function->Entry() == Address);

  const model::Architecture::Values A = Model.Architecture();
  auto CommentIndicator = model::Architecture::getAssemblyCommentIndicator(A);

  const model::Configuration &Configuration = Model.Configuration();
  uint64_t LineWidth = Configuration.CommentLineWidth();

  std::string R = ptml::functionComment(B,
                                        TheFunction,
                                        Model,
                                        CommentIndicator,
                                        0,
                                        LineWidth,
                                        NameBuilder);
  R += yield::ptml::functionAssembly(B, *Function, Model);
  R = B.getTag(ptml::tags::Div, std::move(R)).toString();
  *Output.getOStream(Object) << R;
}

} // namespace revng::pypeline::piperuns