/*
 * 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.
 */
#include <tvm/ffi/reflection/registry.h>
#include <tvm/relax/analysis.h>
#include <tvm/relax/script/builder/ir.h>
#include <tvm/relax/type.h>
#include <tvm/tirx/op.h>

#include "./utils.h"

namespace tvm {
namespace script {
namespace ir_builder {
namespace relax {

///////////////////////////////// Vars //////////////////////////////////

using tvm::script::ir_builder::details::Namer;

TVM_STATIC_IR_FUNCTOR(Namer, vtable)
    .set_dispatch<tvm::relax::VarNode>([](const ffi::ObjectRef& node, ffi::String name) -> void {
      using tvm::relax::VarNode;
      VarNode* var = const_cast<VarNode*>(node.as<VarNode>());
      var->name_hint = name;
    });

TVM_STATIC_IR_FUNCTOR(Namer, vtable)
    .set_dispatch<tvm::relax::DataflowVarNode>([](const ffi::ObjectRef& node,
                                                  ffi::String name) -> void {
      using tvm::relax::DataflowVarNode;
      DataflowVarNode* var = const_cast<DataflowVarNode*>(node.as<DataflowVarNode>());
      var->name_hint = name;
    });

/////////////////////////////// Function ////////////////////////////////

FunctionFrame Function(bool is_pure, bool is_private) {
  ffi::ObjectPtr<FunctionFrameNode> n = ffi::make_object<FunctionFrameNode>();
  const IRBuilder& ir_builder = IRBuilder::Current();
  ffi::Optional<tvm::IRModule> mod = std::nullopt;
  if (const ffi::Optional<ir::IRModuleFrame> mod_frame =
          ir_builder->GetLastFrame<ir::IRModuleFrame>()) {
    mod = tvm::IRModule(mod_frame.value()->functions);
  }
  n->block_builder = tvm::relax::BlockBuilder::Create(
      /*mod=*/mod, tvm::relax::BlockBuilder::DisableOperatorSpecificNormalizationForTVMScript());
  n->is_pure = is_pure;
  n->is_private = is_private;
  return FunctionFrame(n);
}

tvm::relax::Var Arg(const ffi::String& name, const tvm::Type& ty) {
  FunctionFrame frame = FindFunctionFrame("R.Arg");
  tvm::relax::Var var(name, ty);
  frame->params.push_back(var);
  frame->block_builder->AddDefinitionToScope(var);

  return var;
}

void FuncName(const ffi::String& name) {
  FunctionFrame frame = FindFunctionFrame("R.func_name");
  if (frame->name.has_value()) {
    TVM_FFI_THROW(ValueError) << "Duplicate function name, previous one is: \""
                              << frame->name.value() << "\"";
  }
  frame->name = name;
}

void FuncAttrs(ffi::Map<ffi::String, ffi::Any> attrs) {
  FunctionFrame frame = FindFunctionFrame("R.func_attr");
  for (const auto& [key, value] : attrs) {
    if (key == tvm::attr::kGlobalSymbol && frame->is_private.value_or(false)) {
      TVM_FFI_THROW(ValueError) << "A private function may not have the kGlobalSymbol (\""
                                << tvm::attr::kGlobalSymbol << "\") attribute.  "
                                << "However, a private function specified the global symbol as "
                                << value;
    }
    if (auto prev = frame->attrs.Get(key)) {
      TVM_FFI_THROW(ValueError) << "Duplicate R.func_attr annotation for key = \"" << key << "\".  "
                                << "Previous value was " << prev.value()
                                << ", with later definition as " << value;
    } else {
      frame->attrs.Set(key, value);
    }
  }
}

void FuncRetType(const tvm::Type& ret_ty) {
  FunctionFrame frame = FindFunctionFrame("R.func_ret_type");
  if (frame->ret_ty.has_value()) {
    TVM_FFI_THROW(ValueError) << "Duplicate function return type, previous one is:\n "
                              << frame->ret_ty.value();
  }
  frame->ret_ty = ret_ty;
}

void FuncRetValue(const tvm::relax::Expr& value) {
  // Step 0. Normalize the value.
  const tvm::relax::BlockBuilder& block_builder = GetBlockBuilder();
  tvm::relax::Expr normalized_value = block_builder->Normalize(value);

  IRBuilder ir_builder = IRBuilder::Current();

  // Step 1. The current Relax TVMScript syntax only allows function return appearing at the end of
  // a function body. Therefore if there is any unended block frame when dealing with function
  // return, we should end the block frame.

  if (auto opt = ir_builder->GetLastFrame<BindingBlockFrame>()) {
    auto block_frame = opt.value();
    for (const auto& var : tvm::relax::FreeVars(normalized_value)) {
      if (var->IsInstance<tvm::relax::DataflowVarNode>()) {
        block_frame->output_vars.push_back(var);
      }
    }
  }
  // Step 2. Add the output value to the function frame.
  FunctionFrame frame = FindFunctionFrame("return");
  TVM_FFI_CHECK(!frame->output.has_value(), ValueError)
      << "Relax functions do not support multiple return statement.  "
      << "However, return of " << normalized_value << " occurred after a return of "
      << frame->output << ".  "
      << "Please make sure function only has a single return statement, "
      << "which appears at the end of function.";

  frame->output = std::move(normalized_value);
}

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef()
      .def("script.ir_builder.relax.Function", Function)
      .def("script.ir_builder.relax.Arg", Arg)
      .def("script.ir_builder.relax.FuncName", FuncName)
      .def("script.ir_builder.relax.FuncAttrs", FuncAttrs)
      .def("script.ir_builder.relax.FuncRetType", FuncRetType)
      .def("script.ir_builder.relax.FuncRetValue", FuncRetValue);
}

///////////////////////////// BindingBlock //////////////////////////////

BindingBlockFrame Dataflow() {
  ffi::ObjectPtr<BindingBlockFrameNode> n = ffi::make_object<BindingBlockFrameNode>();
  n->is_dataflow = true;
  n->block_ended = false;
  return BindingBlockFrame(n);
}

BindingBlockFrame BindingBlock() {
  ffi::ObjectPtr<BindingBlockFrameNode> n = ffi::make_object<BindingBlockFrameNode>();
  n->is_dataflow = false;
  n->block_ended = false;
  return BindingBlockFrame(n);
}

void DataflowBlockOutput(const ffi::Array<tvm::relax::Var>& vars) {
  // Step 1. Check that we're in a Dataflow block that is not ended.
  ffi::Optional<BindingBlockFrame> block_frame =
      IRBuilder::Current()->GetLastFrame<BindingBlockFrame>();
  TVM_FFI_CHECK(block_frame.has_value() && block_frame.value()->is_dataflow, ValueError)
      << "`R.output` should appear inside a dataflow block. However, the current "
         "innermost block is not a dataflow block.";
  TVM_FFI_CHECK(!block_frame.value()->block_ended, ValueError)
      << "It is not allowed for a dataflow block to have multiple output operation.";

  // Step 2. Mark the block frame ended of construction, so that any followup binding after this
  // mark in the dataflow block will lead to an error.
  block_frame.value()->block_ended = true;

  // Step 3. All the output variables must be global variables and must be emitted by this dataflow
  // block.
  const ffi::Array<tvm::relax::Var>& emitted_vars = block_frame.value()->emitted_vars;
  for (const tvm::relax::Var& var : vars) {
    TVM_FFI_CHECK(std::find_if(emitted_vars.begin(), emitted_vars.end(),
                               [&](const tvm::relax::Var& emitted) {
                                 return emitted.same_as(var);
                               }) != emitted_vars.end(),
                  ValueError)
        << "An output variable is not emitted by this dataflow block. Please make sure "
           "all dataflow block output variables are emitted exactly by this block.";
    block_frame.value()->output_vars.push_back(var);
  }
}

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef()
      .def("script.ir_builder.relax.Dataflow", Dataflow)
      .def("script.ir_builder.relax.BindingBlock", BindingBlock)
      .def("script.ir_builder.relax.DataflowBlockOutput", DataflowBlockOutput);
}

/////////////////////////////// Bindings ///////////////////////////////

tvm::relax::Var Emit(const tvm::relax::Expr& expr, const ffi::Optional<tvm::Type>& annotate_ty) {
  using tvm::relax::GetType;
  BindingBlockFrame block_frame = CheckBindingBlockFrameExistAndUnended();
  const tvm::relax::BlockBuilder& block_builder = GetBlockBuilder();
  if (annotate_ty.has_value()) {
    const auto& ty = annotate_ty.value();
    if (expr->ty.IsMissing()) {
      tvm::relax::UpdateType(expr, ty);
    } else {
      TVM_FFI_ICHECK(tvm::relax::TypeBaseCheck(ty, GetType(expr)) !=
                     tvm::relax::BaseCheckResult::kFailL0)
          << "Invalid annotation. Got rhs value type: " << GetType(expr) << ", given type: " << ty;
    }
  }
  tvm::relax::Var var = block_builder->Emit(expr);
  block_frame->emitted_vars.push_back(var);
  return var;
}

tvm::relax::Var EmitMatchCast(const tvm::relax::Expr& value, const tvm::Type& ty) {
  BindingBlockFrame block_frame = CheckBindingBlockFrameExistAndUnended();
  const tvm::relax::BlockBuilder& block_builder = GetBlockBuilder();

  tvm::relax::Var var = block_builder->EmitMatchCast(value, ty);
  block_frame->emitted_vars.push_back(var);
  return var;
}

tvm::relax::Var EmitVarBinding(const tvm::relax::VarBinding& binding) {
  BindingBlockFrame block_frame = CheckBindingBlockFrameExistAndUnended();
  const tvm::relax::BlockBuilder& block_builder = GetBlockBuilder();
  block_builder->EmitNormalized(binding);
  block_frame->emitted_vars.push_back(binding->var);
  return binding->var;
}

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef()
      .def("script.ir_builder.relax.Emit", Emit)
      .def("script.ir_builder.relax.EmitMatchCast", EmitMatchCast)
      .def("script.ir_builder.relax.EmitVarBinding", EmitVarBinding);
}

/////////////////////////////// SeqExpr ///////////////////////////////

SeqExprFrame SeqExpr() {
  ffi::ObjectPtr<SeqExprFrameNode> n = ffi::make_object<SeqExprFrameNode>();
  return SeqExprFrame(n);
}

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef().def("script.ir_builder.relax.SeqExpr", SeqExpr);
}

///////////////////////////// If Then Else /////////////////////////////

IfFrame If(tvm::relax::Expr condition) {
  ffi::ObjectPtr<IfFrameNode> n = ffi::make_object<IfFrameNode>();
  n->condition = condition;
  n->then_expr = std::nullopt;
  n->else_expr = std::nullopt;
  return IfFrame(n);
}

ThenFrame Then() {
  ffi::ObjectPtr<ThenFrameNode> n = ffi::make_object<ThenFrameNode>();
  return ThenFrame(n);
}

ElseFrame Else() {
  ffi::ObjectPtr<ElseFrameNode> n = ffi::make_object<ElseFrameNode>();
  return ElseFrame(n);
}

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef()
      .def("script.ir_builder.relax.If", If)
      .def("script.ir_builder.relax.Then", Then)
      .def("script.ir_builder.relax.Else", Else);
}

}  // namespace relax
}  // namespace ir_builder
}  // namespace script
}  // namespace tvm
