#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* getConstantInsertionBlock(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 = dyn_cast<func::FuncOp>(anchorOp))
    return &funcOp.getBody().front();
  if (auto moduleOp = dyn_cast<ModuleOp>(anchorOp))
    return moduleOp.getBody();
  if (auto funcOp = anchorOp->getParentOfType<func::FuncOp>())
    return &funcOp.getBody().front();
  if (auto moduleOp = anchorOp->getParentOfType<ModuleOp>())
    return moduleOp.getBody();
  return anchorOp->getBlock();
}

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

  OpBuilder::InsertionGuard guard(rewriter);
  rewriter.setInsertionPointToStart(hostBlock);
  return arith::ConstantOp::create(rewriter, anchorOp->getLoc(), type, cast<TypedAttr>(value)).getResult();
}

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

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

Value getOrCreateIndexConstant(RewriterBase& rewriter, Operation* anchorOp, int64_t value) {
  Builder builder(anchorOp->getContext());
  return getOrCreateConstant(rewriter, anchorOp, builder.getIndexAttr(value), builder.getIndexType());
}

std::optional<int64_t> matchConstantIndexValue(Value value) {
  if (!value || !value.getType().isIndex())
    return std::nullopt;

  if (auto constant = value.getDefiningOp<arith::ConstantIndexOp>())
    return constant.value();

  if (auto constant = value.getDefiningOp<arith::ConstantOp>())
    if (auto intAttr = dyn_cast<IntegerAttr>(constant.getValue()); intAttr && intAttr.getType().isIndex())
      return intAttr.getInt();

  return std::nullopt;
}

std::optional<int64_t> matchConstantIndexValue(OpFoldResult value) {
  if (auto attr = dyn_cast<Attribute>(value))
    if (auto intAttr = dyn_cast<IntegerAttr>(attr); intAttr && intAttr.getType().isIndex())
      return intAttr.getInt();
  if (auto operand = dyn_cast<Value>(value))
    return matchConstantIndexValue(operand);
  return std::nullopt;
}

} // namespace onnx_mlir