/*
 * 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 "./utils.h"

namespace tvm {
namespace script {
namespace printer {

Doc PrintBlock(IRDocsifier d, tirx::SBlock block, AccessPath block_p,  //
               ffi::Optional<tirx::SBlockRealize> opt_realize,
               ffi::Optional<AccessPath> opt_realize_p) {
  With<TIRFrame> frame(d, block);
  TVM_FFI_ICHECK_EQ(opt_realize.has_value(), opt_realize_p.has_value());
  const tirx::SBlockRealizeNode* realize =
      opt_realize.has_value() ? opt_realize.value().get() : nullptr;
  AccessPath realize_p = *opt_realize_p;

  // Step 1. Handle block var and block bindings
  // Step 1.1. Obtain all loop var defined along path
  std::unordered_map<const tirx::VarNode*, tirx::For> loop_vars;
  for (Frame f : d->frames) {
    if (const auto* tir_f = f.as<TIRFrameNode>()) {
      if (auto for_loop = tir_f->tirx.as<tirx::For>()) {
        for (ffi::Optional<tirx::For> loop = for_loop; loop;
             loop = loop.value()->body.as<tirx::For>()) {
          loop_vars.insert(std::make_pair(loop.value()->loop_var.get(), loop.value()));
        }
      }
    }
  }

  std::vector<int> remap_vars_indices;
  auto add_remapped_iter_var = [&](int i) -> bool {
    if (realize && d->cfg->syntax_sugar) {
      tirx::ExprDeepEqual expr_equal;
      tirx::IterVar iter_var = block->iter_vars[i];
      PrimExpr value = realize->iter_values[i];
      if (iter_var->iter_type == tirx::IterVarType::kDataPar ||
          iter_var->iter_type == tirx::IterVarType::kCommReduce) {
        if (auto var = value.as<tirx::PrimVar>()) {
          if (loop_vars.count(var.value().get())) {
            tirx::For for_loop = loop_vars.at(var.value().get());
            if (expr_equal(for_loop->min, iter_var->dom->min) &&
                expr_equal(for_loop->extent, iter_var->dom->extent)) {
              remap_vars_indices.push_back(i);
              return true;
            }
          }
        }
      }
    }
    return false;
  };

  auto print_single_iter_var = [&](int i) {
    tirx::IterVar iter_var = block->iter_vars[i];
    AccessPath iter_var_p = block_p->Attr("iter_var")->ArrayItem(i);
    ExprDoc rhs = TIR(d, "axis");
    if (iter_var->iter_type == tirx::IterVarType::kDataPar) {
      rhs = rhs->Attr("spatial");
    } else if (iter_var->iter_type == tirx::IterVarType::kCommReduce) {
      rhs = rhs->Attr("reduce");
    } else if (iter_var->iter_type == tirx::IterVarType::kOrdered) {
      rhs = rhs->Attr("scan");
    } else if (iter_var->iter_type == tirx::IterVarType::kOpaque) {
      rhs = rhs->Attr("opaque");
    } else {
      TVM_FFI_THROW(ValueError) << "Unknown IterVarType in block signature: "
                                << tirx::IterVarType2String(iter_var->iter_type);
    }
    ExprDoc dom{ffi::UnsafeInit()};
    if (tirx::is_zero(iter_var->dom->min)) {
      ExprDoc extent = d->AsDoc<ExprDoc>(iter_var->dom->extent,  //
                                         iter_var_p->Attr("dom")->Attr("extent"));
      dom = extent;
    } else {
      ExprDoc min = d->AsDoc<ExprDoc>(iter_var->dom->min, iter_var_p->Attr("dom")->Attr("min"));
      ExprDoc max = d->AsDoc<ExprDoc>(iter_var->dom->min + iter_var->dom->extent,
                                      iter_var_p->Attr("dom")->Attr("extent"));
      dom = TupleDoc({min, max});
    }
    if (realize) {
      ExprDoc binding = d->AsDoc<ExprDoc>(realize->iter_values[i],  //
                                          realize_p->Attr("iter_values")->ArrayItem(i));
      rhs = rhs->Call({dom, binding});
    } else {
      rhs = rhs->Call({dom});
    }
    (*frame)->stmts.push_back(AssignDoc(DefineVar(iter_var->var, *frame, d), rhs, std::nullopt));
  };

  auto print_remapped_iter_var = [&]() {
    if (remap_vars_indices.size()) {
      int m = remap_vars_indices.size();
      if (m == 1) {
        print_single_iter_var(remap_vars_indices[0]);
        remap_vars_indices.clear();
        return;
      }
      ffi::Array<ExprDoc> lhs;
      ffi::Array<ExprDoc> loop_var_doc;
      lhs.reserve(m);
      loop_var_doc.reserve(m);
      std::string binding_type = "";
      ffi::Array<AccessPath> binding_paths;
      for (int i : remap_vars_indices) {
        tirx::IterVar iter_var = block->iter_vars[i];
        AccessPath iter_var_p = block_p->Attr("iter_vars")->ArrayItem(i);
        lhs.push_back(DefineVar(iter_var->var, *frame, d));
        loop_var_doc.push_back(d->AsDoc<ExprDoc>(realize->iter_values[i],
                                                 realize_p->Attr("iter_values")->ArrayItem(i)));
        binding_paths.push_back(iter_var_p->Attr("iter_type"));
        binding_type += iter_var->iter_type == tirx::IterVarType::kDataPar ? "S" : "R";
      }
      ExprDoc rhs = TIR(d, "axis")->Attr("remap");
      ExprDoc binding_str = LiteralDoc::Str(binding_type, std::nullopt);
      binding_str->source_paths = std::move(binding_paths);
      rhs = rhs->Call({binding_str, ListDoc(loop_var_doc)});
      (*frame)->stmts.push_back(AssignDoc(TupleDoc(lhs), rhs, std::nullopt));
      remap_vars_indices.clear();
    }
  };

  // Step 1.2. Construct all block var bindings
  int n_vars = block->iter_vars.size();
  for (int i = 0; i < n_vars; ++i) {
    if (!add_remapped_iter_var(i)) {
      print_remapped_iter_var();
      print_single_iter_var(i);
    }
  }
  print_remapped_iter_var();

  // Step 2. Handle block predicate
  if (realize) {
    PrimType predicate_ty = realize->predicate.ty();
    TVM_FFI_ICHECK(realize->predicate.defined() &&
                   predicate_ty.MatchesCode(DLDataTypeCode::kDLBool));
    if (!tirx::is_one(realize->predicate)) {
      (*frame)->stmts.push_back(ExprStmtDoc(
          TIR(d, "where")
              ->Call({d->AsDoc<ExprDoc>(realize->predicate, realize_p->Attr("predicate"))})));
    }
  }
  // Step 3. Handle block read/write regions
  {
    ffi::Array<ExprDoc> reads;
    for (int i = 0, n = block->reads.size(); i < n; ++i) {
      reads.push_back(d->AsDoc<ExprDoc>(block->reads[i], block_p->Attr("reads")->ArrayItem(i)));
    }
    (*frame)->stmts.push_back(ExprStmtDoc(TIR(d, "reads")->Call(reads)));
    ffi::Array<ExprDoc> writes;
    for (int i = 0, n = block->writes.size(); i < n; ++i) {
      writes.push_back(d->AsDoc<ExprDoc>(block->writes[i], block_p->Attr("writes")->ArrayItem(i)));
    }
    (*frame)->stmts.push_back(ExprStmtDoc(TIR(d, "writes")->Call(writes)));
  }
  // Step 4. Handle block attributes
  if (!block->annotations.empty()) {
    (*frame)->stmts.push_back(ExprStmtDoc(
        TIR(d, "sblock_attr")
            ->Call({d->AsDoc<ExprDoc>(block->annotations, block_p->Attr("annotations"))})));
  }
  // Step 5. Handle `alloc_buffer`
  for (int i = 0, n = block->alloc_buffers.size(); i < n; ++i) {
    tirx::BufferVar buffer = block->alloc_buffers[i];
    AccessPath buffer_p = block_p->Attr("alloc_buffers")->ArrayItem(i);
    IdDoc lhs = DefineBuffer(buffer, *frame, d);
    ExprDoc rhs = BufferDecl(buffer, "sblock_alloc_buffer", {}, buffer_p, *frame, d,
                             BufferVarDefinition::DataPointer);
    (*frame)->stmts.push_back(AssignDoc(lhs, rhs, std::nullopt));
  }
  // Step 6. Handle `match_buffer`
  for (int i = 0, n = block->match_buffers.size(); i < n; ++i) {
    tirx::MatchBufferRegion buffer_region = block->match_buffers[i];
    AccessPath buffer_region_p = block_p->Attr("match_buffers")->ArrayItem(i);
    StmtDoc doc = d->AsDoc<StmtDoc>(buffer_region, buffer_region_p);
    (*frame)->stmts.push_back(doc);
  }
  // Step 7. Handle init block
  if (block->init.has_value()) {
    tirx::Stmt init = block->init.value();
    With<TIRFrame> init_frame(d, init);
    AsDocBody(init, block_p->Attr("init"), init_frame->get(), d);
    (*frame)->stmts.push_back(
        ScopeDoc(std::nullopt, TIR(d, "init")->Call({}), (*init_frame)->stmts));
  }
  // Step 8. Handle block body
  AsDocBody(block->body, block_p->Attr("body"), frame->get(), d);
  ffi::Array<ffi::String> kwargs_keys;
  ffi::Array<ExprDoc> kwargs_values;
  if (!realize) {
    kwargs_keys.push_back("no_realize");
    kwargs_values.push_back(LiteralDoc::Boolean(true, std::nullopt));
  }
  return ScopeDoc(std::nullopt,
                  TIR(d, "sblock")  //
                      ->Call({LiteralDoc::Str(block->name_hint, block_p->Attr("name_hint"))},
                             kwargs_keys, kwargs_values),
                  (*frame)->stmts);
}

TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
    .set_dispatch<tirx::SBlockRealize>(
        "", [](tirx::SBlockRealize realize, AccessPath p, IRDocsifier d) -> Doc {
          Doc doc = PrintBlock(d, realize->block, p->Attr("block"), realize, p);
          // since we do not have d->AsDoc for realize->block,
          // we should add possible doc decoration manually.
          AddDocDecoration<ScopeDoc>(doc, realize->block, p->Attr("block"), d->cfg);
          return doc;
        });

TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
    .set_dispatch<tirx::SBlock>("", [](tirx::SBlock block, AccessPath p, IRDocsifier d) -> Doc {
      return PrintBlock(d, block, p, std::nullopt, std::nullopt);
    });

TVM_REGISTER_SCRIPT_AS_REPR(tirx::SBlockNode, ReprPrintTIR);
TVM_REGISTER_SCRIPT_AS_REPR(tirx::SBlockRealizeNode, ReprPrintTIR);

TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
    .set_dispatch<tirx::ScopeIdDefStmt>(
        "", [](tirx::ScopeIdDefStmt stmt, AccessPath p, IRDocsifier d) -> Doc {
          // Render as ``(var1, var2, ...) = T.cta_id([ext], preferred=[...])``
          // (or the appropriate API name for the binding).
          TVM_FFI_ICHECK(!d->frames.empty());
          tirx::ScopeIdDef def = stmt->def;
          AccessPath def_p = p->Attr("def");
          ffi::Array<ExprDoc> lhs;
          for (auto scope_id : def->def_ids) {
            lhs.push_back(DefineVar(scope_id, d->frames.back(), d));
          }
          ffi::Array<ExprDoc> rhs_args;
          if (def->scope != tirx::ScopeBinding::kClusterCtaPair && def->extents.has_value()) {
            rhs_args.push_back(d->AsDoc<ExprDoc>(def->extents.value(), def_p->Attr("extents")));
          }
          ffi::Array<ffi::String> kwarg_keys;
          ffi::Array<ExprDoc> kwarg_vals;
          if (def->preferred_extents.has_value()) {
            kwarg_keys.push_back("preferred");
            kwarg_vals.push_back(d->AsDoc<ExprDoc>(def->preferred_extents.value(),
                                                   def_p->Attr("preferred_extents")));
          }
          ExprDoc rhs = TIR(d, ScopeIdApiName(def->scope))->Call(rhs_args, kwarg_keys, kwarg_vals);
          return AssignDoc(TupleDoc(lhs), rhs, std::nullopt);
        });

TVM_SCRIPT_REPR(tirx::ScopeIdDefStmtNode, ReprPrintTIR);

TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
    .set_dispatch<tirx::ExecScope>(
        "", [](tirx::ExecScope exec_scope, AccessPath p, IRDocsifier d) -> Doc {
          Doc doc =
              TIR(d, "ExecScope")->Call({LiteralDoc::Str(exec_scope->name(), p->Attr("name"))});
          return doc;
        });
TVM_SCRIPT_REPR(tirx::ExecScopeNode, ReprPrintTIR);

TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
    .set_dispatch<tirx::ScopeIdDef>(
        "", [](tirx::ScopeIdDef def, AccessPath p, IRDocsifier d) -> Doc {
          auto [parent, cur] = tirx::ScopeBindingToStringPair(def->scope);
          ExprDoc extents_doc = def->extents.has_value()
                                    ? d->AsDoc<ExprDoc>(def->extents.value(), p->Attr("extents"))
                                    : LiteralDoc::None(p->Attr("extents"));
          Doc doc = TIR(d, "ScopeIdDef")
                        ->Call({d->AsDoc<ExprDoc>(def->def_ids, p->Attr("def_ids")), extents_doc,
                                LiteralDoc::Str(parent, p->Attr("parent")),
                                LiteralDoc::Str(cur, p->Attr("cur"))});
          return doc;
        });
TVM_SCRIPT_REPR(tirx::ScopeIdDefNode, ReprPrintTIR);

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