#pragma once

#include <optional>

#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/STLFunctionalExtras.h"

#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/Value.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Support/LogicalResult.h"

#include "src/Accelerators/PIM/Common/PimCommon.hpp"

namespace onnx_mlir::spatial {
class SpatBlueprintOp;
}

namespace onnx_mlir {

mlir::FailureOr<mlir::IntegerAttr>
getTensorSizeInBytesAttr(mlir::Builder& builder, mlir::Operation* anchor, mlir::Value value);

template <class T>
size_t rangeLength(const mlir::iterator_range<T> range) {
  return std::distance(range.begin(), range.end());
}

/**
 * Retrieves the earliest operation that uses the given value within the value's
 * block.
 *
 * @param value The value for which to find the earliest user operation.
 * @return The earliest user operation that uses the given value within the
 * current block.
 */
mlir::Operation* getEarliestUserWithinBlock(mlir::Value value);

mlir::SmallVector<mlir::Value> getOpOperandsSortedByUses(mlir::Operation* operation);

mlir::Value getBestOutputTensorFromOperandsOrAllocate(mlir::RewriterBase& rewriter, mlir::Operation* operation);

mlir::LogicalResult validateFragmentAssemblyMetadata(onnx_mlir::spatial::SpatBlueprintOp blueprint,
                                                     int64_t resultRank,
                                                     size_t operandCount,
                                                     llvm::ArrayRef<int64_t> operandIndices,
                                                     llvm::ArrayRef<int64_t> sourceOffsets,
                                                     llvm::ArrayRef<int64_t> flatOffsets,
                                                     llvm::ArrayRef<int64_t> flatSizes,
                                                     llvm::ArrayRef<int64_t> flatStrides);

mlir::FailureOr<mlir::SmallVector<int64_t, 4>>
getStaticSliceOffsetsForElementOffset(mlir::Operation* anchor,
                                      mlir::ShapedType sourceType,
                                      llvm::ArrayRef<int64_t> fragmentShape,
                                      int64_t sourceElementOffset,
                                      llvm::StringRef fieldName);

mlir::LogicalResult
forEachContiguousDestinationChunk(llvm::ArrayRef<int64_t> destShape,
                                  llvm::ArrayRef<int64_t> baseOffsets,
                                  llvm::ArrayRef<int64_t> sizes,
                                  llvm::function_ref<mlir::LogicalResult(llvm::ArrayRef<int64_t>, int64_t, int64_t)>
                                    callback);

int64_t getFragmentAssemblySourceElementOffset(mlir::RankedTensorType sourceType,
                                               int64_t sourceSlot,
                                               int64_t sourceOffset);

struct FragmentAssemblyCopy {
  mlir::Value source;
  mlir::RankedTensorType sourceType;
  unsigned hostTargetIndex = 0;
  int64_t lane = 0;
  int64_t sourceByteOffset = 0;
  int64_t hostByteOffset = 0;
  int64_t byteSize = 0;
};

struct FragmentAssemblyCopyRun {
  mlir::Value source;
  mlir::RankedTensorType sourceType;
  unsigned hostTargetIndex = 0;
  int64_t count = 0;
  int64_t sourceStepBytes = 0;
  int64_t hostStepBytes = 0;
  int64_t byteSize = 0;
  mlir::SmallVector<int64_t, 8> sourceStartBytesByLane;
  mlir::SmallVector<int64_t, 8> hostStartBytesByLane;
};

mlir::FailureOr<mlir::SmallVector<FragmentAssemblyCopyRun, 8>>
groupFragmentAssemblyCopyRuns(llvm::ArrayRef<FragmentAssemblyCopy> copies, uint32_t laneCount = 1);

mlir::FailureOr<mlir::Value> emitFragmentAssemblyCopyRuns(mlir::IRRewriter& rewriter,
                                                          mlir::Location loc,
                                                          llvm::ArrayRef<FragmentAssemblyCopyRun> runs,
                                                          mlir::Value hostTarget,
                                                          mlir::Operation* anchor,
                                                          std::optional<mlir::Value> laneArg = std::nullopt,
                                                          mlir::Value baseHostOffset = {});

inline mlir::tensor::EmptyOp
createEmptyTensorFromShaped(mlir::IRRewriter& rewriter, mlir::Location loc, mlir::ShapedType shapedType) {
  return mlir::tensor::EmptyOp::create(rewriter, loc, shapedType.getShape(), shapedType.getElementType());
}

} // namespace onnx_mlir