#include "llvm/Support/raw_os_ostream.h"

#include <fstream>

#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp"
#include "src/Accelerators/PIM/Common/Support/FileSystemUtils.hpp"

namespace onnx_mlir {

std::fstream openDialectDumpFileWithExtension(const std::string& name, llvm::StringRef destination,  llvm::StringRef extension) {
  std::string outputDir = getOutputDir();
  if (outputDir.empty())
    return {};

  std::string dialectsDir = (outputDir + destination).str();
  createDirectory(dialectsDir);
  return std::fstream(dialectsDir + "/" + name + "." + extension.str(), std::ios::out);
}

void dumpModule(mlir::ModuleOp moduleOp, const std::string& name, bool assumeVerified) {
  std::fstream file = openDialectDumpFileWithExtension(name, "/dialects", "mlir");
  if (!file.is_open())
    return;

  llvm::raw_os_ostream os(file);
  mlir::OpPrintingFlags flags;
  flags.elideLargeElementsAttrs().enableDebugInfo(false, false);
  if (assumeVerified)
    flags.assumeVerified();
  moduleOp.print(os, flags);
  os.flush();
  file.close();
}

} // namespace onnx_mlir