/*
 * 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/runtime/device_api.h>

#include "./utils.h"

namespace tvm {
namespace script {
namespace printer {

bool IsSimpleBuffer(const tirx::Buffer& buf) {
  if (!buf->strides.empty()) {
    return false;
  }
  for (const PrimExpr& shp_i : buf->shape) {
    if (!tirx::UndefinedVars(shp_i).empty()) {
      return false;
    }
  }
  for (const PrimExpr& stride_i : buf->strides) {
    if (!tirx::UndefinedVars(stride_i).empty()) {
      return false;
    }
  }
  if (!tirx::UndefinedVars(buf->elem_offset).empty()) {
    return false;
  } else if (buf->elem_offset->IsInstance<IntImmNode>()) {
    IntImm elem_offset = Downcast<IntImm>(buf->elem_offset);
    if (elem_offset->value != 0) {
      return false;
    }
  }
  return buf.scope() == "global" && buf->data_alignment == runtime::kAllocAlignment &&
         buf->offset_factor == 1 && buf->buffer_type == tirx::BufferType::kDefault &&
         !buf->axis_separators.size();
}

int CountVarOccurrence(const tirx::PrimFunc& f, const tirx::Var& v) {
  OccurrenceCounter counter(v.get());
  counter(f->body);
  for (const tirx::Var& v : f->params) {
    counter(v);
  }
  for (const auto& pair : f->buffer_map) {
    counter(pair.first);
    counter.VisitBuffer(pair.second.get());
  }
  return counter.count;
}

TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
    .set_dispatch<tirx::PrimFunc>("", [](tirx::PrimFunc func, AccessPath p, IRDocsifier d) -> Doc {
      With<TIRFrame> f(d, func);
      (*f)->AddDispatchToken(d, "tirx");
      IdDoc func_name = IdDoc(FindFunctionName(d, func).value_or("main"));
      d->SetCommonPrefix(func, [](const ObjectRef& obj) {
        return obj->IsInstance<tirx::VarNode>() || obj->IsInstance<tirx::BufferNode>();
      });
      int n_args = func->params.size();
      std::unordered_map<const tirx::VarNode*, int> buffer_data_counter;
      for (const auto& pair : func->buffer_map) {
        const tirx::VarNode* data_var = pair.second->data.get();
        if (!buffer_data_counter.count(data_var)) {
          buffer_data_counter.insert({data_var, 0});
        }
        ++buffer_data_counter.at(data_var);
      }
      // Step 1. Handle `func->params`
      ffi::Array<AssignDoc> args;
      args.reserve(n_args);
      std::unordered_set<const tirx::BufferNode*> buffer_inlined;
      for (int i = 0; i < n_args; ++i) {
        tirx::Var var = func->params[i];
        AccessPath var_p = p->Attr("params")->ArrayItem(i);
        if (d->cfg->syntax_sugar && CountVarOccurrence(func, var) == 2 &&
            func->buffer_map.count(var)) {
          tirx::Buffer buffer = func->buffer_map[var];
          if (IsSimpleBuffer(buffer) && buffer_data_counter.at(buffer->data.get()) == 1) {
            AccessPath buffer_p = p->Attr("buffer_map")->MapItem(var);
            IdDoc lhs = DefineBuffer(buffer, *f, d);
            ExprDoc annotation = BufferAttn(buffer, buffer_p, *f, d);
            args.push_back(AssignDoc(lhs, std::nullopt, annotation));
            buffer_inlined.insert(buffer.get());
            continue;
          }
        }
        ExprDoc a = d->AsDoc<ExprDoc>(var->type_annotation, var_p->Attr("type_annotation"));
        args.push_back(AssignDoc(DefineVar(var, *f, d), std::nullopt, a));
      }
      // Step 2. Handle `func->attrs`
      if (func->attrs.defined() && !func->attrs->dict.empty()) {
        // for global symbol, don't display it if it matches the func name
        if (func->attrs->dict.count(tvm::attr::kGlobalSymbol) &&
            Downcast<ffi::String>(func->attrs->dict.at(tvm::attr::kGlobalSymbol)) ==
                func_name->name) {
          ffi::Map<ffi::String, Any> new_attrs;
          for (auto kv : func->attrs->dict) {
            if (kv.first != tvm::attr::kGlobalSymbol) {
              new_attrs.Set(kv.first, kv.second);
            }
          }
          if (!new_attrs.empty()) {
            (*f)->stmts.push_back(ExprStmtDoc(
                TIR(d, "func_attr")  //
                    ->Call({d->AsDoc<ExprDoc>(DictAttrs(new_attrs), p->Attr("attrs"))})));
          }
        } else {
          (*f)->stmts.push_back(
              ExprStmtDoc(TIR(d, "func_attr")  //
                              ->Call({d->AsDoc<ExprDoc>(func->attrs, p->Attr("attrs"))})));
        }
      }
      // Step 3. Handle `func->buffer_map`
      for (int i = 0; i < n_args; ++i) {
        tirx::Var param = func->params[i];
        if (func->buffer_map.count(param)) {
          tirx::Buffer buffer = func->buffer_map[param];
          if (buffer_inlined.count(buffer.get())) {
            continue;
          }
          ExprDoc param_doc = args[i]->lhs;
          AccessPath buffer_p = p->Attr("buffer_map")->MapItem(param);
          ExprDoc lhs = DefineBuffer(buffer, *f, d);
          ExprDoc rhs = BufferDecl(buffer, "match_buffer", {param_doc}, buffer_p, *f, d,
                                   BufferVarDefinition::MatchBuffer);
          (*f)->stmts.push_back(AssignDoc(lhs, rhs, std::nullopt));
        }
      }
      // Step 4. Handle `func->body`
      ffi::Optional<tirx::SBlock> implicit_root_block = [&]() -> ffi::Optional<tirx::SBlock> {
        const tirx::SBlockRealizeNode* root_block_realize =
            func->body.as<tirx::SBlockRealizeNode>();
        if (root_block_realize && !root_block_realize->iter_values.size() &&
            tirx::is_one(root_block_realize->predicate)) {
          tirx::SBlock root_block = root_block_realize->block;
          if (!root_block->annotations.size() && !root_block->match_buffers.size() &&
              !root_block->reads.size() && !root_block->writes.size() &&
              !root_block->init.defined()) {
            const tirx::SBlockRealizeNode* block_realize =
                root_block->body.as<tirx::SBlockRealizeNode>();
            if (root_block->alloc_buffers.size() ||
                (block_realize && block_realize->block->iter_vars.size()) ||
                (!block_realize && tirx::ContainsNode<tirx::SBlockRealizeNode>(root_block->body))) {
              return root_block;
            }
          }
        }
        return std::nullopt;
      }();
      if (d->cfg->syntax_sugar && implicit_root_block) {
        tirx::SBlock root_block = implicit_root_block.value();
        AccessPath root_block_p = p->Attr("body")->Attr("block");
        (*f)->stmts.push_back(CommentDoc("with T.sblock(\"root\"):"));
        // Handle root block `alloc_buffer`
        for (int i = 0, n = root_block->alloc_buffers.size(); i < n; ++i) {
          tirx::Buffer buffer = root_block->alloc_buffers[i];
          AccessPath buffer_p = root_block_p->Attr("alloc_buffers")->ArrayItem(i);
          IdDoc lhs = DefineBuffer(buffer, *f, d);
          ExprDoc rhs = BufferDecl(buffer, "sblock_alloc_buffer", {}, buffer_p, *f, d,
                                   BufferVarDefinition::DataPointer);
          (*f)->stmts.push_back(AssignDoc(lhs, rhs, std::nullopt));
        }
        AsDocBody(root_block->body, root_block_p->Attr("body"), f->get(), d);
      } else {
        AsDocBody(func->body, p->Attr("body"), f->get(), d);
      }
      ffi::Optional<ExprDoc> ret_type = std::nullopt;
      if (func->ret_type.defined()) {
        const auto* as_tuple = func->ret_type.as<TupleTypeNode>();
        if (!as_tuple || as_tuple->fields.size()) {
          ret_type = d->AsDoc<ExprDoc>(func->ret_type, p->Attr("ret_type"));
        }
      }
      // Step 5. Determine if we need to display the private annotation in the decorator
      ExprDoc decorator = TIR(d, "prim_func");
      // mark private if there is no global symbol
      if (!func->attrs.defined() || !func->attrs->dict.count(tvm::attr::kGlobalSymbol)) {
        ffi::Array<ExprDoc> pos_args;
        decorator = decorator->Call(pos_args, {"private"},
                                    {LiteralDoc::Boolean(true, ffi::Optional<AccessPath>())});
      }

      return HeaderWrapper(d, FunctionDoc(
                                  /*name=*/func_name,
                                  /*args=*/args,
                                  /*decorators=*/{decorator},
                                  /*return_type=*/ret_type,
                                  /*body=*/(*f)->stmts));
    });

TVM_SCRIPT_REPR(tirx::PrimFuncNode, ReprPrintTIR);

TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
    .set_dispatch<tvm::GlobalVar>(                                            //
        "tirx", [](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_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
    .set_dispatch<tvm::IRModule>(                                              //
        "tirx", [](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 TIR.";
          return doc.value();
        });

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