#include "tensorflow/compiler/mlir/op_or_arg_name_mapper.h"
#include "tensorflow/lite/toco/toco_flags.pb.h"

namespace tflite {
// Options for exporting to Flatbuffer.
struct FlatbufferExportOptions {
  // TocoFlags proto. The following fields are migrated.
  // bool emit_builtin_tflite_ops  -> !toco_flags.force_select_tf_ops()
  // bool emit_select_tf_ops       -> toco_flags.enable_select_tf_ops()

    return absl::InternalError(absl::StrCat(
        "Failed to sparsify: ", statusHandler.ConsumeStatus().message()));
  }

  // Export the results to the builder
  std::string result;
  tflite::FlatbufferExportOptions options;
  options.toco_flags.set_force_select_tf_ops(false);
  options.toco_flags.set_enable_select_tf_ops(true);
  options.toco_flags.set_allow_custom_ops(true);

  }

  std::string result;
  if (output_mlir) {
    llvm::raw_string_ostream os(result);
    module.print(os);
    os.flush();
  } else {
    tflite::FlatbufferExportOptions options;
    options.toco_flags.set_force_select_tf_ops(false);
    options.toco_flags.set_allow_custom_ops(true);
    if (enable_select_tf_ops) {
      options.toco_flags.set_enable_select_tf_ops(true);

    LOG(ERROR) << "Failed to quantize: " << err;
    return kTfLiteError;
  }

  // Export the results to the builder
  std::string result;
  tflite::FlatbufferExportOptions options;
  options.toco_flags.set_force_select_tf_ops(false);
  options.toco_flags.set_enable_select_tf_ops(true);
  options.toco_flags.set_allow_custom_ops(true);

    const std::string err(statusHandler.ConsumeStatus().message());
    LOG(ERROR) << "Failed to quantize: " << err;
    return kTfLiteError;
  }

  // Export the results.
  tflite::FlatbufferExportOptions options;
  options.toco_flags.set_force_select_tf_ops(false);
  options.toco_flags.set_enable_select_tf_ops(true);
  options.toco_flags.set_allow_custom_ops(true);

  if (main.getFunctionType().getNumInputs() != 0)
    LOG(QFATAL) << "NYI: Only nullary functions supported.";

  // Convert to flatbuffer.
  std::string serialized_flatbuffer;
  tflite::FlatbufferExportOptions options;
  options.toco_flags.set_force_select_tf_ops(!emit_builtin_tflite_ops);
  options.toco_flags.set_enable_select_tf_ops(emit_select_tf_ops);
  options.toco_flags.set_allow_custom_ops(emit_custom_ops);

    op_or_arg_name_mapper =
        std::make_unique<tensorflow::OpOrArgStripNameMapper>();
  } else {
    op_or_arg_name_mapper =
        std::make_unique<tensorflow::OpOrArgLocNameMapper>();
  }
  tflite::FlatbufferExportOptions options;
  options.toco_flags.set_force_select_tf_ops(!emit_builtin_tflite_ops);
  options.toco_flags.set_enable_select_tf_ops(emit_select_tf_ops);
  options.toco_flags.set_allow_custom_ops(emit_custom_ops);

    return status_handler.Combine(
        absl::InvalidArgumentError("VHLO lowering failed"));
  }

  // Write MLIR Stablehlo dialect into FlatBuffer
  OpOrArgLocNameMapper op_or_arg_name_mapper;
  tflite::FlatbufferExportOptions options;
  options.toco_flags = toco_flags;
  options.saved_model_tags = saved_model_tags;
  options.op_or_arg_name_mapper = &op_or_arg_name_mapper;
  options.metadata[tflite::kModelUseStablehloTensorKey] = "true";

  }
  return control_edges;
}

}  // namespace

namespace tflite {

bool MlirToFlatBufferTranslateFunction(mlir::ModuleOp module,
                                       const FlatbufferExportOptions& options,
                                       std::string* serialized_flatbuffer,
                                       bool serialize_stablehlo_ops) {
  auto maybe_translated = Translator::Translate(