/*
 * 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/cast.h>
#include <tvm/ir/expr.h>

#include "../../../tirx/script/printer/utils.h"
#include "./utils.h"

namespace tvm {
namespace script {
namespace printer {

/*! \brief Find the innermost Relax function frame, or the innermost Relax frame. */
RelaxFrameNode* GetRelaxFrame(IRDocsifier d) {
  RelaxFrameNode* f = nullptr;
  for (auto it = d->frames.rbegin(); it != d->frames.rend(); ++it) {
    if (const auto* relax_frame = (*it).as<RelaxFrameNode>()) {
      if (relax_frame->is_func) {
        return const_cast<RelaxFrameNode*>(relax_frame);
      } else if (f == nullptr) {
        f = const_cast<RelaxFrameNode*>(relax_frame);
      }
    }
  }
  return f;
}

Doc PrintCanonicalVar(Var n, AccessPath n_p, IRDocsifier d) {
  if (!n->ty.as<PrimTypeNode>()) {
    return PrintRelaxVar(n, n_p, d);
  }
  PrimVar prim_var = n.as_or_throw<PrimVar>();
  if (!d->IsVarDefined(n)) {
    PrimType n_ty = n->ty.as_or_throw<PrimType>();
    TVM_FFI_CHECK(!n_ty.IsScalableVector() && !n_ty.IsFixedLengthVector(), TypeError)
        << "Relax only uses scalar TIR variables,"
        << "but received TIR variable " << n << " with dtype " << n_ty->dtype;
    RelaxFrameNode* f = GetRelaxFrame(d);
    // There should be at least one Relax frame
    if (f == nullptr) {
      TVM_FFI_THROW(IndexError) << "No relax environment is found when printing a TIR var under "
                                   "relax's dispatch token";
    }
    // If the Relax function frame is collecting func vars
    if (f->func_vars) {
      TVM_FFI_ICHECK(f->is_func);
      f->func_vars->insert(n.get());
      if (!f->prim_params->count(n.get())) {
        f->type_vars->insert(n.get());
      }
    }
    Frame frame = d->frames.front();
    IdDoc var = d->Define(n, frame, n->name.empty() ? "v" : n->name);
    var->source_paths.push_back(n_p);
    if (!f->func_vars || f->prim_params->count(n.get()) || !f->type_vars->count(n.get())) {
      frame->stmts.push_back(AssignDoc(var, PrintVarCreation(prim_var, n_p, d), std::nullopt));
    }
  }
  if (ffi::Optional<ExprDoc> doc = d->GetVarDoc(n)) {
    return doc.value();
  }
  TVM_FFI_THROW(IndexError) << "Variable is not defined in the environment: " << n;
  TVM_FFI_UNREACHABLE();
}

TVM_FFI_STATIC_INIT_BLOCK() { IRDocsifier::vtable().set_dispatch<Var>("relax", PrintCanonicalVar); }

TVM_FFI_STATIC_INIT_BLOCK() {
  IRDocsifier::vtable().set_dispatch<tvm::IntImm>(                        //
      "relax", [](tvm::IntImm n, AccessPath n_p, IRDocsifier d) -> Doc {  //
        // TODO(@junrushao): support non-int64 cases
        if (n->ty.as_or_throw<PrimType>().MatchesElementType(DLDataTypeCode::kDLBool, 8)) {
          return LiteralDoc::Boolean(static_cast<bool>(n->value), n_p);
        } else {
          return LiteralDoc::Int(n, n_p);
        }
      });
}

TVM_FFI_STATIC_INIT_BLOCK() {
  IRDocsifier::vtable().set_dispatch<tvm::GlobalVar>(                        //
      "relax", [](tvm::GlobalVar n, AccessPath n_p, IRDocsifier d) -> Doc {  //
        if (ffi::Optional<ExprDoc> doc = d->GetVarDoc(n)) {
          return doc.value();
        } else {
          IdDoc ret(n->name_hint);
          ret->source_paths.push_back(n_p);
          return ret;
        }
      });
}

TVM_FFI_STATIC_INIT_BLOCK() {
  IRDocsifier::vtable().set_dispatch<tvm::IRModule>(                          //
      "relax", [](tvm::IRModule mod, AccessPath n_p, IRDocsifier d) -> Doc {  //
        ffi::Optional<ExprDoc> doc = d->GetVarDoc(mod);
        TVM_FFI_ICHECK(doc) << "Unable to print IRModule before definition in Relax.";
        if (d->cfg->module_alias.empty()) {
          // Use Module Name directly
          return doc.value();
        }
        RelaxFrameNode* f = GetRelaxFrame(d);
        TVM_FFI_CHECK(f != nullptr && f->is_func, IndexError)
            << "No relax environment is found when printing a module alias var "
               "under relax's dispatch token";
        if (!f->module_alias_printed) {
          // If the module_alias is not defined before, define it.
          f->stmts.push_back(AssignDoc(IdDoc(d->cfg->module_alias), doc.value(), std::nullopt));
          f->module_alias_printed = true;
        }
        return IdDoc(d->cfg->module_alias);
      });
}

TVM_FFI_STATIC_INIT_BLOCK() {
  IRDocsifier::vtable().set_dispatch<Range>(
      "relax", [](Range range, AccessPath p, IRDocsifier d) -> Doc {
        return Relax(d, "Range")
            ->Call({
                d->AsDoc<ExprDoc>(range->min, p->Attr("min")),
                d->AsDoc<ExprDoc>(range->extent + range->min, p->Attr("extent")),
            });
      });
}

}  // namespace printer
}  // namespace script
}  // namespace tvm