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

#include "ConstantUtils.hpp"
#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"

using namespace mlir;

namespace onnx_mlir {

Block* getHostConstantBlock(Operation* anchorOp) {
  assert(anchorOp && "expected a valid anchor operation");

  for (Operation* current = anchorOp; current; current = current->getParentOp())
    if (isa<spatial::SpatCompute, spatial::SpatComputeBatch, pim::PimCoreOp, pim::PimCoreBatchOp>(current))
      return current->getBlock();

  if (auto funcOp = anchorOp->getParentOfType<func::FuncOp>())
    return &funcOp.getBody().front();
  if (auto moduleOp = anchorOp->getParentOfType<ModuleOp>())
    return moduleOp.getBody();
  return anchorOp->getBlock();
}

Value getOrCreateHostConstant(Operation* anchorOp, Attribute value, Type type, OperationFolder& folder) {
  assert(anchorOp && "expected a valid anchor operation");
  Block* hostBlock = getHostConstantBlock(anchorOp);
  for (Operation& op : *hostBlock) {
    auto constantOp = dyn_cast<arith::ConstantOp>(&op);
    if (!constantOp || constantOp.getType() != type || constantOp.getValue() != value)
      continue;
    return constantOp.getResult();
  }

  auto* arithDialect = anchorOp->getContext()->getOrLoadDialect<arith::ArithDialect>();
  return folder.getOrCreateConstant(hostBlock, arithDialect, value, type);
}

Value getOrCreateHostConstantLike(arith::ConstantOp constantOp, OperationFolder& folder) {
  return getOrCreateHostConstant(constantOp.getOperation(), constantOp.getValue(), constantOp.getType(), folder);
}

Value getOrCreateHostIndexConstant(Operation* anchorOp, int64_t value, OperationFolder& folder) {
  Builder builder(anchorOp->getContext());
  return getOrCreateHostConstant(anchorOp, builder.getIndexAttr(value), builder.getIndexType(), folder);
}

Value getOrCreateHostI32Constant(Operation* anchorOp, int32_t value, OperationFolder& folder) {
  Builder builder(anchorOp->getContext());
  return getOrCreateHostConstant(anchorOp, builder.getI32IntegerAttr(value), builder.getI32Type(), folder);
}

Value getOrCreateHostI64Constant(Operation* anchorOp, int64_t value, OperationFolder& folder) {
  Builder builder(anchorOp->getContext());
  return getOrCreateHostConstant(anchorOp, builder.getI64IntegerAttr(value), builder.getI64Type(), folder);
}

} // namespace onnx_mlir