#include "tensorflow/compiler/mlir/tensorflow_to_stablehlo/python/pywrap_tensorflow_to_stablehlo_lib.h"

namespace py = pybind11;

namespace {

using mlir::tensorflow_to_stablehlo::pywrap::PywrapSavedModelToStablehlo;
using mlir::tensorflow_to_stablehlo::pywrap::PywrapTfModuleToStablehlo;

}  // namespace

PYBIND11_MODULE(pywrap_tensorflow_to_stablehlo, m) {
  m.doc() = "TensorFlow to StableHLO APIs.";

//
// Returns:
//   An absl::StatusOr containing the serialized bytecode of the StableHLO
//   module on success, or an error status on failure.
absl::StatusOr<std::string> PywrapSavedModelToStablehlo(
    absl::string_view input_path,
    const std::vector<std::string>& exported_model_signatures,
    const std::vector<std::string>& tag_names,
    absl::string_view input_arg_shapes_str);

  OpPrintingFlags printing_flags;
  module.print(os, printing_flags);

  output->os() << result;
  output->keep();

  return output_filename;
}

absl::StatusOr<std::string> PywrapSavedModelToStablehlo(
    absl::string_view input_path,
    const std::vector<std::string>& exported_model_signatures,
    const std::vector<std::string>& tag_names,
    absl::string_view input_arg_shapes_str) {