#pragma once

#include "mlir/Dialect/MemRef/IR/MemRef.h"

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

#include <limits>
#include <optional>
#include <string>

#include "src/Accelerators/PIM/Compiler/PimCodeGen.hpp"
#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"

namespace onnx_mlir {

struct LocalAllocInterval {
  size_t id = 0;
  mlir::memref::AllocOp alloc;
  MemoryValueKey key;
  uint64_t start = 0;
  uint64_t end = 0;
  size_t size = 0;
  mlir::Operation *startOp = nullptr;
  mlir::Operation *endOp = nullptr;
  mlir::Operation *firstTouchOp = nullptr;
  mlir::Operation *lastTouchOp = nullptr;
  uint64_t firstTouchPosition = 0;
  uint64_t lastTouchPosition = 0;
  bool startUsedAllocFallback = false;
  bool endUsedFallback = false;
  bool hasRuntimeUse = false;
  bool insideNestedRegion = false;
  bool escapesLoop = false;
  std::string fallbackReason;
  llvm::SmallVector<std::string, 8> aliasesFollowed;
  size_t slotPlanIndex = std::numeric_limits<size_t>::max();
  size_t physicalSlotId = std::numeric_limits<size_t>::max();
  size_t assignedAddress = 0;
  size_t physicalSlotSize = 0;
};

struct PlannedPhysicalSlot {
  size_t id = std::numeric_limits<size_t>::max();
  size_t requiredSize = 0;
  size_t size = 0;
  size_t address = 0;
  llvm::SmallVector<size_t, 8> intervalIndices;
};

llvm::SmallVector<LocalAllocInterval, 0> buildLocalAllocIntervals(mlir::Operation *coreLikeOp,
                                                                  std::optional<unsigned> lane);

llvm::SmallVector<PlannedPhysicalSlot, 0> planPhysicalSlots(llvm::MutableArrayRef<LocalAllocInterval> intervals);

MemoryPlanArtifacts buildMemoryPlanArtifacts(mlir::Operation *coreLikeOp,
                                             std::optional<unsigned> lane,
                                             llvm::ArrayRef<LocalAllocInterval> intervals,
                                             llvm::ArrayRef<PlannedPhysicalSlot> slots,
                                             size_t addressLimit,
                                             PimMemoryReportLevel reportLevel);

} // namespace onnx_mlir