#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"

namespace mlir {
namespace TF {

class TensorFlowRegistryEffectInterfaceFallback;

class TensorFlowDialect final : public Dialect {
 public:
  explicit TensorFlowDialect(MLIRContext *context);
  ~TensorFlowDialect() override;

  static StringRef getDialectNamespace() { return "tf"; }

  // Overrides to redirect to tf_type dialect.

  let description = [{
    Runs the pre calibration passes for post-training quantization with default
    configuration.
  }];
  let dependentDialects = [
    "mlir::stablehlo::StablehloDialect", "mlir::TF::TensorFlowDialect",
    "mlir::func::FuncDialect", "mlir::tf_executor::TensorFlowExecutorDialect",
    "mlir::mhlo::MhloDialect", "mlir::vhlo::VhloDialect",
  ];
}

  let description = [{
    Convert TF Quant ops to MHLO quant ops.
  }];

  let constructor = "mlir::quant::stablehlo::CreateConvertTFQuantOpsToMHLOPass()";
  let dependentDialects = ["TF::TensorFlowDialect", "chlo::ChloDialect",
                           "mhlo::MhloDialect", "tf_type::TFTypeDialect",
                           "quant::QuantizationDialect"];
}

using mlir::OperationState;
using mlir::UnknownLoc;
using mlir::chlo::ChloDialect;
using mlir::TF::TensorFlowDialect;
using tensorflow::tf2xla::internal::IsInBridgeAcceptableDialects;

class SharedUtilsTest : public ::testing::Test {};

TEST_F(SharedUtilsTest, IsInFunctionalDialectPasses) {
  MLIRContext context;
  context.loadDialect<TensorFlowDialect>();
  OpBuilder opBuilder(&context);

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;

 private:
  void runOnOperation() override {

    default_device_ = std::string(default_device);
  }

  void runOnOperation() override {
    Builder builder(&getContext());
    Dialect* tf = getContext().getLoadedDialect<TensorFlowDialect>();
    getOperation().walk([&](Operation* op) {
      if (auto device_attr = op->getAttrOfType<StringAttr>(kDeviceAttr)) {
        // We assign default device to ops with device attribute that is empty.

#define GET_OP_CLASSES
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/tf_quant_ops.cc.inc"

namespace mlir {
namespace quant {

void RegisterOpsHook(TF::TensorFlowDialect &dialect) {
  dialect.addOperations<
#define GET_OP_LIST
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/tf_quant_ops.cc.inc"
      >();
}

void RegisterOps() {

  if (candidate_op == nullptr) return false;

  // Don't prune if candidate op might have side effects.
  if (isa_and_nonnull<TF::TensorFlowDialect>(candidate_op->getDialect())) {
    if (TF::TensorFlowDialect::CanHaveSideEffects(candidate_op)) {
      return false;
    }
  } else if (!isMemoryEffectFree(candidate_op)) {
    return false;
  }

PYBIND11_MODULE(tfr_wrapper, m) {
  m.def("verify", [](std::string input) {
    mlir::DialectRegistry registry;
    registry.insert<mlir::arith::ArithDialect, mlir::scf::SCFDialect,
                    mlir::TF::TensorFlowDialect, mlir::func::FuncDialect,
                    mlir::shape::ShapeDialect, mlir::TFR::TFRDialect>();
    mlir::MLIRContext ctx(registry);
    ctx.loadAllAvailableDialects();

    tfg::ResultElementTypeFn result_element_type_fn,
    SmallVectorImpl<ShapedTypeComponents>& inferred_return_shapes) {
  assert(op->getName().getDialectNamespace() ==
         TensorFlowDialect::getDialectNamespace());
  return tfg::InferReturnTypeComponentsForTFOp(
      location, op, op->getOperands(), graph_version, operand_as_constant_fn,
      op_result_as_shape_fn, result_element_type_fn,