#pragma once

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

#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"

#include <cstdint>
#include <string>

#include "src/Accelerators/PIM/Common/IR/WeightUtils.hpp"

namespace onnx_mlir {

struct WeightFileRequest {
  size_t coreId = 0;
  llvm::SmallVector<ResolvedWeightView, 8> weights;
};

struct WeightEmissionResult {
  llvm::DenseMap<size_t, llvm::SmallVector<std::string, 8>> mapCoreWeightToFileName;
  uint64_t totalWeightBytes = 0;
};

WeightEmissionResult createAndPopulateWeightFolder(llvm::ArrayRef<WeightFileRequest> requests,
                                                   llvm::StringRef outputDirPath);

} // namespace onnx_mlir