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

#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/Support/FileSystem.h"
#include "llvm/Support/raw_ostream.h"

#include <algorithm>
#include <cassert>
#include <cstring>
#include <vector>

#include "src/Accelerators/PIM/Common/IR/WeightUtils.hpp"
#include "src/Accelerators/PIM/Compiler/PimArtifactWriter.hpp"
#include "src/Accelerators/PIM/Compiler/PimBinaryFormat.hpp"
#include "src/Accelerators/PIM/Compiler/PimCodeGen.hpp"
#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"

using namespace llvm;
using namespace mlir;

namespace onnx_mlir {

OnnxMlirCompilerErrorCodes writeHostCoreArtifacts(StringRef outputDirPath) {
  std::error_code errorCode;
  std::string outputHostCorePath = outputDirPath.str() + "/core_0.pim";
  raw_fd_ostream hostFileStream(outputHostCorePath, errorCode, sys::fs::OF_None);
  if (errorCode) {
    errs() << "Error while opening host core file `" << outputHostCorePath << "`: " << errorCode.message() << '\n';
    return InvalidOutputFileAccess;
  }

  pim_binary::writeHeader(hostFileStream);
  pim_binary::InstructionRecord noop;
  noop.opcode = pim_binary::Opcode::sldi;
  pim_binary::writeInstructionRecord(hostFileStream, noop);
  pim_binary::writeInstructionRecord(hostFileStream, noop);
  pim_binary::patchInstructionCount(hostFileStream, 2);
  hostFileStream.close();

  if (pimEmitJson.getValue()) {
    std::string outputHostJsonPath = outputDirPath.str() + "/core_0.json";
    raw_fd_ostream hostJsonStream(outputHostJsonPath, errorCode);
    if (errorCode) {
      errs() << "Error while opening host core json file `" << outputHostJsonPath << "`: " << errorCode.message()
             << '\n';
      return InvalidOutputFileAccess;
    }
    // The host core json contains two no-op-like instructions to satisfy pimsim-nn
    hostJsonStream << "[{\"imm\":0,\"op\":\"sldi\",\"rd\":0},{\"imm\":0,\"op\":\"sldi\",\"rd\":0}]";
    hostJsonStream.close();
  }
  return CompilerSuccess;
}

OnnxMlirCompilerErrorCodes
writeMemoryBinary(ModuleOp moduleOp, func::FuncOp funcOp, PimAcceleratorMemory& memory, StringRef outputDirPath) {
  auto memoryFilePath = (outputDirPath + "/memory.bin").str();
  std::error_code errorCode;
  raw_fd_ostream memoryFileStream(memoryFilePath, errorCode, sys::fs::OF_None);
  if (errorCode) {
    errs() << "Error while opening memory file " << memoryFilePath << ": " << errorCode.message() << '\n';
    return InvalidOutputFileAccess;
  }

  std::vector<char> memoryBuffer(memory.hostMem.getFirstAvailableAddress(), 0);

  SmallPtrSet<Operation*, 16> writtenGlobals;
  funcOp.walk([&](memref::GetGlobalOp getGlobalOp) {
    if (hasWeightAlways(getGlobalOp))
      return;
    auto globalOp = lookupGlobalForGetGlobal(moduleOp, getGlobalOp);
    if (!globalOp)
      return;
    if (!writtenGlobals.insert(globalOp.getOperation()).second)
      return;
    auto initialValue = globalOp.getInitialValue();
    if (!initialValue)
      return;
    auto denseAttr = dyn_cast<DenseElementsAttr>(*initialValue);
    if (!denseAttr)
      return;

    MemEntry memEntry = memory.hostMem.getMemEntry(getGlobalOp.getResult());
    ArrayRef<char> rawData = denseAttr.getRawData();
    char* dst = memoryBuffer.data() + memEntry.address;

    if (denseAttr.isSplat()) {
      size_t elementSize = rawData.size();
      assert(elementSize * getGlobalOp.getType().getNumElements() == memEntry.size && "Data size mismatch");
      for (size_t offset = 0; offset < memEntry.size; offset += elementSize)
        std::memcpy(dst + offset, rawData.data(), std::min(elementSize, memEntry.size - offset));
    }
    else {
      assert(rawData.size() == memEntry.size && "Data size mismatch");
      std::memcpy(dst, rawData.data(), rawData.size());
    }
  });

  memoryFileStream.write(memoryBuffer.data(), memoryBuffer.size());
  memoryFileStream.close();
  return CompilerSuccess;
}

OnnxMlirCompilerErrorCodes writeConfigJson(func::FuncOp funcOp,
                                           PimAcceleratorMemory& memory,
                                           size_t maxCoreId,
                                           json::Object xbarsPerArrayGroup,
                                           StringRef outputDirPath) {
  json::Object configJson;

  configJson["core_cnt"] = maxCoreId + 1;
  configJson["adc_count"] = 16;
  configJson["cell_precision"] = 2;
  configJson["xbar_array_count"] = crossbarCountInCore.getValue();
  configJson["xbar_size"] = {crossbarSize.getValue(), crossbarSize.getValue()};
  configJson["array_group_map"] = std::move(xbarsPerArrayGroup);

  json::Array inputsAddresses;
  for (BlockArgument input : funcOp.getArguments())
    inputsAddresses.push_back(memory.getValueAddress(input));
  configJson["inputs_addresses"] = std::move(inputsAddresses);

  json::Array outputsAddresses;
  for (func::ReturnOp returnOp : funcOp.getOps<func::ReturnOp>())
    for (mlir::Value output : returnOp.getOperands())
      outputsAddresses.push_back(memory.getValueAddress(output));
  configJson["outputs_addresses"] = std::move(outputsAddresses);

  auto configPath = (outputDirPath + "/config.json").str();
  std::error_code errorCode;
  raw_fd_ostream jsonOS(configPath, errorCode);
  if (errorCode) {
    errs() << "Error while opening config file: " << errorCode.message() << '\n';
    return InvalidOutputFileAccess;
  }
  jsonOS << json::Value(std::move(configJson)) << '\n';
  jsonOS.close();
  return CompilerSuccess;
}

} // namespace onnx_mlir