#include "tsl/platform/errors.h"
#include "tsl/platform/statusor.h"

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(

    }
  } else {
    tensor_names->insert(tensor_info.name());
  }
}

// Gets the union of all inputs and outputs of all SignatureDefs in the bundle
void GetSignatureDefsInputsAndOutputs(
    const SavedModelBundle& saved_model_bundle,
    std::unordered_set<string>* inputs, std::unordered_set<string>* outputs) {
  for (auto& sigdef_elem : saved_model_bundle.meta_graph_def.signature_def()) {

// XlaCallModuleOp. Returns a non-OK status if quantization fails, or any of
// `saved_model_bundle` or `quantization_py_function_lib` is a nullptr.
absl::StatusOr<mlir::ModuleOp> RunQuantization(
    const SavedModelBundle* saved_model_bundle,
    absl::string_view saved_model_dir,
    const std::unordered_set<std::string>& saved_model_tags,
    const stablehlo::quantization::QuantizationConfig& quantization_config,

                      const RunOptions& run_options, const string& export_dir,
                      const std::unordered_set<string>& tags,
                      SavedModelBundle* const bundle) {
  return LoadSavedModelGeneric<SavedModelBundle>(session_options, run_options,
                                                 export_dir, tags, bundle);
}

Status RestoreSession(const RunOptions& run_options,

    mlir::OwningOpRef<mlir::ModuleOp> module,
    const mlir::TFL::PassConfig& pass_config,
    const std::unordered_set<std::string>& saved_model_tags, string* result,
    SavedModelBundle* saved_model_bundle,
    const quantization::PyFunctionLibrary* quantization_py_function_lib);

// Give a warning for any unused flags that have been specified.
void WarningUnusedFlags(const toco::ModelFlags& model_flags,

    bool include_variables_in_initializers, bool upgrade_legacy,
    bool show_debug_info, TF_Status* status) {
  // Load the saved model into a SavedModelBundle.

  std::unordered_set<string> tag_set =
      absl::StrSplit(tags, ',', absl::SkipEmpty());

  tensorflow::SavedModelBundle bundle;
  auto load_status =
      tensorflow::LoadSavedModel({}, {}, saved_model_path, tag_set, &bundle);
  if (!load_status.ok()) {

namespace odml {

absl::StatusOr<OwningOpRef<mlir::ModuleOp>> ImportSavedModelOrMLIR(
    const std::string& input_path, MLIRContext* context,
    llvm::SourceMgr* source_mgr,
    std::unique_ptr<tensorflow::SavedModelBundle>* saved_model_bundle) {
  if (absl::EndsWith(input_path, ".mlir")) {
    auto file_or_err = llvm::MemoryBuffer::getFileOrSTDIN(input_path.c_str());
    if (std::error_code error = file_or_err.getError()) {

// and represents a computation defined by a TF subgraph.
// These Signatures were primarily designed to be interoperable with the legacy
// TF 1 Session-based C++ SavedModelBundle loading APIs:
// https://github.com/tensorflow/tensorflow/blob/26c4ee0c833e74f94d0102d8b005c41a28b44445/tensorflow/cc/saved_model/loader.h#L96-L108
// SignatureDefFunctions have different semantics from regular TF2

#include "tensorflow/core/public/session.h"
#include "tsl/platform/errors.h"
#include "tsl/platform/statusor.h"

namespace mlir {
namespace {

// Extract the mlir TF module and optionally a ::tensorflow::SavedModelBundle
// from a saved model or from an mlir file.
absl::StatusOr<quant::stablehlo::ImportedMlirModuleOp> ImportSavedModelOrTfMlir(
    absl::string_view input_path, MLIRContext* context,

    const absl::string_view saved_model_dir,
    const std::unordered_set<std::string>& saved_model_tags,
    QuantizationConfig* quantization_config,
    const PyFunctionLibrary* quantization_py_function_lib,
    const SavedModelBundle* saved_model_bundle, mlir::PassManager& pass_manager,
    mlir::StatusScopedDiagnosticHandler& status_handler, ModuleOp& module) {
  // TODO: b/194747383 - We need to valid that indeed the "main" func is
  // presented.