/*========================== begin_copyright_notice ============================

Copyright (C) 2022-2023 Intel Corporation

SPDX-License-Identifier: MIT

============================= end_copyright_notice ===========================*/
#include <igc/Options/Options.h>

#include <llvm/ADT/SmallVector.h>
#include <llvm/ADT/StringExtras.h>
#include <llvm/ADT/StringRef.h>
#include <llvm/ADT/Twine.h>
#include <llvm/IR/LLVMContext.h>
#include <llvm/IRReader/IRReader.h>
#include <llvm/Support/CommandLine.h>
#include <llvm/Support/FileSystem.h>
#include <llvm/Support/MD5.h>
#include <llvm/Support/MemoryBuffer.h>
#include <llvm/Support/ToolOutputFile.h>
#include <llvmWrapper/Support/TargetRegistry.h>
#include <llvmWrapper/Target/TargetMachine.h>

#include <llvm/Transforms/Utils/Cloning.h>

#include <Probe/Assertion.h>

using namespace llvm;

static cl::opt<std::string>
    SymbolPrefix("symb", cl::desc("prefix symbol for emmiting BiF functions "),
                 cl::value_desc("prefix"), cl::init("VCBuiltins64RawData"));

// Structure that holds list of platforms and an ordinal number of a
// corresponding unique BiF.
struct UniquePltf {
  size_t Num;
  std::vector<std::string> Platforms;
};

// Transient structure to hold platform name and path to the platform-specific
// emulation BiF.
struct PlatformModule {
  std::string Platform;
  std::string ModulePath;
  PlatformModule(const std::string &PlatformIn, const std::string &ModulePathIn)
      : Platform{PlatformIn}, ModulePath{ModulePathIn} {};
};

// Gets contents of file, reporting an error if problems are detected.
std::unique_ptr<MemoryBuffer> vcbGetFile(StringRef fileName) {
  auto FileOrError = MemoryBuffer::getFileOrSTDIN(fileName);
  if (!FileOrError) {
    errs() << FileOrError.getError().message();
    report_fatal_error("vcb: can't open file " + fileName);
  }
  return std::move(FileOrError.get());
}

// Generates C++ source file with a procedure that allows VC backend
// to select platform-specific emulation BiF module during the runtime.
// Arguments:
//  \p HashedUniquePltfs - mapping between bitcode and list of platforms
//                         for which this bitcode represents platform-specific
//                         emulation BiF.
//  \p Output - path to the output .cpp file.
//
// The generated code looks like this:
//
// ```
//      static unsigned char VCEmulation64RawDataPLTF0 = { ... }
//      static unsigned char VCEmulation64RawDataPLTF1 = { ... }
//
//      llvm::StringRef getVCEmulation64RawDataImpl(llvm::StringRef CPUStr) {
//        if (CPUStr.equals("PLTF_A") || ... )
//          return {reinterpret_cast<const char*>(VCEmulation64RawDataPLTF0),
//                  VCEmulation64RawDataPLTF0_size};
//        if (CPUStr.equals("PLTF_B") || ... )
//          return ...
//        ...
//      }
// ```
static std::string renderPlatformLiteral(StringRef Platform) {
  return (Twine("\"") + Platform + "\"").str();
};

void generateBifSelectionProcedure(
    std::map<std::string, UniquePltf> &HashedUniquePltfs, std::string &Output) {
  int FD;
  auto EC = llvm::sys::fs::openFileForWrite(Output, FD);
  if (EC)
    report_fatal_error(llvm::StringRef("vcb : can't open output file " + Output));
  raw_fd_ostream OS{FD, /*shouldClose=*/true};

  OS << "// This file is auto generated by vcb tool, DO NOT EDIT\n\n";

  OS << "#include \"IGC/common/StringMacros.hpp\"\n";
  OS << "#include \"llvm/ADT/StringRef.h\"\n";
  OS << "\n";

  // For each unique bitcode generate a C array that contains binary data
  // representing the bitcode
  for (const auto &[ByteCode, UniPltf] : HashedUniquePltfs) {
    std::string Pltf = SymbolPrefix + "PLTF" + std::to_string(UniPltf.Num);
    OS << "static unsigned char " << Pltf << "[] = {";
    bool FirstIn = true;
    for (size_t i = 0; i < ByteCode.size(); i++) {
      if (!FirstIn)
        OS << ",";
      FirstIn = false;
      uint8_t Num = ByteCode[i];
      OS << " 0x" << utohexstr(Num);
    }

    OS << "\n    };\n\n"
       << "unsigned int " << Pltf << "_size = " << ByteCode.size() << ";\n\n";
  }
  OS << "llvm::StringRef get" << SymbolPrefix
     << "Impl(llvm::StringRef CPUStr) {\n";

  // Generate a selection procedure that for each supported platform
  // (taken from configuration file) selects a BLOB that represents a
  // platform-specific emulation BiF corresponding to that platform.
  for (const auto &[ByteCode, UniPltf] : HashedUniquePltfs) {
    const auto &PltfList = UniPltf.Platforms;
    std::vector<std::string> PlatformCompareExpressions;
    llvm::transform(PltfList, std::back_inserter(PlatformCompareExpressions),
                    [](const auto &Pltf) {
                      return (Twine("CPUStr.equals(") +
                              renderPlatformLiteral(Pltf) + ")")
                          .str();
                    });
    OS << "  if (" << llvm::join(PlatformCompareExpressions, "\n    || ")
       << ")\n"
       << "      return  {reinterpret_cast<const char*>(" << SymbolPrefix
       << "PLTF" << UniPltf.Num << "),\n"
       << "              " << SymbolPrefix << "PLTF" << UniPltf.Num
       << "_size};\n";
  }
  OS << "return \"\";\n";
  OS << "};\n\n";
  OS.close();
}

// Parses input configuration file (see \fn vcbCompileUnique for the format)
// and returns a list of PlatformModule objects.
std::vector<PlatformModule> parseResponseFile(std::string &InputFilename) {
  std::vector<PlatformModule> RetVal;
  std::unique_ptr<MemoryBuffer> File = vcbGetFile(InputFilename);

  auto [Paths, Platforms] = File->getBuffer().split("\n");
  SmallVector<StringRef, 50> VecPaths;
  SmallVector<StringRef, 50> VecPlatforms;
  Platforms.split(VecPlatforms, ";");
  Paths.split(VecPaths, ";");
  if (VecPlatforms.size() != VecPaths.size())
    report_fatal_error("vcb: incorrect Responce file - the number of "
                       "Platforms and Paths is different");
  auto ZippedVecPlatformsPaths = llvm::zip(VecPlatforms, VecPaths);
  llvm::transform(ZippedVecPlatformsPaths, std::back_inserter(RetVal),
                  [](const auto &ZippedPlatformArg) {
                    const auto &[Platform, Path] = ZippedPlatformArg;
                    return PlatformModule(Platform.trim().str(),
                                          Path.trim().str());
                  });

  return RetVal;
}

// Generates .cpp code that is used by VC backend to obtain platform-specific
// precompiled emulation BiF during runtime. One of the primary entry points
// to vcb tool.
//
// Bitcodes corresponding to different architectures that have the same
// binary representation are coalesced to reduce the size.
//
// The generated function has the following protorype:
//      llvm::StringRef getVCEmulation64RawDataImpl(llvm::StringRef CPUStr);
//
// Arguments:
//    \p InputFilename - path to configifuration file that contains information
//                       where to search for platform-specific precompiled BiF.
//    \p OutputFilename - name of the output .cpp file.
//
// Note:
//    Configuration file is expected to have two lines:
//    1. <BiF_Path1;Bif_Path2;Bif_Path3;...> - colon-separted paths to each
//        platform-specific BiF.
//    2. <Plafrom1;Platform2;...> - colon-separated platform names
//        corresponding to files listed in the previous line.
void vcbCompileUnique(std::string InputFilename, std::string OutputFilename) {
  // Parse configuration.
  std::map<std::string, UniquePltf> HashedUniquePltfs;
  std::vector<PlatformModule> VecPltfModule = parseResponseFile(InputFilename);

  // Create a mapping between coalesced BiF binary data and corresponding
  // platforms.
  for (const auto &[Platform, ModulePath] : VecPltfModule) {
    std::string CompiledBytecode = vcbGetFile(ModulePath)->getBuffer().str();
    if (HashedUniquePltfs.find(CompiledBytecode) == HashedUniquePltfs.end()) {
      HashedUniquePltfs[std::move(CompiledBytecode)] =
          UniquePltf{HashedUniquePltfs.size(), {Platform}};
    } else {
      HashedUniquePltfs[std::move(CompiledBytecode)].Platforms.push_back(
          Platform);
    }
  }
  // Finally, produce a resulting .cpp file.
  generateBifSelectionProcedure(HashedUniquePltfs, OutputFilename);
}
