#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"

#include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/HostFoldability.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/HostLegality.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"

using namespace mlir;

namespace onnx_mlir {

LogicalResult verifyONNXToSpatialHostLegality(func::FuncOp funcOp) {
  pim::CappedDiagnosticReporter diagnostics;

  for (Operation& op : funcOp.getFunctionBody().front()) {
    if (isa<func::ReturnOp, spatial::SpatCompute, spatial::SpatComputeBatch>(&op))
      continue;
    if (isHostFoldableOp(&op))
      continue;

    diagnostics.report(&op, [](Operation* illegalOp) {
      illegalOp->emitOpError("non-foldable top-level runtime op remains after ONNX-to-Spatial; lower it inside "
                             "spat.compute");
    });
  }

  diagnostics.emitSuppressedSummary(funcOp, "ONNX-to-Spatial host legality failures");

  return success(!diagnostics.hasFailure());
}

} // namespace onnx_mlir