#pragma once

#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/Support/LogicalResult.h"

#include "llvm/ADT/StringRef.h"
#include "llvm/Support/JSON.h"
#include "llvm/Support/raw_ostream.h"

#include <cstddef>
#include <cstdint>
#include <vector>

#include "onnx-mlir/Compiler/OMCompilerTypes.h"
#include "src/Accelerators/PIM/Compiler/PimBinaryFormat.hpp"

namespace onnx_mlir {

class PimAcceleratorMemory;

class PimInstructionWriter {
public:
  explicit PimInstructionWriter(llvm::raw_pwrite_stream& stream);
  explicit PimInstructionWriter(llvm::raw_fd_ostream& stream);

  mlir::LogicalResult append(const pim_binary::InstructionRecord& record);
  mlir::LogicalResult finalize();

private:
  static constexpr size_t kBufferSize = 256 * 1024;
  mlir::LogicalResult flushBuffer();
  bool consumeStreamError();

  llvm::raw_pwrite_stream& stream;
  llvm::raw_fd_ostream* fileStream = nullptr;
  std::vector<char> buffer;
  size_t bufferedBytes = 0;
  uint32_t count = 0;
  bool finalized = false;
  bool hasFailure = false;
};

OnnxMlirCompilerErrorCodes writeMemoryBinary(mlir::ModuleOp moduleOp,
                                             mlir::func::FuncOp funcOp,
                                             PimAcceleratorMemory& memory,
                                             llvm::StringRef outputDirPath);
OnnxMlirCompilerErrorCodes writeConfigJson(mlir::func::FuncOp funcOp,
                                           PimAcceleratorMemory& memory,
                                           size_t maxCoreId,
                                           llvm::json::Object xbarsPerArrayGroup,
                                           llvm::StringRef outputDirPath);

} // namespace onnx_mlir