#include <functional>
#include <string>

#include "tensorflow/compiler/mlir/tensorflow/utils/bridge_logger.h"

namespace tensorflow {

class DataDumperLoggerConfig : public ::tensorflow::BridgeLoggerConfig {
 public:
  explicit DataDumperLoggerConfig(
      std::function<std::string(const std::string &, mlir::Operation *op)>
          get_filename,

        get_filename,
    const std::string &pass_prefix, bool print_module_scope,
    bool print_after_only_on_change, mlir::OpPrintingFlags op_printing_flags)
    : ::tensorflow::BridgeLoggerConfig(
          print_module_scope, print_after_only_on_change, op_printing_flags),
      get_filename_(get_filename),
      pass_prefix_(pass_prefix) {}

void DataDumperLoggerConfig::printBeforeIfEnabled(

    // Print the whole module after each pass which requires disabling
    // multi-threading as well.
    module_op.getContext()->disableMultithreading();
    pm.enableIRPrinting(std::make_unique<tensorflow::BridgeLoggerConfig>(
        /*print_module_scope=*/true));
  }

  mlir::StatusScopedDiagnosticHandler diag_handler(module_op.getContext());
  if (failed(pm.run(module_op))) return diag_handler.ConsumeStatus();