/*
 * 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 <algorithm>

#include "../../../tirx/transform/ir_utils.h"  // For `GetPtrStorageScope`
#include "./utils.h"

namespace tvm {
namespace script {
namespace printer {

Doc DoConciseScoping(const ffi::Optional<ExprDoc>& lhs, const ExprDoc& rhs,
                     ffi::Array<StmtDoc>* stmts, bool concise_scoping) {
  if (concise_scoping) {
    if (lhs.has_value()) {
      stmts->insert(stmts->begin(), AssignDoc(lhs.value(), rhs, std::nullopt));
    } else {
      stmts->insert(stmts->begin(), ExprStmtDoc(rhs));
    }
    return StmtBlockDoc(*stmts);
  } else {
    return ScopeDoc(lhs, rhs, *stmts);
  }
}

bool AllowConciseScoping(const IRDocsifier& d, const ffi::ObjectRef& obj) {
  if (d->cfg.defined()) {
    if (d->cfg->obj_to_annotate.count(obj)) {
      // if the object requires annotation, do not fold this frame
      return false;
    }
  }
  TVM_FFI_ICHECK(!d->frames.empty());
  if (const auto* f = d->frames.back().as<TIRFrameNode>()) {
    return f->allow_concise_scoping;
  }
  TVM_FFI_THROW(NotImplementedError) << "fragment printing";
  TVM_FFI_UNREACHABLE();
}

bool IsAncestorOfAllVarUse(const tirx::Stmt& node, const ffi::ObjectRef& var,
                           const IRDocsifier& d) {
  if (!d->common_prefix.count(var.get())) {
    return false;
  }
  const std::vector<const ffi::Object*>& path = d->common_prefix.at(var.get());
  for (auto it = path.rbegin(); it != path.rend(); ++it) {
    if (*it == node.get()) {
      return true;
    }
  }
  return false;
}

TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
    .set_dispatch<tirx::TilePrimitiveCall>(
        "", [](tirx::TilePrimitiveCall op_call, AccessPath p, IRDocsifier d) -> Doc {
          static const OpAttrMap<tirx::TScriptPrinterName>& op_names =
              Op::GetAttrMap<tirx::TScriptPrinterName>("TScriptPrinterName");
          auto op = op_call->op;
          if (op_names.count(op) == 0) {
            LOG(WARNING) << "No TScriptPrinterName attribute for " << op->name;
          }

          static const auto& category_map = Op::GetAttrMap<tirx::TIRxOpCategory>("TIRxOpCategory");
          bool is_tile_primitive = category_map.get(op, ffi::String("")) == "tile_primitive";
          TVM_FFI_ICHECK(is_tile_primitive)
              << "Only tile primitive ops can be used in tirx::TilePrimitiveCall";
          ffi::String name = op_names.get(op, op->name);
          // Per-call execution scope is printed as a namespace prefix on the op,
          // e.g. ``T.warp.copy(...)``. ``warpgroup`` prints as ``wg``. The
          // default ``thread`` scope prints through the explicit tile namespace,
          // e.g. ``T.tile.copy(...)``, so canonical script only needs the full
          // TIRx dialect import. ``Tx`` remains a handwritten shorthand for
          // ``T.tile`` and ``T.<scope>`` tile calls.
          auto scope_ns = [](tirx::ScopeKind k) -> ffi::Optional<ffi::String> {
            switch (k) {
              case tirx::ScopeKind::kWarp:
                return ffi::String("warp");
              case tirx::ScopeKind::kWarpgroup:
                return ffi::String("wg");
              case tirx::ScopeKind::kCta:
                return ffi::String("cta");
              case tirx::ScopeKind::kCluster:
                return ffi::String("cluster");
              default:  // kThread -> no prefix
                return std::nullopt;
            }
          };
          auto scoped_callee = [&](const ffi::String& op_name) -> ExprDoc {
            ffi::Optional<ffi::String> ns = scope_ns(op_call->scope->kind);
            if (ns.has_value()) {
              return TIRx(d, ns.value())->Attr(op_name);
            }
            return TIRx(d, "tile")->Attr(op_name);
          };
          if (!op.same_as(tirx::compose_op())) {
            // Trim trailing None args (e.g. optional bias=None, scale=None)
            size_t n_args = op_call->args.size();
            while (n_args > 0 &&
                   op_call->args[n_args - 1].type_index() == ffi::TypeIndex::kTVMFFINone) {
              --n_args;
            }
            // Detect in-place unary ops: after trimming Nones, if exactly 2 args
            // and args[0]/args[1] refer to the same buffer region, collapse to 1 arg
            bool inplace_unary = false;
            if (n_args == 2) {
              auto dst_opt = op_call->args[0].as<tirx::BufferRegion>();
              auto src_opt = op_call->args[1].as<tirx::BufferRegion>();
              if (dst_opt.has_value() && src_opt.has_value() &&
                  dst_opt.value()->buffer.same_as(src_opt.value()->buffer) &&
                  StructuralEqual()(dst_opt.value()->region, src_opt.value()->region)) {
                inplace_unary = true;
              }
            }
            ffi::Array<Doc> args;
            for (size_t i = 0; i < n_args; ++i) {
              if (inplace_unary && i == 1) continue;  // skip duplicate src
              args.push_back(d->AsDoc<Doc>(op_call->args[i], p->Attr("args")->ArrayItem(i)));
            }
            ffi::Optional<ExprDoc> disp = std::nullopt;
            if (op_call->dispatch.has_value()) {
              disp = LiteralDoc::Str(op_call->dispatch.value(), p->Attr("dispatch"));
            }
            return OpCallDoc(scoped_callee(name), args,
                             d->AsDoc<DictDoc>(op_call->workspace, p->Attr("workspace")),
                             d->AsDoc<DictDoc>(op_call->config, p->Attr("config")), disp);
          } else {
            With<TIRFrame> f(d, op_call);
            ffi::Array<tirx::Stmt> stmts;
            for (size_t i = 0, n = op_call->args.size(); i < n; ++i) {
              stmts.push_back(op_call->args[i].as_or_throw<tirx::Stmt>());
            }
            tirx::SeqStmt seq_stmt(stmts);
            AsDocBody(seq_stmt, p->Attr("args"), f->get(), d);
            // Build kwargs: workspace, dispatch, then flatten config
            ffi::Array<ffi::String> kw_keys;
            ffi::Array<ExprDoc> kw_values;
            if (!op_call->workspace.empty()) {
              kw_keys.push_back("workspace");
              kw_values.push_back(d->AsDoc<DictDoc>(op_call->workspace, p->Attr("workspace")));
            }
            if (op_call->dispatch.has_value()) {
              kw_keys.push_back("dispatch");
              kw_values.push_back(LiteralDoc::Str(op_call->dispatch.value(), p->Attr("dispatch")));
            }
            using POO = std::pair<ffi::String, ffi::Any>;
            std::vector<POO> items{op_call->config.begin(), op_call->config.end()};
            std::sort(items.begin(), items.end(),
                      [](const POO& a, const POO& b) { return a.first < b.first; });
            for (const auto& kv : items) {
              kw_keys.push_back(kv.first);
              kw_values.push_back(
                  d->AsDoc<ExprDoc>(kv.second, p->Attr("config")->MapItem(kv.first)));
            }
            return ScopeDoc(std::nullopt, scoped_callee("compose_op")->Call({}, kw_keys, kw_values),
                            (*f)->stmts);
          }
        });
TVM_SCRIPT_REPR(tirx::TilePrimitiveCallNode, ReprPrintTIR);

TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
    .set_dispatch<tirx::Evaluate>("", [](tirx::Evaluate eval, AccessPath p, IRDocsifier d) -> Doc {
      ExprDoc value = d->AsDoc<ExprDoc>(eval->value, p->Attr("value"));
      const auto* call = eval->value.as<CallNode>();
      if (call && !call->op.same_as(tirx::builtin::buffer_data())) {
        return ExprStmtDoc(value);
      }
      return ExprStmtDoc(TIR(d, "evaluate")->Call({value}));
    });

TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
    .set_dispatch<tirx::Return>("", [](tirx::Return stmt, AccessPath p, IRDocsifier d) -> Doc {
      ExprDoc value = d->AsDoc<ExprDoc>(stmt->value, p->Attr("value"));
      return ReturnDoc(value);
    });

TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
    .set_dispatch<tirx::Bind>("", [](tirx::Bind stmt, AccessPath p, IRDocsifier d) -> Doc {
      // Step 1. Type annotation
      TVM_FFI_ICHECK(!stmt->var->ty.IsMissing())
          << "Type annotation is required for variable: " << stmt->var->name;
      ffi::Optional<ExprDoc> type_doc = d->AsDoc<ExprDoc>(stmt->var->ty,  //
                                                          p->Attr("var")->Attr("ty"));
      if (const auto* tuple_type = stmt->var->ty.as<TupleTypeNode>()) {
        if (tuple_type->fields.empty()) {
          type_doc = std::nullopt;
        }
      }
      // Step 2. RHS
      ExprDoc rhs = d->AsDoc<ExprDoc>(stmt->value, p->Attr("value"));
      // Step 3. LHS - Bind is flat, define var if new, otherwise just assign
      if (!d->IsVarDefined(stmt->var)) {
        TVM_FFI_ICHECK(!d->frames.empty());
        ExprDoc lhs = DefineVar(stmt->var, d->frames.back(), d);
        ExprDoc let_ann = type_doc.has_value()
                              ? ExprDoc(IndexDoc(TIR(d, "let"), {type_doc.value()}))
                              : TIR(d, "let");
        return AssignDoc(lhs, rhs, let_ann);
      } else {
        ExprDoc lhs = d->AsDoc<ExprDoc>(stmt->var, p->Attr("var"));
        return AssignDoc(lhs, rhs, std::nullopt);
      }
    });

TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
    .set_dispatch<tirx::AssertStmt>(
        "", [](tirx::AssertStmt stmt, AccessPath p, IRDocsifier d) -> Doc {
          ExprDoc cond = d->AsDoc<ExprDoc>(stmt->condition, p->Attr("condition"));
          // Always emit the canonical tuple form: assert cond, ("Kind", ["part0", "part1", ...])
          ffi::Array<ExprDoc> parts;
          auto parts_path = p->Attr("message_parts");
          for (size_t i = 0; i < stmt->message_parts.size(); ++i) {
            parts.push_back(d->AsDoc<ExprDoc>(stmt->message_parts[i], parts_path->ArrayItem(i)));
          }
          ExprDoc kind_doc = d->AsDoc<ExprDoc>(stmt->error_kind, p->Attr("error_kind"));
          return AssertDoc(cond, TupleDoc({kind_doc, ListDoc(parts)}));
        });

TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
    .set_dispatch<tirx::While>("", [](tirx::While stmt, AccessPath p, IRDocsifier d) -> Doc {
      ExprDoc cond = d->AsDoc<ExprDoc>(stmt->condition, p->Attr("condition"));
      With<TIRFrame> f(d, stmt);
      AsDocBody(stmt->body, p->Attr("body"), f->get(), d);
      return WhileDoc(cond, (*f)->stmts);
    });

TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
    .set_dispatch<tirx::Break>("", [](tirx::Break stmt, AccessPath p, IRDocsifier d) -> Doc {
      return BreakDoc();
    });

TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
    .set_dispatch<tirx::Continue>("", [](tirx::Continue stmt, AccessPath p, IRDocsifier d) -> Doc {
      return ContinueDoc();
    });

namespace {

/*!
 * \brief Find the parent buffer named by a child's explicit data projection.
 * \param child The child buffer.
 * \param data The child's explicit source pointer, if any.
 * \param d The IRDocsifier.
 * \return A list of candidate parent buffers.
 */
std::vector<tirx::BufferVar> FindParentBuffers(const tirx::BufferVar& child,
                                               const ffi::Optional<Expr>& data,
                                               const IRDocsifier& d) {
  if (!data.has_value()) {
    return {};
  }
  const auto* call = data.value().as<CallNode>();
  if (call == nullptr || !call->op.same_as(tirx::builtin::buffer_data()) ||
      call->args.size() != 1) {
    return {};
  }
  auto parent_var = call->args[0].as<tirx::Var>();
  if (!parent_var.has_value() || !parent_var.value()->ty.as<tirx::BufferTypeNode>()) {
    return {};
  }
  tirx::BufferVar parent(parent_var.value());
  if (parent.same_as(child) || !d->GetVarDoc(parent).has_value()) {
    return {};
  }
  return {parent};
}

/*!
 * \brief Check if a layout is the default layout for a given shape.
 */
bool IsDefaultLayout(const ffi::Optional<tirx::Layout>& layout, const ffi::Array<PrimExpr>& shape) {
  if (!layout.has_value()) return false;
  return StructuralEqual()(layout.value(), tirx::TileLayoutNode::DefaultLayout(shape));
}

/*!
 * \brief Try to produce a DeclBuffer sugar expression for the given child buffer
 *        with respect to a specific parent buffer.
 *
 * Returns std::nullopt if no sugar pattern matches.
 */
ffi::Optional<ExprDoc> TryDeclBufferSugarWithParent(const tirx::BufferVar& child,
                                                    const AccessPath& p, const IRDocsifier& d,
                                                    const tirx::BufferVar& parent,
                                                    bool require_same_layout) {
  ffi::Optional<ExprDoc> parent_doc = d->GetVarDoc(parent);
  if (!parent_doc.has_value()) return std::nullopt;
  ExprDoc pdoc = parent_doc.value();

  tirx::ExprDeepEqual expr_equal;

  // Check elem_offset equality
  bool same_elem_offset = expr_equal(child->elem_offset, parent->elem_offset);
  // Check dtype equality
  bool same_dtype = (child->dtype == parent->dtype);
  // Check shape equality
  bool same_shape = (child->shape.size() == parent->shape.size());
  if (same_shape) {
    for (size_t i = 0; i < child->shape.size(); ++i) {
      if (!expr_equal(child->shape[i], parent->shape[i])) {
        same_shape = false;
        break;
      }
    }
  }
  bool same_strides = (child->strides.size() == parent->strides.size());
  if (same_strides) {
    for (size_t i = 0; i < child->strides.size(); ++i) {
      if (!expr_equal(child->strides[i], parent->strides[i])) {
        same_strides = false;
        break;
      }
    }
  }

  bool child_is_default = IsDefaultLayout(child->layout, child->shape);
  bool parent_is_default = IsDefaultLayout(parent->layout, parent->shape);

  // NOTE: an earlier sugar printed rank-preserving aliases with a different
  // elem_offset as ``parent[slices]``. That print is not roundtrippable: it
  // reparses as a BufferRegion, not a Buffer, so any later Buffer use of the
  // alias (stores, views) breaks. Such aliases now print as plain
  // T.decl_buffer, which reparses exactly.

  // Differences in these Buffer fields cannot be expressed by the alias sugar
  // below, so conservatively fall back to T.decl_buffer.
  // Shape, strides, elem_offset, dtype, and layout are checked by each helper
  // because those are the fields that individual transformations may change.
  //
  // The explicit data projection was checked by FindParentBuffers.  name/span
  // do not participate in structural equality, and BufferTypeNode has no
  // axis-separators field (unlike tir::Buffer).
  bool same_common_metadata = child.scope() == parent.scope() &&
                              child->data_alignment == parent->data_alignment &&
                              child->offset_factor == parent->offset_factor &&
                              StructuralEqual()(child->allocated_addr, parent->allocated_addr);
  if (!same_common_metadata) return std::nullopt;

  // --- (b) Local: parent has thread axes and child spans its physical storage ---
  if (same_elem_offset && same_dtype && same_strides && !parent_is_default &&
      parent->layout.has_value() && child->layout.has_value()) {
    if (auto* parent_tile = parent->layout.value().as<tirx::TileLayoutNode>()) {
      if (parent_tile->HasThreadAxis()) {
        // Compute the raw physical storage span after filtering thread axes.
        std::vector<tirx::Iter> storage_shard;
        std::vector<tirx::Iter> storage_replica;
        ffi::Map<tirx::Axis, PrimExpr> storage_offset;
        for (const auto& iter : parent_tile->shard) {
          if (!iter->axis->IsThreadAxis()) {
            storage_shard.push_back(iter);
          }
        }
        for (const auto& iter : parent_tile->replica) {
          if (!iter->axis->IsThreadAxis()) {
            storage_replica.push_back(iter);
          }
        }
        for (const auto& [axis, off] : parent_tile->offset) {
          if (!axis->IsThreadAxis()) {
            storage_offset.Set(axis, off);
          }
        }
        tirx::TileLayout expected_storage(
            ffi::Array<tirx::Iter>(storage_shard.begin(), storage_shard.end()),
            ffi::Array<tirx::Iter>(storage_replica.begin(), storage_replica.end()), storage_offset);

        PrimExpr storage_span = expected_storage->GetSpan(ffi::Optional<ffi::String>());
        PrimExpr storage_size = expected_storage->GetSize(ffi::Optional<ffi::String>());
        PrimExpr child_total = IntImm::Int32(1);
        for (const PrimExpr& dim : child->shape) {
          child_total = child_total * dim;
        }
        arith::Analyzer analyzer;
        bool default_physical =
            child_is_default && analyzer->CanProveEqual(child_total, storage_span);
        bool child_has_thread_axis = false;
        if (const auto* child_tile = child->layout.value().as<tirx::TileLayoutNode>()) {
          child_has_thread_axis = child_tile->HasThreadAxis();
        }
        bool explicit_override = !default_physical && !child_has_thread_axis;
        if (default_physical || explicit_override) {
          PrimExpr expected_extent = default_physical ? storage_span : storage_size;
          bool auto_shape =
              child->shape.size() == 1 && analyzer->CanProveEqual(child->shape[0], expected_extent);
          ffi::Array<ExprDoc> args;
          if (!auto_shape) {
            for (size_t i = 0; i < child->shape.size(); ++i) {
              args.push_back(d->AsDoc<ExprDoc>(child->shape[i],
                                               p->Attr("buffer")->Attr("shape")->ArrayItem(i)));
            }
          }
          ffi::Array<ffi::String> kwargs_keys;
          ffi::Array<ExprDoc> kwargs_values;
          if (explicit_override) {
            kwargs_keys.push_back("layout");
            kwargs_values.push_back(
                d->AsDoc<ExprDoc>(child->layout.value(), p->Attr("buffer")->Attr("layout")));
          }
          return pdoc->Attr("local")->Call(args, kwargs_keys, kwargs_values);
        }
      }
    }
  }

  // --- (c) View(dtype): different dtype, same elem_offset ---
  if (same_elem_offset && !same_dtype && child->shape.size() == parent->shape.size()) {
    // Verify shape compatibility with dtype reinterpret cast
    int child_bits = child->dtype.bits();
    int parent_bits = parent->dtype.bits();
    bool shapes_compatible = true;
    // All dims except last must match
    for (size_t i = 0; i + 1 < child->shape.size(); ++i) {
      if (!expr_equal(child->shape[i], parent->shape[i])) {
        shapes_compatible = false;
        break;
      }
    }
    if (shapes_compatible && !child->shape.empty()) {
      auto* child_last = child->shape.back().as<IntImmNode>();
      auto* parent_last = parent->shape.back().as<IntImmNode>();
      if (child_last && parent_last) {
        if (child_bits > parent_bits) {
          // Cast up: child_last = parent_last / ratio
          int ratio = child_bits / parent_bits;
          shapes_compatible = (parent_last->value == child_last->value * ratio);
        } else {
          // Cast down: child_last = parent_last * ratio
          int ratio = parent_bits / child_bits;
          shapes_compatible = (child_last->value == parent_last->value * ratio);
        }
      } else {
        shapes_compatible = false;
      }
    }
    // Also verify the parent's layout is compatible with the pack/unpack operation
    if (shapes_compatible && parent->layout.has_value()) {
      if (auto* ptile = parent->layout.value().as<tirx::TileLayoutNode>()) {
        if (!ptile->shard.empty() && child_bits > parent_bits) {
          // Cast up requires pack: last shard iter must have stride=1
          // and extent divisible by ratio
          const auto& last_iter = ptile->shard.back();
          auto* last_stride = last_iter->stride.as<IntImmNode>();
          auto* last_extent = last_iter->extent.as<IntImmNode>();
          int ratio = child_bits / parent_bits;
          if (!last_stride || last_stride->value != 1 || !last_extent ||
              last_extent->value % ratio != 0) {
            shapes_compatible = false;
          }
        }
      }
    }
    if (shapes_compatible) {
      ExprDoc dtype_doc =
          LiteralDoc::Str(DType2Str(child->dtype->dtype), p->Attr("buffer")->Attr("dtype"));
      return pdoc->Attr("view")->Call({dtype_doc});
    }
  }

  // --- (d) Permute: child shape is a permutation of parent shape, same elem_offset ---
  if (same_elem_offset && same_dtype && !same_shape &&
      child->shape.size() == parent->shape.size()) {
    // Try to find a permutation
    std::vector<int> perm(child->shape.size(), -1);
    std::vector<bool> used(parent->shape.size(), false);
    bool is_permutation = true;
    for (size_t i = 0; i < child->shape.size(); ++i) {
      bool found = false;
      for (size_t j = 0; j < parent->shape.size(); ++j) {
        if (!used[j] && expr_equal(child->shape[i], parent->shape[j])) {
          perm[i] = j;
          used[j] = true;
          found = true;
          break;
        }
      }
      if (!found) {
        is_permutation = false;
        break;
      }
    }
    // Check it's not identity
    bool is_identity = is_permutation;
    if (is_permutation) {
      for (size_t i = 0; i < perm.size(); ++i) {
        if (perm[i] != static_cast<int>(i)) {
          is_identity = false;
          break;
        }
      }
    }
    if (is_permutation && !is_identity) {
      // Verify the layout matches permutation by comparing shard iters directly
      bool layout_matches = false;
      if (parent->layout.has_value() && child->layout.has_value()) {
        auto* parent_tile = parent->layout.value().as<tirx::TileLayoutNode>();
        auto* child_tile = child->layout.value().as<tirx::TileLayoutNode>();
        if (parent_tile && child_tile && parent_tile->shard.size() == child_tile->shard.size()) {
          StructuralEqual seq;
          layout_matches = true;
          for (size_t i = 0; i < perm.size(); ++i) {
            if (!seq(child_tile->shard[i], parent_tile->shard[perm[i]])) {
              layout_matches = false;
              break;
            }
          }
          // Also check replica and offset are unchanged
          if (layout_matches) {
            layout_matches = seq(child_tile->replica, parent_tile->replica) &&
                             seq(child_tile->offset, parent_tile->offset);
          }
        }
      }
      if (layout_matches) {
        ffi::Array<ExprDoc> args;
        for (int idx : perm) {
          args.push_back(LiteralDoc::Int(idx, p->Attr("buffer")->Attr("shape")));
        }
        return pdoc->Attr("permute")->Call(args);
      }
    }
  }

  // --- (e) Partition: child has 2*parent_ndim dims with grid+tile strides ---
  if (same_elem_offset && same_dtype && !parent->shape.empty() &&
      child->shape.size() == 2 * parent->shape.size() && !child->strides.empty() &&
      child->strides.size() == 2 * parent->shape.size()) {
    size_t ndim = parent->shape.size();
    // Compute parent's row-major strides
    std::vector<int64_t> parent_rm_strides(ndim);
    int64_t stride = 1;
    bool all_const = true;
    for (int i = static_cast<int>(ndim) - 1; i >= 0; --i) {
      parent_rm_strides[i] = stride;
      if (auto* s = parent->shape[i].as<IntImmNode>()) {
        stride *= s->value;
      } else {
        all_const = false;
        break;
      }
    }
    if (all_const) {
      bool is_partition = true;
      for (size_t i = 0; i < ndim; ++i) {
        auto* grid_dim = child->shape[i].as<IntImmNode>();
        auto* tile_dim = child->shape[ndim + i].as<IntImmNode>();
        auto* parent_dim = parent->shape[i].as<IntImmNode>();
        auto* grid_stride = child->strides[i].as<IntImmNode>();
        auto* tile_stride = child->strides[ndim + i].as<IntImmNode>();
        if (!grid_dim || !tile_dim || !parent_dim || !grid_stride || !tile_stride) {
          is_partition = false;
          break;
        }
        // grid × tile == parent dim
        if (grid_dim->value * tile_dim->value != parent_dim->value) {
          is_partition = false;
          break;
        }
        // inner strides match parent's row-major strides
        if (tile_stride->value != parent_rm_strides[i]) {
          is_partition = false;
          break;
        }
        // grid stride == tile_dim × inner stride
        if (grid_stride->value != tile_dim->value * tile_stride->value) {
          is_partition = false;
          break;
        }
      }
      if (is_partition) {
        ffi::Array<ExprDoc> tuple_elems;
        for (size_t i = 0; i < ndim; ++i) {
          tuple_elems.push_back(
              d->AsDoc<ExprDoc>(child->shape[i], p->Attr("buffer")->Attr("shape")->ArrayItem(i)));
        }
        return pdoc->Attr("partition")->Call({}, {"num_tiles"}, {TupleDoc(tuple_elems)});
      }
    }
  }

  // --- (f) View(*shape, layout=L): different shape/layout, same dtype and elem_offset ---
  if (same_elem_offset && same_dtype && !same_shape) {
    // Buffer.view(...) copies the parent's strides onto the child (see
    // python/tvm/tirx/buffer.py:view). If parent has strides but child
    // doesn't (or vice versa), the sugar can't faithfully round-trip
    // through view — fall back to T.decl_buffer where strides is an
    // explicit kwarg.
    if (!same_strides) return std::nullopt;

    ffi::Array<ExprDoc> args;
    ffi::Array<ffi::String> kwargs_keys;
    ffi::Array<ExprDoc> kwargs_values;
    for (size_t i = 0; i < child->shape.size(); ++i) {
      args.push_back(
          d->AsDoc<ExprDoc>(child->shape[i], p->Attr("buffer")->Attr("shape")->ArrayItem(i)));
    }
    // Check if layout differs
    bool same_layout = false;
    if (child->layout.has_value() && parent->layout.has_value()) {
      same_layout = StructuralEqual()(child->layout.value(), parent->layout.value());
    } else if (!child->layout.has_value() && !parent->layout.has_value()) {
      same_layout = true;
    }
    // First pass prefers a parent whose layout matches structurally, so the
    // sugar prints as a bare reshape instead of restating the layout.
    if (require_same_layout && !same_layout) return std::nullopt;
    // Default layouts are shape-specific objects, but a default-to-default
    // reshape is still represented by view(*shape) without an explicit layout.
    if (!same_layout && !(child_is_default && parent_is_default)) {
      // Buffer.view(..., layout=None) means "inherit the parent layout", so it
      // cannot reconstruct a layout-less child from a laid-out parent.
      if (!child->layout.has_value()) return std::nullopt;
      kwargs_keys.push_back("layout");
      kwargs_values.push_back(
          d->AsDoc<ExprDoc>(child->layout.value(), p->Attr("buffer")->Attr("layout")));
    }
    return pdoc->Attr("view")->Call(args, kwargs_keys, kwargs_values);
  }

  return std::nullopt;
}

/*!
 * \brief Try to produce a DeclBuffer sugar expression, trying all parent buffer candidates.
 */
ffi::Optional<ExprDoc> TryDeclBufferSugar(const tirx::BufferVar& child, const AccessPath& p,
                                          const ffi::Optional<Expr>& data, const IRDocsifier& d) {
  auto parents = FindParentBuffers(child, data, d);
  for (const auto& parent : parents) {
    if (auto sugar = TryDeclBufferSugarWithParent(child, p, d, parent,
                                                  /*require_same_layout=*/true)) {
      return sugar;
    }
  }
  for (const auto& parent : parents) {
    if (auto sugar = TryDeclBufferSugarWithParent(child, p, d, parent,
                                                  /*require_same_layout=*/false)) {
      return sugar;
    }
  }
  return std::nullopt;
}

Doc DeclBufferDoc(tirx::DeclBuffer stmt, AccessPath p, IRDocsifier d,
                  BufferVarDefinition var_definitions) {
  // Try sugar detection when syntax_sugar is enabled
  if (d->cfg->syntax_sugar) {
    if (auto sugar = TryDeclBufferSugar(stmt->buffer, p, stmt->data, d)) {
      ExprDoc lhs = DefineBuffer(stmt->buffer, d->frames.back(), d);
      return AssignDoc(lhs, sugar.value(), std::nullopt);
    }
  }
  ExprDoc rhs = BufferDecl(stmt->buffer, "decl_buffer", {}, p->Attr("buffer"), d->frames.back(), d,
                           var_definitions, stmt->data);
  ExprDoc lhs = DefineBuffer(stmt->buffer, d->frames.back(), d);
  return AssignDoc(lhs, rhs, std::nullopt);
}
}  // namespace

TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
    .set_dispatch<tirx::DeclBuffer>(  //
        "", [](tirx::DeclBuffer stmt, AccessPath p, IRDocsifier d) -> Doc {
          return DeclBufferDoc(stmt, p, d, BufferVarDefinition::None);
        });

namespace {
Doc AllocBufferDoc(tirx::AllocBuffer stmt, AccessPath p, IRDocsifier d) {
  if (d->cfg->syntax_sugar && stmt->buffer.IsScalar(true)) {
    ExprDoc lhs = DefineBuffer(stmt->buffer, d->frames.back(), d);
    ExprDoc type_ann = TIR(d, DType2Str(stmt->buffer->dtype->dtype));
    return AssignDoc(lhs, std::nullopt, type_ann);
  }
  ExprDoc rhs = BufferDecl(stmt->buffer, "alloc_buffer", {}, p->Attr("buffer"), d->frames.back(), d,
                           BufferVarDefinition::DataPointer);
  // alloc_buffer carries an `annotations` field on the IR node that BufferDecl
  // doesn't know about. When non-empty, append it as an `annotations=...`
  // kwarg on the emitted call so round-trip preserves the annotation map.
  if (!stmt->annotations.empty()) {
    if (const auto* call = rhs.as<CallDocNode>()) {
      ffi::Array<ffi::String> new_keys = call->kwargs_keys;
      ffi::Array<ExprDoc> new_values = call->kwargs_values;
      new_keys.push_back("annotations");
      new_values.push_back(d->AsDoc<ExprDoc>(stmt->annotations, p->Attr("annotations")));
      rhs = CallDoc(call->callee, call->args, new_keys, new_values);
    }
  }
  ExprDoc lhs = DefineBuffer(stmt->buffer, d->frames.back(), d);
  return AssignDoc(lhs, rhs, std::nullopt);
}

}  // namespace

TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
    .set_dispatch<tirx::AllocBuffer>(  //
        "", [](tirx::AllocBuffer stmt, AccessPath p, IRDocsifier d) -> Doc {
          return AllocBufferDoc(stmt, p, d);
        });

TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
    .set_dispatch<tirx::IfThenElse>(  //
        "", [](tirx::IfThenElse stmt, AccessPath p, IRDocsifier d) -> Doc {
          ExprDoc cond = d->AsDoc<ExprDoc>(stmt->condition, p->Attr("condition"));
          ffi::Array<StmtDoc> then_branch;
          ffi::Array<StmtDoc> else_branch;
          if (stmt->then_case.defined()) {
            With<TIRFrame> f(d, stmt->then_case);
            AsDocBody(stmt->then_case, p->Attr("then_case"), f->get(), d);
            then_branch = (*f)->stmts;
          }
          if (stmt->else_case.has_value()) {
            With<TIRFrame> f(d, stmt->else_case.value());
            AsDocBody(stmt->else_case.value(), p->Attr("else_case"), f->get(), d);
            else_branch = (*f)->stmts;
          }
          return IfDoc(cond, then_branch, else_branch);
        });

TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
    .set_dispatch<tirx::SeqStmt>("", [](tirx::SeqStmt stmt, AccessPath p, IRDocsifier d) -> Doc {
      With<TIRFrame> f(d, stmt);
      AsDocBody(stmt, p, f->get(), d);
      return StmtBlockDoc((*f)->stmts);
    });

void InsertEnvThread(const tirx::IterVar& iter_var, const AccessPath& iter_var_p,
                     const IRDocsifier& d) {
  Frame f = FindLowestVarDef(iter_var->var, d).value();
  DefineVar(iter_var->var, f, d);
  ExprDoc rhs = TIR(d, "env_thread")
                    ->Call({LiteralDoc::Str(iter_var->thread_tag,  //
                                            iter_var_p->Attr("thread_tag"))});
  ExprDoc lhs = d->AsDoc<ExprDoc>(iter_var->var, iter_var_p->Attr("var"));
  f->stmts.push_back(AssignDoc(lhs, rhs, std::nullopt));
}

ExprDoc DocsifyLaunchThread(const tirx::AttrStmt& attr_stmt, const AccessPath& attr_stmt_p,
                            ffi::Optional<tirx::Var>* define_var, const IRDocsifier& d) {
  tirx::IterVar iter_var = attr_stmt->node.as_or_throw<tirx::IterVar>();
  AccessPath iter_var_p = attr_stmt_p->Attr("node");

  ExprDoc var_doc{ffi::UnsafeInit()};
  if (d->IsVarDefined(iter_var->var)) {
    var_doc = d->AsDoc<ExprDoc>(iter_var->var, iter_var_p->Attr("var"));
  } else if (IsAncestorOfAllVarUse(attr_stmt, iter_var->var, d)) {
    var_doc = LiteralDoc::Str(iter_var->thread_tag, iter_var_p->Attr("thread_tag"));
    *define_var = iter_var->var;
  } else {
    InsertEnvThread(iter_var, iter_var_p, d);
    var_doc = d->AsDoc<ExprDoc>(iter_var->var, iter_var_p->Attr("var"));
  }
  return TIR(d, "launch_thread")
      ->Call({
          var_doc,
          d->AsDoc<ExprDoc>(attr_stmt->value, attr_stmt_p->Attr("value")),
      });
}

/*! \brief Check whether an AttrStmt has node=0 (the dict-attr pattern). */
static bool IsDictAttrPattern(const tirx::AttrStmt& stmt) {
  if (auto int_value = stmt->node.as<int64_t>()) {
    return int_value.value() == 0;
  }
  return false;
}

TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
    .set_dispatch<tirx::AttrStmt>(  //
        "", [](tirx::AttrStmt stmt, AccessPath stmt_p, IRDocsifier d) -> Doc {
          bool concise = AllowConciseScoping(d, stmt);
          ffi::Optional<ExprDoc> lhs = std::nullopt;
          ffi::Optional<ExprDoc> rhs = std::nullopt;
          ffi::Optional<tirx::Var> define_var = std::nullopt;
          tirx::Stmt body = stmt->body;
          AccessPath body_p = stmt_p->Attr("body");
          if (stmt->attr_key == "thread_extent" || stmt->attr_key == "virtual_thread") {
            if (stmt->node.as<tirx::IterVarNode>()) {
              rhs = DocsifyLaunchThread(stmt, stmt_p, &define_var, d);
            }
          }
          if (stmt->attr_key == "tirx_hint") {
            if (auto map_node = stmt->node.as<ffi::Map<ffi::String, ffi::Any>>()) {
              ffi::Array<ExprDoc> args;
              ffi::Array<ffi::String> kwargs_keys;
              ffi::Array<ExprDoc> kwargs_values;
              for (const auto& [k, v] : map_node.value()) {
                if (k == "message") {
                  auto s = v.as<ffi::String>().value();
                  args.push_back(LiteralDoc::Str(s, stmt_p->Attr("node")));
                } else {
                  kwargs_keys.push_back(k);
                  kwargs_values.push_back(d->AsDoc<ExprDoc>(v, stmt_p->Attr("node")));
                }
              }
              rhs = TIR(d, "hint")->Call(args, kwargs_keys, kwargs_values);
            }
          }
          if (!rhs.has_value()) {
            // Try to collapse consecutive dict-attr-pattern AttrStmts into T.attr({...})
            if (IsDictAttrPattern(stmt)) {
              ffi::Array<ExprDoc> keys;
              ffi::Array<ExprDoc> values;
              tirx::AttrStmt cur = stmt;
              AccessPath cur_p = stmt_p;
              while (true) {
                keys.push_back(LiteralDoc::Str(cur->attr_key, cur_p->Attr("attr_key")));
                values.push_back(d->AsDoc<ExprDoc>(cur->value, cur_p->Attr("value")));
                if (auto next = cur->body.as<tirx::AttrStmt>()) {
                  if (IsDictAttrPattern(next.value())) {
                    cur = next.value();
                    cur_p = cur_p->Attr("body");
                    continue;
                  }
                }
                body = cur->body;
                body_p = cur_p->Attr("body");
                break;
              }
              rhs = TIR(d, "attr")->Call({DictDoc(keys, values)});
            } else {
              rhs = TIR(d, "attr")->Call({
                  d->AsDoc<ExprDoc>(stmt->node, stmt_p->Attr("node")),
                  LiteralDoc::Str(stmt->attr_key, stmt_p->Attr("attr_key")),
                  d->AsDoc<ExprDoc>(stmt->value, stmt_p->Attr("value")),
              });
            }
          }
          With<TIRFrame> f(d, stmt);
          if (define_var.has_value()) {
            lhs = DefineVar(define_var.value(), *f, d);
          }
          AsDocBody(body, body_p, f->get(), d);
          return DoConciseScoping(lhs, rhs.value(), &(*f)->stmts, concise);
        });

TVM_SCRIPT_REPR(tirx::BindNode, ReprPrintTIR);
TVM_SCRIPT_REPR(tirx::AttrStmtNode, ReprPrintTIR);
TVM_SCRIPT_REPR(tirx::AssertStmtNode, ReprPrintTIR);
TVM_SCRIPT_REPR(tirx::WhileNode, ReprPrintTIR);
TVM_SCRIPT_REPR(tirx::AllocBufferNode, ReprPrintTIR);
TVM_SCRIPT_REPR(tirx::ReturnNode, ReprPrintTIR);
TVM_SCRIPT_REPR(tirx::BreakNode, ReprPrintTIR);
TVM_SCRIPT_REPR(tirx::ContinueNode, ReprPrintTIR);
TVM_SCRIPT_REPR(tirx::DeclBufferNode, ReprPrintTIR);
TVM_SCRIPT_REPR(tirx::SeqStmtNode, ReprPrintTIR);
TVM_SCRIPT_REPR(tirx::IfThenElseNode, ReprPrintTIR);
TVM_SCRIPT_REPR(tirx::EvaluateNode, ReprPrintTIR);
}  // namespace printer
}  // namespace script
}  // namespace tvm
