/*
 * 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 make_packed_api.cc Lower PrimFunc to use the packed function API.
 */
#include <tvm/ffi/cast.h>
#include <tvm/ffi/extra/module.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/access_path.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/device_api.h>
#include <tvm/target/target.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/buffer.h>
#include <tvm/tirx/builtin.h>
#include <tvm/tirx/expr.h>
#include <tvm/tirx/stmt_functor.h>
#include <tvm/tirx/transform.h>

#include <unordered_set>
#include <utility>
#include <vector>

#include "ir_utils.h"
#include "tvm_ffi_binder.h"

namespace tvm {
namespace tirx {

namespace {
class ReturnRewriter : public StmtMutator {
 public:
  explicit ReturnRewriter(Var ret_var) : ret_var_(ret_var) {}

  Stmt VisitStmt_(const ForNode* node) override {
    if (node->kind == ForKind::kParallel) in_parallel_ += 1;
    Stmt ret = StmtMutator::VisitStmt_(node);
    if (node->kind == ForKind::kParallel) in_parallel_ -= 1;
    return ret;
  }

  Stmt VisitStmt_(const EvaluateNode* node) override {
    Stmt ret = StmtMutator::VisitStmt_(node);
    const EvaluateNode* eval = ret.as<EvaluateNode>();
    TVM_FFI_ICHECK(eval);
    if (const CallNode* call = eval->value.as<CallNode>()) {
      if (call->op.same_as(builtin::ret())) {
        TVM_FFI_ICHECK_EQ(in_parallel_, 0) << "tirx.ret cannot be used in parallel scope.";
        TVM_FFI_ICHECK_EQ(call->args.size(), 1) << "tirx.ret expect a single argument.";
        ret = WriteToOut(call->args[0]);
      }
    }
    return ret;
  }

 private:
  struct ConvertedInfo {
    int type_index{-1};
    Expr expr;
  };

  ConvertedInfo ConvertForFFI(Expr val) {
    ConvertedInfo info;

    // convert val's data type to FFI data type, return type code
    if (val->ty.as<PointerTypeNode>()) {
      info.type_index = ffi::TypeIndex::kTVMFFIOpaquePtr;
      info.expr = val;
      return info;
    }

    PrimExpr prim_val = val.as_or_throw<PrimExpr>();
    PrimType dtype = prim_val.ty();
    if (dtype.MatchesCode(DLDataTypeCode::kDLBool)) {
      info.type_index = ffi::TypeIndex::kTVMFFIBool;
      info.expr = Cast(PrimType::Int(64), prim_val);

    } else if (dtype.MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt)) {
      info.type_index = ffi::TypeIndex::kTVMFFIInt;
      info.expr = Cast(PrimType::Int(64), prim_val);
    } else if (dtype.code() == DLDataTypeCode::kDLFloat) {
      info.type_index = ffi::TypeIndex::kTVMFFIFloat;
      info.expr = Cast(PrimType::Float(64), prim_val);
    } else if (dtype.IsVoid()) {
      info.type_index = ffi::TypeIndex::kTVMFFINone;
      info.expr = prim_val;
    } else {
      TVM_FFI_THROW(InternalError) << "data type " << dtype->dtype << " not supported yet";
    }
    return info;
  }

  Stmt WriteToOut(Expr val) {
    auto info = ConvertForFFI(val);
    Stmt store_tindex = tirx::Evaluate(
        Call(PrimType::Int(32), tirx::builtin::tvm_struct_set(),
             {ret_var_, IntImm::Int32(0), IntImm::Int32(tirx::builtin::kTVMFFIAnyTypeIndex),
              IntImm::Int32(info.type_index)})
            .as_or_throw<PrimExpr>());
    Stmt store_zero_padding =
        tirx::Evaluate(Call(PrimType::Int(32), tirx::builtin::tvm_struct_set(),
                            {ret_var_, IntImm::Int32(0),
                             IntImm::Int32(tirx::builtin::kTVMFFIAnyZeroPadding), IntImm::Int32(0)})
                           .as_or_throw<PrimExpr>());
    Stmt store_val =
        tirx::Evaluate(Call(PrimType::Int(32), tirx::builtin::tvm_struct_set(),
                            {ret_var_, IntImm::Int32(0),
                             IntImm::Int32(tirx::builtin::kTVMFFIAnyUnionValue), info.expr})
                           .as_or_throw<PrimExpr>());
    Stmt ret_zero = Evaluate(tvm::ret(0));
    return SeqStmt({store_tindex, store_zero_padding, store_val, ret_zero});
  }

  Var ret_var_;
  int in_parallel_{0};
};

class SubroutineCallRewriter : public StmtExprMutator {
 public:
  static ffi::Optional<Stmt> Apply(const ffi::Map<GlobalVar, ffi::String>& packed_func_methods,
                                   Stmt stmt) {
    SubroutineCallRewriter rewriter(packed_func_methods);
    stmt = rewriter.VisitStmt(std::move(stmt));
    if (rewriter.made_change_) {
      return stmt;
    } else {
      return std::nullopt;
    }
  }

 private:
  explicit SubroutineCallRewriter(const ffi::Map<GlobalVar, ffi::String>& packed_func_methods)
      : packed_func_methods(packed_func_methods) {}

  Expr VisitExpr_(const CallNode* op) override {
    auto node = StmtExprMutator::VisitExpr_(op).as_or_throw<Call>();

    if (auto* gvar_ptr = node->op.as<GlobalVarNode>()) {
      auto gvar = ffi::GetRef<GlobalVar>(gvar_ptr);
      if (auto symbol = packed_func_methods.Get(gvar)) {
        ffi::Array<Expr> cpacked_args;
        cpacked_args.push_back(tirx::StringImm(symbol.value()));
        for (const Expr& arg : node->args) {
          cpacked_args.push_back(arg);
        }

        // push an empty handle to be compatible with current cpacked convention
        cpacked_args.push_back(tirx::ConstHandle(0));
        made_change_ = true;
        return Call(node->ty, tirx::builtin::tvm_call_cpacked(), cpacked_args);
      }
    }

    return node;
  }
  const ffi::Map<GlobalVar, ffi::String>& packed_func_methods;
  bool made_change_{false};
};

}  // namespace

/* \brief Return the global_symbol of the function, if it should be updated
 *
 * \param func The function to be inspected
 *
 * \returns The global_symbol to be used for the function at call
 * sites, or std::nullopt if the function is to remain unchanged.
 */
ffi::Optional<ffi::String> RequiresPackedAPI(const PrimFunc& func) {
  // A function with an explicit calling convention has already been
  // lowered, and should not be modified.
  if (auto opt = func->GetAttr<CallingConv>(tvm::attr::kCallingConv)) {
    if (opt.value() != CallingConv::kDefault) {
      return std::nullopt;
    }
  }

  // Internal function calls do not need the ffi::Function API
  auto global_symbol = func->GetAttr<ffi::String>(tvm::attr::kGlobalSymbol);
  if (!global_symbol.has_value()) {
    return std::nullopt;
  }

  return global_symbol.value();
}

PrimFunc MakePackedAPI(PrimFunc func) {
  auto global_symbol = RequiresPackedAPI(func);
  if (!global_symbol.has_value()) {
    return func;
  }
  std::string name_hint = global_symbol.value();

  Target target = [&]() {
    auto opt = func->GetAttr<Target>(tvm::attr::kTarget);
    TVM_FFI_ICHECK(opt)
        << "MakePackedAPI required the function to be annotated with tvm::attr::kTarget ("
        << tvm::attr::kTarget << "), but the function only has attributes " << func->attrs;
    return opt.value();
  }();
  int target_device_type = target->GetTargetDeviceType();

  // A function without a host target has already been lowered.
  Target target_host;
  if (auto opt = target->GetHost()) {
    target_host = opt.value();
  } else {
    return func;
  }

  auto* func_ptr = func.CopyOnWrite();
  const Stmt nop = Evaluate(0);

  // Data field definitions
  Var v_self_handle("self_handle", PointerType::VoidPointerTy());
  Var v_packed_args("args", PointerType::VoidPointerTy());
  Var v_num_packed_args("num_args", PrimType::Int(32));
  Var v_result("result", PointerType::VoidPointerTy());

  // The device context
  Var device_id("dev_id");
  IntImm device_type(PrimType::Int(32), target_device_type);

  // Create TVMFFIABIBuilder and decode all packed args
  TVMFFIABIBuilder binder(name_hint, func_ptr->params, func_ptr->buffer_map, v_packed_args,
                          v_num_packed_args, device_type, device_id.as_or_throw<PrimExpr>());
  binder.DecodeAllParams();

  auto result = binder.Finalize();
  bool need_set_device = result.var_defs.count(device_id.get());

  std::vector<Stmt> seq_check;

  // signature: (void* handle, TVMFFIAny* packed_args, int num_args, TVMFFIAny* v_result)
  ffi::Array<Var> args{v_self_handle, v_packed_args, v_num_packed_args, v_result};

  // reset global symbol to attach prefix
  func = WithAttrs(std::move(func), {{tvm::attr::kCallingConv, CallingConv::kCPackedFunc},
                                     {tvm::attr::kTarget, target_host},
                                     {tvm::attr::kGlobalSymbol,
                                      ffi::symbol::tvm_ffi_symbol_prefix + global_symbol.value()}});

  Stmt body = ReturnRewriter(v_result)(func_ptr->body);
  body = AttrStmt(IntImm::Int32(0), attr::compute_scope, StringImm(name_hint + "_compute_"), body);
  // Set device context
  if (need_set_device) {
    ffi::Any node = ffi::String("default");
    seq_check.push_back(AttrStmt(node, attr::device_id, device_id.as_or_throw<PrimExpr>(), nop));
    seq_check.push_back(AttrStmt(node, attr::device_type, device_type, nop));

    if (runtime::DeviceAPI::NeedSetDevice(target_device_type)) {
      Stmt set_device = Evaluate(Call(PrimType::Int(32), builtin::tvm_call_packed(),
                                      {StringImm(runtime::symbol::tvm_set_device), device_type,
                                       device_id.as_or_throw<PrimExpr>()})
                                     .as_or_throw<PrimExpr>());
      body = SeqStmt({set_device, body});
    }
  }

  // Return error code of zero on success
  body = SeqStmt({body, Evaluate(ret(IntImm::Int32(0)))});

  body = MergeNest({std::move(result.init_nest), seq_check, std::move(result.asserts),
                    std::move(result.decl_buffers)},
                   body);
  func_ptr->body = body;
  func_ptr->params = args;

  ffi::Array<Var> undefined = UndefinedVars(func_ptr->body, func_ptr->params);
  TVM_FFI_ICHECK_EQ(undefined.size(), 0)
      << "In PrimFunc " << name_hint << " variables " << undefined
      << " are used, but are not passed in as API arguments";

  func_ptr->buffer_map = ffi::Map<Var, Buffer>();
  func_ptr->ret_type = PrimType::Int(32);

  // return the function.
  return func;
}

namespace transform {

Pass MakePackedAPI() {
  auto pass_func = [](IRModule mod, PassContext ctx) {
    ffi::Map<GlobalVar, ffi::String> packed_func_methods;
    for (const auto& [gvar, base_func] : mod->functions) {
      if (auto opt = base_func.as<PrimFunc>()) {
        auto prim_func = opt.value();
        if (auto global_symbol = RequiresPackedAPI(prim_func)) {
          packed_func_methods.Set(gvar, global_symbol.value());
        }
      }
    }

    IRModuleNode* mptr = mod.CopyOnWrite();
    IRModule updates;

    for (const auto& [gvar, base_func] : mptr->functions) {
      if (auto opt = base_func.as<PrimFunc>()) {
        auto func = opt.value();
        auto orig_func = func;

        if (auto body = SubroutineCallRewriter::Apply(packed_func_methods, func->body)) {
          func.CopyOnWrite()->body = body.value();
        }

        func = MakePackedAPI(std::move(func));

        if (!func.same_as(orig_func)) {
          updates->Add(gvar, func);
        }
      }
    }

    if (updates->functions.size()) {
      mod.CopyOnWrite()->Update(updates);
    }
    return mod;
  };

  return tvm::transform::CreateModulePass(pass_func, 0, "tirx.MakePackedAPI", {});
}

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef().def("tirx.transform.MakePackedAPI", []() { return MakePackedAPI(); });
}
}  // namespace transform
}  // namespace tirx
}  // namespace tvm
