/*
 * 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_c_host.cc
 */
#include "codegen_c_host.h"
#include "support/check.h"
#include <tvm/ir/cast.h>
#include <tvm/runtime/logging.h>

#include <tvm/ffi/extra/module.h>
#include <tvm/target/codegen.h>

#include <algorithm>
#include <string>
#include <unordered_set>
#include <utility>
#include <vector>

// For escaping strings embedded into generated C sources
#include "support/str_escape.h"

namespace tvm {
namespace tl {

using namespace ffi;

CodeGenCHost::CodeGenCHost() {
  module_name_ = name_supply_->FreshName(symbol::tvm_ffi_library_ctx);
}

void CodeGenCHost::Init(bool output_ssa, bool emit_asserts,
                        bool emit_fwd_func_decl, std::string target_str,
                        const std::unordered_set<std::string> &devices) {
  emit_asserts_ = emit_asserts;
  emit_fwd_func_decl_ = emit_fwd_func_decl;
  declared_globals_.clear();
  decl_stream << "// tilelang target: " << target_str << "\n";
  decl_stream << "#define TVM_EXPORTS\n";
  decl_stream << "#include \"tvm/runtime/base.h\"\n";
  decl_stream << "#include \"tvm/runtime/c_backend_api.h\"\n";
  decl_stream << "#include \"tvm/ffi/c_api.h\"\n";
  decl_stream << "#include <math.h>\n";
  // snprintf for richer assert messages with actual values
  decl_stream << "#include <stdio.h>\n";
  decl_stream << "#include <stdbool.h>\n";
  // Portable alignment specifier: MSVC uses __declspec(align(N)), every other
  // compiler nvcc invokes (gcc/clang) accepts __attribute__((aligned(N))).
  decl_stream << "#if defined(_MSC_VER)\n";
  decl_stream << "#define TL_ALIGN(N) __declspec(align(N))\n";
  decl_stream << "#else\n";
  decl_stream << "#define TL_ALIGN(N) __attribute__((aligned(N)))\n";
  decl_stream << "#endif\n";

  decl_stream << "#ifdef __OBJC__\n";
  decl_stream << "#include \"tvm/runtime/device_api.h\"\n";
  decl_stream << "#include \"tvm/ffi/function.h\"\n";

  decl_stream << "#include <Metal/Metal.h>\n";
  decl_stream << "#include <Foundation/Foundation.h>\n";

  decl_stream << "#include <torch/mps.h>\n";
  decl_stream << "#endif\n";

  CodeGenCHost::InitGlobalContext();
  tvm::codegen::CodeGenC::Init(output_ssa);
}

void CodeGenCHost::InitGlobalContext() {
  decl_stream << "void* " << symbol::tvm_ffi_library_ctx << " = NULL;\n";
}

void CodeGenCHost::DefineModuleName() {
  decl_stream << "void* " << module_name_ << " = NULL;\n";
}

void CodeGenCHost::AddFunction(const tvm::GlobalVar &gvar,
                               const tvm::tirx::PrimFunc &func) {
  return AddFunction(gvar, func, /*emit_fwd_func_decl=*/false);
}

void CodeGenCHost::AddFunction(const tvm::GlobalVar &gvar,
                               const tvm::tirx::PrimFunc &func,
                               bool emit_fwd_func_decl) {
  auto global_symbol = func->GetAttr<String>(tvm::attr::kGlobalSymbol);
  if (global_symbol) {
    function_names_.push_back(global_symbol.value());
  }

  emit_fwd_func_decl_ = emit_fwd_func_decl;
  tvm::codegen::CodeGenC::AddFunction(gvar, func);
  if (func->HasNonzeroAttr(tvm::tirx::attr::kIsEntryFunc) && !has_main_func_) {
    ICHECK(global_symbol.has_value())
        << "CodeGenCHost: The entry func must have the global_symbol "
           "attribute, "
        << "but function " << gvar << " only has attributes " << func->attrs;
    if (global_symbol.value() != symbol::tvm_ffi_main) {
      function_names_.push_back(symbol::tvm_ffi_main);
      stream << "// CodegenC: NOTE: Auto-generated entry function\n";
      PrintFuncPrefix(stream);
      PrintType(func->ret_type, stream);
      stream << " " << symbol::tvm_ffi_main
             << "(void* self, void* args,int num_args, void* result) {\n";
      stream << "  return " << static_cast<std::string>(global_symbol.value())
             << "(self, args, num_args, result);\n";
      stream << "}\n";
    }
    has_main_func_ = true;
  }
}

void CodeGenCHost::GenerateForwardFunctionDeclarations(
    String global_symbol, const Array<tvm::Type> &arg_types,
    const tvm::Type &ret_type) {
  if (!emit_fwd_func_decl_) {
    return;
  }
  for (auto &func_already_defined : GetFunctionNames()) {
    if (global_symbol == func_already_defined) {
      return;
    }
  }
  this->PrintFuncPrefix(fwd_decl_stream);
  this->PrintType(ret_type, fwd_decl_stream);
  fwd_decl_stream << " " << global_symbol << "(";
  for (size_t i = 0; i < arg_types.size(); ++i) {
    if (i > 0) {
      fwd_decl_stream << ", ";
    }
    tvm::codegen::CodeGenSourceBase::PrintType(arg_types[i], fwd_decl_stream);
  }
  fwd_decl_stream << ");\n";
}

void CodeGenCHost::PrintFuncPrefix(std::ostream &os) { // NOLINT(*)
  os << "#ifdef __cplusplus\n"
     << "extern \"C\"\n"
     << "#endif\n";
}

void CodeGenCHost::PrintType(tvm::DataType t, std::ostream &os) { // NOLINT(*)
  int lanes = t.lanes();
  if (t.is_handle()) {
    ICHECK_EQ(lanes, 1) << "does not support vector types";
    os << "void*";
    return;
  }
  if (t.is_void()) {
    os << "void";
    return;
  }
  if (t == tvm::DataType::Bool()) {
    os << "bool";
    return;
  }
  bool fail = false;
  if (t.is_float()) {
    switch (t.bits()) {
    case 16:
      os << "half";
      break;
    case 32:
      os << "float";
      break;
    case 64:
      os << "double";
      break;
    default:
      fail = true;
      break;
    }
    if (!fail && lanes == 1)
      return;
    if (!fail && (lanes >= 2 && lanes <= 16)) {
      os << lanes;
      return;
    }
  }
  if (t.is_bfloat16()) {
    os << "__bf16";
    return;
  }
  if (t.is_int() || t.is_uint()) {
    if (t.is_uint()) {
      os << 'u';
    }
    switch (t.bits()) {
    case 8:
      os << "int8_t";
      break;
    case 16:
      os << "int16_t";
      break;
    case 32:
      os << "int32_t";
      break;
    case 64:
      os << "int64_t";
      break;
    case 1:
      os << "int32_t";
      break;
    default:
      fail = true;
      break;
    }
    if (!fail && lanes == 1)
      return;
    if (!fail && (lanes >= 2 && lanes <= 16)) {
      os << lanes;
      return;
    }
  }
  LOG(FATAL) << "Cannot convert type " << t << " to C type";
}

void CodeGenCHost::VisitExpr_(const tvm::tirx::BroadcastNode *op,
                              std::ostream &os) { // NOLINT(*)
  std::string v = PrintExpr(op->value);
  int lanes = op->dtype.lanes();
  os << "((";
  PrintType(op->dtype, os);
  os << ")(";
  for (int i = 0; i < lanes; ++i) {
    if (i != 0)
      os << ", ";
    os << v;
  }
  os << "))";
}

void CodeGenCHost::PrintGetFuncFromBackend(
    const std::string &func_name, const std::string &packed_func_name) {
  this->PrintIndent();
  this->stream << "if (" << packed_func_name << " == NULL) {\n";
  int packed_func_if_scope = this->BeginScope();
  this->PrintIndent();
  this->stream << "if (TVMBackendGetFuncFromEnv(" << module_name_ << ", \""
               << func_name << "\""
               << ", &" << packed_func_name << ") != 0) {\n";
  int get_func_env_scope = this->BeginScope();
  this->PrintIndent();
  this->stream << "return -1;\n";
  this->EndScope(get_func_env_scope);
  this->PrintIndent();
  this->stream << "}\n";
  this->EndScope(packed_func_if_scope);
  this->PrintIndent();
  this->stream << "}\n";
}

void CodeGenCHost::PrintCallPacked(const tvm::tirx::CallNode *op) {
  using namespace tvm::tirx;
  const StringImmNode *func_name = op->args[0].as<StringImmNode>();
  ICHECK(func_name != nullptr)
      << "tvm_call_[c]packed_lowered expects first argument as function name";
  int64_t begin = op->args[2].as<IntImmNode>()->value;
  int64_t end = op->args[3].as<IntImmNode>()->value;
  int64_t num_args = end - begin;
  ICHECK_GE(num_args, 0);

  std::string packed_func_name;
  if (op->op.same_as(builtin::tvm_call_packed_lowered())) {
    packed_func_name = GetPackedName(op);
    this->PrintGetFuncFromBackend(func_name->value, packed_func_name);
  } else {
    // directly use the original symbol
    ICHECK(op->op.same_as(builtin::tvm_call_cpacked_lowered()));
    packed_func_name = symbol::tvm_ffi_symbol_prefix + func_name->value;
  }

  std::string args_stack = PrintExpr(op->args[1]);
  this->PrintIndent();
  std::string result = name_supply_->FreshName("result");
  if (is_in_metal_context) {
    this->stream << "__block ";
  }
  this->stream << "TVMFFIAny " << result << ";\n";
  this->PrintIndent();
  // must make sure type_index is set to none
  this->stream << result << ".type_index = kTVMFFINone;\n";
  this->PrintIndent();
  this->stream << result << ".zero_padding = 0;\n";
  this->PrintIndent();
  this->stream << result << ".v_int64 = 0;\n";

  int metal_scope;
  std::string metal_result;
  if (is_in_metal_context) {
    metal_result = name_supply_->FreshName("metal_ret");
    this->PrintLine("__block int ", metal_result, " = 0;");
    this->PrintLine("auto serialQueue = torch::mps::get_dispatch_queue();");
    this->PrintLine("dispatch_sync(serialQueue, ^() {");
    metal_scope = this->BeginScope();

    this->PrintLine("const id<MTLCommandBuffer> commandBuffer = "
                    "torch::mps::get_command_buffer();");
    this->PrintLine(
        "const auto f = tvm::ffi::Function::GetGlobal(\"metal.SetStream\");");
    this->PrintLine("(*f)(static_cast<TVMStreamHandle>(commandBuffer));");
  }

  this->PrintIndent();
  if (op->op.same_as(builtin::tvm_call_packed_lowered())) {
    this->stream << "if (TVMFFIFunctionCall(" << packed_func_name << ", ";
  } else {
    this->stream << "if (" << packed_func_name << "(NULL, ";
  }
  this->stream << "(TVMFFIAny*) " << args_stack << ", " << num_args << ", "
               << "&" << result << ") != 0) {\n";
  int func_call_scope = this->BeginScope();
  if (is_in_metal_context) {
    this->PrintLine(metal_result, " = -1;");
  } else {
    this->PrintLine("return -1;");
  }
  this->EndScope(func_call_scope);

  this->PrintLine("}");

  if (is_in_metal_context) {
    this->EndScope(metal_scope);
    this->PrintLine("});");
    this->PrintLine("if (", metal_result, " != 0) return ", metal_result, ";");
  }
}

std::string CodeGenCHost::GetPackedName(const tvm::tirx::CallNode *op) {
  using namespace tvm::tirx;
  const StringImmNode *s = op->args[0].as<StringImmNode>();
  ICHECK(s != nullptr)
      << "tvm_call_packed_lowered expects first argument as function name";
  std::string func_name = s->value;
  std::string packed_func_name = func_name + "_packed";
  std::string unique_name;
  auto it = declared_globals_.find(packed_func_name);
  if (it != declared_globals_.end()) {
    unique_name = it->second;
  } else {
    unique_name = name_supply_->FreshName(packed_func_name);
    declared_globals_[packed_func_name] = unique_name;
    decl_stream << "static void* " << unique_name << " = NULL;\n";
  }
  return unique_name;
}

void CodeGenCHost::VisitExpr_(const tvm::tirx::CallNode *op,
                              std::ostream &os) { // NOLINT(*)
  using namespace tvm::tirx;
  if (op->op.same_as(builtin::tvm_stack_alloca())) {
    std::string stack_name = name_supply_->FreshName("stack");
    const std::string &type = op->args[0].as<StringImmNode>()->value;
    const IntImmNode *num = op->args[1].as<IntImmNode>();
    ICHECK(num != nullptr);
    static_assert(alignof(TVMFFIAny) % alignof(DLTensor) == 0, "invariant");
    size_t unit = sizeof(TVMFFIAny);
    size_t size = 0;
    if (type == "shape") {
      size = (num->value * sizeof(Shape::index_type) + unit - 1) / unit;
    } else if (type == "tvm_ffi_any") {
      size = (num->value * sizeof(TVMFFIAny) + unit - 1) / unit;
    } else if (type == "array") {
      size = (num->value * sizeof(DLTensor) + unit - 1) / unit;
    } else {
      LOG(FATAL) << "Unknown stack alloca type " << type;
    }
    this->PrintIndent();
    if (type == "tvm_ffi_any") {
      // TMA descriptors are materialized through tvm_ffi_any stack allocas and
      // are passed by value to the kernel as ``__grid_constant__ const
      // CUtensorMap``.
      this->stream << "TL_ALIGN(128) TVMFFIAny " << stack_name << "[" << size
                   << "];\n";
    } else {
      this->stream << "TVMFFIAny " << stack_name << "[" << size << "];\n";
    }
    os << stack_name;
  } else if (op->op.same_as(builtin::tvm_call_packed_lowered())) {
    this->PrintCallPacked(op);
  } else if (op->op.same_as(builtin::tvm_call_cpacked_lowered())) {
    this->PrintCallPacked(op);
  } else if (op->op.same_as(builtin::tvm_throw_last_error())) {
    this->PrintIndent();
    this->stream << "return -1;\n";
  } else {
    tvm::codegen::CodeGenC::VisitExpr_(op, os);
  }
}

void CodeGenCHost::VisitStmt_(
    const tvm::tirx::AssertStmtNode *op) { // NOLINT(*)
  if (emit_asserts_) {
    std::string cond = PrintExpr(op->condition);
    PrintIndent();
    stream << "if (!(" << cond << ")) {\n";
    int assert_if_scope = this->BeginScope();
    {
      // Join message_parts into a single string
      std::string joined_msg;
      bool all_literal = true;
      for (const auto &part : op->message_parts) {
        joined_msg += part->value;
      }
      if (!joined_msg.empty()) {
        std::string esc_msg = tvm::support::StrEscape(
            joined_msg.c_str(), joined_msg.length(), /*use_octal_escape=*/true,
            /*escape_whitespace_special_chars=*/true);

        // Only print expected/got values for equality when message is available
        if (const auto *eq = op->condition.as<tvm::tirx::EQNode>()) {
          std::string lhs = PrintExpr(eq->a);
          std::string rhs = PrintExpr(eq->b);
          PrintIndent();
          stream << "char __tvm_assert_msg_buf[512];\n";
          PrintIndent();
          stream << "snprintf(__tvm_assert_msg_buf, 512, \"%s; expected: %lld, "
                    "got: %lld\", \""
                 << esc_msg << "\", (long long)(" << lhs << "), (long long)("
                 << rhs << "));\n";
          PrintIndent();
          stream << "TVMFFIErrorSetRaisedFromCStr(\"RuntimeError\", "
                    "__tvm_assert_msg_buf);\n";
        } else {
          PrintIndent();
          stream << "TVMFFIErrorSetRaisedFromCStr(\"RuntimeError\", \""
                 << esc_msg << "\");\n";
        }
      } else {
        PrintIndent();
        stream << "TVMFFIErrorSetRaisedFromCStr(\"RuntimeError\", "
                  "\"assertion failed\");\n";
      }
    }
    PrintIndent();
    stream << "return -1;\n";
    this->EndScope(assert_if_scope);
    PrintIndent();
    stream << "}\n";
  }
}

void CodeGenCHost::VisitStmt_(const tvm::tirx::AttrStmtNode *op) {
  bool enter_metal_ctx = op->attr_key == "metal_context";
  if (enter_metal_ctx) {
    ICHECK(!is_in_metal_context) << "Nested metal context";
    is_in_metal_context = true;
  }
  tvm::codegen::CodeGenC::VisitStmt_(op);
  if (enter_metal_ctx) {
    is_in_metal_context = false;
  }
}

void CodeGenCHost::VisitExpr_(const tvm::tirx::MinNode *op,
                              std::ostream &os) { // NOLINT(*)
  PrintTernaryCondExpr(op, "<", os);
}

void CodeGenCHost::VisitExpr_(const tvm::tirx::MaxNode *op,
                              std::ostream &os) { // NOLINT(*)
  PrintTernaryCondExpr(op, ">", os);
}

template <typename T>
inline void CodeGenCHost::PrintTernaryCondExpr(const T *op, const char *compare,
                                               std::ostream &os) { // NOLINT(*)
  std::ostringstream temp_a;
  VisitExpr(op->a, temp_a);
  std::string a_id = SSAGetID(temp_a.str(), op->a.dtype());
  std::ostringstream temp_b;
  VisitExpr(op->b, temp_b);
  std::string b_id = SSAGetID(temp_b.str(), op->b.dtype());

  os << "((" << a_id << ") " << compare << " (" << b_id << ") "
     << "? (" << a_id << ") : (" << b_id << "))";
}

} // namespace tl
} // namespace tvm

namespace tvm {
namespace tl {

using tvm::codegen::CodeGenSourceBase;
using tvm::codegen::CSourceModuleCreate;

// Build function that mirrors TVM's host C codegen, registered under a
// TileLang-specific name.
// NOTE(chaofan): Different from codegen_c / BuildTileLangC, this CodeGen class
// is only used to generate C host code for TileLang when the host codegen is
// enabled. If you use TileLang to generate CPU code (in this case, C is the
// device code) , it will be generated by BuildTileLangC.
Module BuildTileLangCHost(IRModule mod, Target target) {
  bool output_ssa = false;
  bool emit_asserts = true;
  bool emit_fwd_func_decl = true;

  std::unordered_set<std::string> devices;
  if (mod->GetAttr<Map<GlobalVar, String>>("device_contexts") != nullptr) {
    Map<GlobalVar, String> device_contexts =
        mod->GetAttr<Map<GlobalVar, String>>("device_contexts").value();
    for (auto const &context : device_contexts) {
      devices.insert(context.second.data());
    }
  }

  CodeGenCHost cg;
  cg.Init(output_ssa, emit_asserts, emit_fwd_func_decl, target->str(), devices);
  cg.SetConstantsByteAlignment(
      target->GetAttr<::tvm::Integer>("constants-byte-alignment").value_or(16));

  auto is_aot_executor_fn = [](::tvm::tirx::PrimFunc const &func) -> bool {
    return func->GetAttr<::tvm::Bool>("runner_function", ::tvm::Bool(false))
        .value();
  };

  std::vector<std::pair<::tvm::GlobalVar, ::tvm::tirx::PrimFunc>> funcs;
  for (auto [gvar, base_func] : mod->functions) {
    ICHECK(base_func->IsInstance<::tvm::tirx::PrimFuncNode>())
        << "TileLangCodegenCHost: Can only take PrimFunc";
    auto prim_func = ::tvm::Downcast<::tvm::tirx::PrimFunc>(base_func);
    funcs.push_back({gvar, prim_func});
  }

  auto sort_key = [&is_aot_executor_fn](const auto &kv) {
    return std::tuple{is_aot_executor_fn(kv.second), kv.first->name_hint};
  };
  std::sort(funcs.begin(), funcs.end(),
            [&sort_key](const auto &kv_a, const auto &kv_b) {
              return sort_key(kv_a) < sort_key(kv_b);
            });

  for (const auto &[gvar, prim_func] : funcs) {
    cg.DeclareFunction(gvar, prim_func);
  }

  for (const auto &[gvar, prim_func] : funcs) {
    cg.AddFunction(gvar, prim_func, emit_fwd_func_decl);
  }

  std::string code = cg.Finish();
  if (const auto f = Function::GetGlobal("tilelang_callback_c_host_postproc")) {
    code = (*f)(code, target).cast<std::string>();
  }
  return ::tvm::codegen::CSourceModuleCreate(code, "c", cg.GetFunctionNames());
}

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = reflection;
  refl::GlobalDef().def("target.build.tilelang_c_host", BuildTileLangCHost);
}

} // namespace tl
} // namespace tvm
