#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 <limits>

#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 {

PimInstructionWriter::PimInstructionWriter(raw_pwrite_stream& stream)
: stream(stream), buffer(kBufferSize) {
  pim_binary::writeHeader(stream);
}

PimInstructionWriter::PimInstructionWriter(raw_fd_ostream& stream)
: PimInstructionWriter(static_cast<raw_pwrite_stream&>(stream)) {
  fileStream = &stream;
}

bool PimInstructionWriter::consumeStreamError() {
  if (!fileStream || !fileStream->has_error())
    return false;
  fileStream->clear_error();
  hasFailure = true;
  return true;
}

LogicalResult PimInstructionWriter::flushBuffer() {
  if (hasFailure)
    return failure();
  if (bufferedBytes != 0) {
    stream.write(buffer.data(), bufferedBytes);
    bufferedBytes = 0;
  }
  return consumeStreamError() ? failure() : success();
}

LogicalResult PimInstructionWriter::append(const pim_binary::InstructionRecord& record) {
  if (hasFailure || finalized || count == std::numeric_limits<uint32_t>::max()) {
    hasFailure = true;
    return failure();
  }

  if (bufferedBytes + pim_binary::kRecordSize > buffer.size() && failed(flushBuffer()))
    return failure();

  pim_binary::EncodedInstruction encoded = pim_binary::encodeInstructionRecord(record);
  std::memcpy(buffer.data() + bufferedBytes, encoded.data(), encoded.size());
  bufferedBytes += encoded.size();
  ++count;
  return success();
}

LogicalResult PimInstructionWriter::finalize() {
  if (finalized)
    return failure();
  finalized = true;
  if (hasFailure || failed(flushBuffer()))
    return failure();

  stream.flush();
  if (consumeStreamError())
    return failure();
  pim_binary::patchInstructionCount(stream, count);
  stream.flush();
  return consumeStreamError() ? failure() : success();
}

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(), std::nullopt});
    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["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