/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

/*!
 * \file codegen.cc
 * \brief Common utilities to generated C style code.
 */
#include <dmlc/memory_io.h>
#include <tvm/ir/module.h>
#include <tvm/runtime/c_runtime_api.h>
#include <tvm/runtime/container.h>
#include <tvm/runtime/module.h>
#include <tvm/runtime/registry.h>
#include <tvm/target/codegen.h>
#include <tvm/target/target.h>
#include <tvm/tir/function.h>
#include <tvm/tir/transform.h>

#include <cstdint>
#include <cstring>
#include <sstream>
#include <unordered_set>
#include <vector>

namespace tvm {
namespace codegen {

runtime::Module Build(IRModule mod, Target target) {
  if (transform::PassContext::Current()
          ->GetConfig<Bool>("tir.disable_assert", Bool(false))
          .value()) {
    mod = tir::transform::SkipAssert()(mod);
  }
  std::string build_f_name;
  if (target->kind->name == "micro_dev") {
    build_f_name = "target.build.c";
  } else {
    build_f_name = "target.build." + target->kind->name;
  }
  // the build function.
  const PackedFunc* bf = runtime::Registry::Get(build_f_name);
  CHECK(bf != nullptr) << build_f_name << " is not enabled";
  return (*bf)(mod, target);
}

/*! \brief Helper class to serialize module */
class ModuleSerializer {
 public:
  explicit ModuleSerializer(runtime::Module mod) : mod_(mod) { Init(); }

  void SerializeModule(dmlc::Stream* stream) {
    // Only have one DSO module and it is in the root, then
    // we will not produce import_tree_.
    bool has_import_tree = true;
    if (DSOExportable(mod_.operator->()) && mod_->imports().empty()) {
      has_import_tree = false;
    }
    uint64_t sz = 0;
    if (has_import_tree) {
      // we will append one key for _import_tree
      // The layout is the same as before: binary_size, key, logic, key, logic...
      sz = mod_vec_.size() + 1;
    } else {
      // Keep the old behaviour
      sz = mod_->imports().size();
    }
    stream->Write(sz);

    for (auto m : mod_vec_) {
      std::string mod_type_key = m->type_key();
      if (!DSOExportable(m)) {
        stream->Write(mod_type_key);
        m->SaveToBinary(stream);
      } else if (has_import_tree) {
        mod_type_key = "_lib";
        stream->Write(mod_type_key);
      }
    }

    // Write _import_tree key if we have
    if (has_import_tree) {
      std::string import_key = "_import_tree";
      stream->Write(import_key);
      stream->Write(import_tree_row_ptr_);
      stream->Write(import_tree_child_indices_);
    }
  }

 private:
  void Init() {
    CreateModuleIndex();
    CreateImportTree();
  }

  // invariance: root module is always at location 0.
  // The module order is collected via DFS
  void CreateModuleIndex() {
    std::unordered_set<const runtime::ModuleNode*> visited{mod_.operator->()};
    std::vector<runtime::ModuleNode*> stack{mod_.operator->()};
    uint64_t module_index = 0;

    while (!stack.empty()) {
      runtime::ModuleNode* n = stack.back();
      stack.pop_back();
      mod2index_[n] = module_index++;
      mod_vec_.emplace_back(n);
      for (runtime::Module m : n->imports()) {
        runtime::ModuleNode* next = m.operator->();
        if (visited.count(next) == 0) {
          visited.insert(next);
          stack.push_back(next);
        }
      }
    }
  }

  void CreateImportTree() {
    for (auto m : mod_vec_) {
      for (runtime::Module im : m->imports()) {
        uint64_t mod_index = mod2index_[im.operator->()];
        import_tree_child_indices_.push_back(mod_index);
      }
      import_tree_row_ptr_.push_back(import_tree_child_indices_.size());
    }
  }

  bool DSOExportable(const runtime::ModuleNode* mod) {
    return !std::strcmp(mod->type_key(), "llvm") || !std::strcmp(mod->type_key(), "c");
  }

  runtime::Module mod_;
  // construct module to index
  std::unordered_map<runtime::ModuleNode*, size_t> mod2index_;
  // index -> module
  std::vector<runtime::ModuleNode*> mod_vec_;
  std::vector<uint64_t> import_tree_row_ptr_{0};
  std::vector<uint64_t> import_tree_child_indices_;
};

namespace {
std::string SerializeModule(const runtime::Module& mod) {
  std::string bin;
  dmlc::MemoryStringStream ms(&bin);
  dmlc::Stream* stream = &ms;

  ModuleSerializer module_serializer(mod);
  module_serializer.SerializeModule(stream);

  return bin;
}
}  // namespace

std::string PackImportsToC(const runtime::Module& mod, bool system_lib) {
  std::string bin = SerializeModule(mod);

  // translate to C program
  std::ostringstream os;
  os << "#ifdef _WIN32\n"
     << "#define TVM_EXPORT __declspec(dllexport)\n"
     << "#else\n"
     << "#define TVM_EXPORT\n"
     << "#endif\n";
  os << "#ifdef __cplusplus\n"
     << "extern \"C\" {\n"
     << "#endif\n";
  os << "TVM_EXPORT extern const unsigned char " << runtime::symbol::tvm_dev_mblob << "[];\n";
  uint64_t nbytes = bin.length();
  os << "const unsigned char " << runtime::symbol::tvm_dev_mblob << "["
     << bin.length() + sizeof(nbytes) << "] = {\n  ";
  os << std::hex;
  size_t nunit = 80 / 4;
  for (size_t i = 0; i < sizeof(nbytes); ++i) {
    // sperators
    if (i != 0) {
      os << ",";
    }
    os << "0x" << ((nbytes >> (i * 8)) & 0xffUL);
  }
  for (size_t i = 0; i < bin.length(); ++i) {
    // sperators
    if ((i + sizeof(nbytes)) % nunit == 0) {
      os << ",\n  ";
    } else {
      os << ",";
    }
    int c = bin[i];
    os << "0x" << (c & 0xff);
  }
  os << "\n};\n";
  if (system_lib) {
    os << "extern int TVMBackendRegisterSystemLibSymbol(const char*, void*);\n";
    os << "static int " << runtime::symbol::tvm_dev_mblob << "_reg_ = "
       << "TVMBackendRegisterSystemLibSymbol(\"" << runtime::symbol::tvm_dev_mblob << "\", (void*)"
       << runtime::symbol::tvm_dev_mblob << ");\n";
  }
  os << "#ifdef __cplusplus\n"
     << "}\n"
     << "#endif\n";
  return os.str();
}

runtime::Module PackImportsToLLVM(const runtime::Module& mod, bool system_lib,
                                  const std::string& target_triple) {
  std::string bin = SerializeModule(mod);

  uint64_t nbytes = bin.length();
  std::string header;
  for (size_t i = 0; i < sizeof(nbytes); ++i) {
    header.push_back(((nbytes >> (i * 8)) & 0xffUL));
  }
  std::string blob = header + bin;
  TVMByteArray blob_byte_array;
  blob_byte_array.size = blob.length();
  blob_byte_array.data = blob.data();

  // Call codegen_blob to generate LLVM module
  std::string codegen_f_name = "codegen.codegen_blob";
  // the codegen function.
  const PackedFunc* codegen_f = runtime::Registry::Get(codegen_f_name);
  CHECK(codegen_f != nullptr) << "codegen.codegen_blob is not presented.";
  return (*codegen_f)(blob_byte_array, system_lib, target_triple);
}

TVM_REGISTER_GLOBAL("target.Build").set_body_typed(Build);

// Export two auxiliary function to the runtime namespace.
TVM_REGISTER_GLOBAL("runtime.ModulePackImportsToC").set_body_typed(PackImportsToC);

TVM_REGISTER_GLOBAL("runtime.ModulePackImportsToLLVM").set_body_typed(PackImportsToLLVM);

}  // namespace codegen
}  // namespace tvm