namespace mlir::quant::stablehlo {

using ::stablehlo::quantization::QuantizationConfig;
using ::tensorflow::MLIRImportOptions;
using ::tensorflow::SavedModelBundle;
using ::tensorflow::SavedModelSignatureDefsToMlirImport;
using ::tensorflow::quantization::PreprocessAndFreezeGraph;

absl::StatusOr<ImportedMlirModuleOp> SavedModelToMlirModuleOp(
    const absl::string_view saved_model_path,

// given MLIR `context`.
// 'saved_model_bundle' if not null, will be initialized with the model bundle.
absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>>
SavedModelSignatureDefsToMlirImport(
    absl::string_view saved_model_dir,
    const std::unordered_set<std::string>& tags,
    absl::Span<std::string> exported_names, mlir::MLIRContext* context,
    MLIRImportOptions options,

    mlir::MLIRContext context;
    tensorflow::MLIRImportOptions import_options;
    import_options.upgrade_legacy = upgrade_legacy;
    auto module_or = tensorflow::SavedModelSignatureDefsToMlirImport(
        input_filename, tags, exported_names, &context, import_options);
    if (!module_or.status().ok()) return 1;

    module_or.value()->print(output->os());

  if (!module_or.status().ok()) {
    LOG(ERROR) << "SavedModel import failed: " << module_or.status();
  }
  return module_or;
}

absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>>
SavedModelSignatureDefsToMlirImport(
    absl::string_view saved_model_dir,
    const std::unordered_set<std::string>& tags,
    absl::Span<std::string> exported_names, mlir::MLIRContext* context,
    MLIRImportOptions options,

  // TODO: b/213406917 - Add support for the object graph based saved model.
  std::vector<std::string> exported_names = signature_keys;
  absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> module =
      SavedModelSignatureDefsToMlirImport(saved_model_path, tags,
                                          absl::MakeSpan(exported_names),
                                          context, import_options, &bundle);
  if (!module.status().ok()) {

    options.upgrade_legacy = specs.upgrade_legacy;
    options.unconditionally_use_set_output_shapes = true;
    options.lift_variables = enable_variable_lifting;
    auto module_or = SavedModelSignatureDefsToMlirImport(
        input_filename, tags, exported_names, context, options,
        saved_model_bundle);

    if (!module_or.status().ok()) return module_or.status();
    return std::move(module_or).value();