#include "tensorflow/core/tfrt/runtime/runtime.h"

namespace tensorflow {

class BackendCompiler {
 public:
  virtual ~BackendCompiler();

  virtual void GetDependentDialects(mlir::DialectRegistry& registry) const {}

  // Compile the `module` in TF dialect. The result module should be also in TF
  // dialect.
  virtual absl::Status CompileTensorflow(

      llvm::cl::ZeroOrMore};

  void getDependentDialects(mlir::DialectRegistry& registry) const override {
    if (!module_) {
      for (const auto& device : device_specs_flag_) {
        auto* hardware = this->GetTargetHardware(device);
        if (hardware == nullptr) continue;
        hardware->GetDependentDialects(registry);
      }
    }
  }
};

  llvm::StringRef getArgument() const final {
    return "tfl-device-transform-nnapi";
  }
  llvm::StringRef getDescription() const final {
    return "Suitable transformation for nnapi only.";
  }
  void getDependentDialects(DialectRegistry& registry) const override {
    registry.insert<TF::TensorFlowDialect>();
  }
  void runOnOperation() override;
};

void DeviceTransformNNAPIPass::runOnOperation() {

namespace mlir {
namespace TFL {
namespace {
using tsl::StatusOr;

// MockSuccessPass reports errors but doesn't fail.
class MockSuccessPass
    : public PassWrapper<MockSuccessPass, OperationPass<ModuleOp>> {
  void getDependentDialects(DialectRegistry& registry) const override {
    registry.insert<TF::TensorFlowDialect>();
  }

 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MockSuccessPass)

  explicit MockSuccessPass() = default;

  }

  StringRef getDescription() const final {
    // This is a brief description of the pass.
    return "Convert Fake Quant op to quant.qcast and quant.dcast pairs";
  }

  void getDependentDialects(DialectRegistry& registry) const override {
    registry.insert<TF::TensorFlowDialect>();
    registry.insert<quant::QuantizationDialect>();
    registry.insert<quantfork::QuantizationForkDialect>();
  }

 public:
  StringRef getArgument() const final { return "smuggle-disallowed-ops-pass"; }
  StringRef getDescription() const final {
    return "Smuggle disallowed ops via stablehlo.custom_calls";
  }
  void getDependentDialects(DialectRegistry& registry) const final {
    registry.insert<mlir::stablehlo::StablehloDialect>();
  }

  void runOnOperation() override {
    RewritePatternSet patterns(&getContext());

  llvm::StringRef getArgument() const final {
    return "tfl-device-transform-gpu";
  }
  llvm::StringRef getDescription() const final {
    return "Suitable transformation for gpu only.";
  }
  void getDependentDialects(DialectRegistry& registry) const override {
    registry.insert<TF::TensorFlowDialect>();
  }
  void runOnOperation() override;
};

void DeviceTransformGPUPass::runOnOperation() {
  auto func = getOperation();

class XlaCallModuleDeserializationPass
    : public impl::XlaCallModuleDeserializationPassBase<
          XlaCallModuleDeserializationPass> {
 public:
  void getDependentDialects(mlir::DialectRegistry &registry) const override {
    XlaCallModuleDeserializationPassBase::getDependentDialects(registry);
    mlir::func::registerAllExtensions(registry);
  }

  void runOnOperation() override {
    ModuleOp module = getOperation();

  }

  StringRef getDescription() const final {
    // This is a brief description of the pass.
    return "Convert tf.CustomAggregator op to quant.Stats";
  }

  void getDependentDialects(DialectRegistry &registry) const override {
    registry.insert<TF::TensorFlowDialect>();
    registry.insert<quant::QuantizationDialect>();
    registry.insert<quantfork::QuantizationForkDialect>();
  }

struct LowerVariableOpsToMlProgramPass
    : public impl::LowerVariableOpsToMlProgramPassBase<
          LowerVariableOpsToMlProgramPass> {
  explicit LowerVariableOpsToMlProgramPass() = default;
  void getDependentDialects(DialectRegistry& registry) const override {
    registry.insert<mlir::tf_saved_model::TensorFlowSavedModelDialect,
                    ml_program::MLProgramDialect>();
  }
  void runOnOperation() override {