/*
 * 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/error.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/script/printer/doc.h>
#include <tvm/tirx/tile_primitive.h>

#include <algorithm>
#include <cmath>
#include <map>
#include <optional>
#include <sstream>
#include <streambuf>
#include <string>

#include "../../../support/str_escape.h"
#include "../../../support/utils.h"
#include "./base_doc_printer.h"

namespace tvm {
namespace script {
namespace printer {
namespace {

/*!
 * \brief Escape an ExprStringDoc child while preserving final output positions.
 *
 * ExprStringDoc must print its child through the active PythonDocPrinter so that the child's
 * precedence, printer configuration, and source-path bookkeeping stay in the current traversal.
 * However, escaping after printing into a temporary buffer would make PrintDoc record positions
 * in that unescaped buffer, rather than the final Python source. This scoped adapter instead
 * replaces the output stream's buffer only while the child is printed. It escapes each write and
 * forwards it immediately to the original final buffer; the caller emits the surrounding quotes
 * outside the scope so that only the child is transformed.
 *
 * The adapter owns neither the output stream nor its original buffer. Both outlive this stack
 * object. The constructor saves the original buffer in destination_ and installs this adapter;
 * the destructor restores the original buffer before the caller continues writing to the stream.
 *
 * PrintDoc uses tellp() before and after every child to record source spans. tellp() reaches
 * seekoff(0, cur, out), which is the only seek this adapter accepts and delegates directly to the
 * original buffer. Because that buffer already contains the preceding output and advances by the
 * escaped byte count, child spans use absolute positions in the final source, including any
 * expansion introduced by StrEscape.
 *
 * A raw newline violates the one-line expression-string contract. xsputn records its presence
 * while still escaping and forwarding the write, and the caller checks saw_newline() before
 * completing the literal. On every exit path, including exception unwinding, the noexcept
 * destructor restores the original buffer.
 */
class ScopedExprStringEscapeBuf : public std::streambuf {
 public:
  explicit ScopedExprStringEscapeBuf(std::ostream* output)
      : output_(output), destination_(output->rdbuf()) {
    output_->rdbuf(this);
  }

  ~ScopedExprStringEscapeBuf() noexcept { output_->rdbuf(destination_); }

  ScopedExprStringEscapeBuf(const ScopedExprStringEscapeBuf&) = delete;
  ScopedExprStringEscapeBuf& operator=(const ScopedExprStringEscapeBuf&) = delete;

  bool saw_newline() const { return saw_newline_; }

 protected:
  std::streamsize xsputn(const char* data, std::streamsize count) final {
    if (count <= 0) return 0;
    saw_newline_ = saw_newline_ || std::find(data, data + count, '\n') != data + count;
    // StrEscape is byte-wise, so each write can be transformed independently and forwarded
    // without retaining or copying the complete output. Report consumed input bytes, not the
    // potentially larger number of escaped bytes written to the destination.
    std::string escaped = support::StrEscape(data, static_cast<size_t>(count));
    destination_->sputn(escaped.data(), escaped.size());
    return count;
  }

  int_type overflow(int_type ch) final {
    if (traits_type::eq_int_type(ch, traits_type::eof())) {
      return traits_type::not_eof(ch);
    }
    char value = traits_type::to_char_type(ch);
    return xsputn(&value, 1) == 1 ? ch : traits_type::eof();
  }

  int sync() final { return destination_->pubsync(); }

  pos_type seekoff(off_type offset, std::ios_base::seekdir direction,
                   std::ios_base::openmode mode) final {
    // Support the current-position query used by tellp(), preserving the destination's absolute
    // escaped offset. ExprStringDoc rendering never needs to reposition the output sequence.
    if (offset == 0 && direction == std::ios_base::cur && (mode & std::ios_base::out)) {
      return destination_->pubseekoff(offset, direction, mode);
    }
    return pos_type(off_type(-1));
  }

 private:
  std::ostream* output_;
  std::streambuf* destination_;
  bool saw_newline_{false};
};

ffi::String RenderInvisiblePathInfo(const ffi::String& script,
                                    const ffi::Array<AccessPath>& requested_paths,
                                    const ffi::Array<ffi::Optional<AccessPath>>& visible_paths) {
  if (requested_paths.empty()) return script;

  std::ostringstream os;
  for (size_t i = 0; i < requested_paths.size(); ++i) {
    if (i != 0) os << "\n";
    const AccessPath& requested_path = requested_paths[i];
    os << "Access path: " << requested_path;

    ffi::Optional<AccessPath> visible_path = std::nullopt;
    if (i < visible_paths.size()) visible_path = visible_paths[i];
    if (!visible_path.has_value()) {
      os << "\nNote: No visible object for this path is rendered in TVMScript.";
    } else if (!visible_path.value()->PathEqual(requested_path) &&
               visible_path.value()->IsPrefixOf(requested_path)) {
      os << "\nNote: The underlined object is the nearest visible parent of this path.";
    }
  }
  os << "\n\n" << script;
  return ffi::String(os.str());
}

}  // namespace

/*!
 * \brief Operator precedence
 *
 * This is based on
 * https://docs.python.org/3/reference/expressions.html#operator-precedence
 */
enum class ExprPrecedence : int32_t {
  /*! \brief Unknown precedence */
  kUnkown = 0,
  /*! \brief Lambda Expression */
  kLambda = 1,
  /*! \brief Conditional Expression */
  kIfThenElse = 2,
  /*! \brief Boolean OR */
  kBooleanOr = 3,
  /*! \brief Boolean AND */
  kBooleanAnd = 4,
  /*! \brief Boolean NOT */
  kBooleanNot = 5,
  /*! \brief Comparisons */
  kComparison = 6,
  /*! \brief Bitwise OR */
  kBitwiseOr = 7,
  /*! \brief Bitwise XOR */
  kBitwiseXor = 8,
  /*! \brief Bitwise AND */
  kBitwiseAnd = 9,
  /*! \brief Shift Operators */
  kShift = 10,
  /*! \brief Addition and subtraction */
  kAdd = 11,
  /*! \brief Multiplication, division, floor division, remainder */
  kMult = 12,
  /*! \brief Positive negative and bitwise NOT */
  kUnary = 13,
  /*! \brief Exponentiation */
  kExp = 14,
  /*! \brief Index access, attribute access, call and atom expression */
  kIdentity = 15,
};

ExprPrecedence GetExprPrecedence(const ExprDoc& doc) {
  // Key is the value of OperationDocNode::Kind
  static const std::vector<ExprPrecedence> op_kind_precedence = []() {
    using OpKind = OperationDocNode::Kind;
    std::map<OpKind, ExprPrecedence> raw_table = {
        {OpKind::kUSub, ExprPrecedence::kUnary},
        {OpKind::kInvert, ExprPrecedence::kUnary},
        {OpKind::kNot, ExprPrecedence::kBooleanNot},
        {OpKind::kAdd, ExprPrecedence::kAdd},
        {OpKind::kSub, ExprPrecedence::kAdd},
        {OpKind::kMult, ExprPrecedence::kMult},
        {OpKind::kDiv, ExprPrecedence::kMult},
        {OpKind::kFloorDiv, ExprPrecedence::kMult},
        {OpKind::kMod, ExprPrecedence::kMult},
        {OpKind::kPow, ExprPrecedence::kExp},
        {OpKind::kLShift, ExprPrecedence::kShift},
        {OpKind::kRShift, ExprPrecedence::kShift},
        {OpKind::kBitAnd, ExprPrecedence::kBitwiseAnd},
        {OpKind::kBitOr, ExprPrecedence::kBitwiseOr},
        {OpKind::kBitXor, ExprPrecedence::kBitwiseXor},
        {OpKind::kLt, ExprPrecedence::kComparison},
        {OpKind::kLtE, ExprPrecedence::kComparison},
        {OpKind::kEq, ExprPrecedence::kComparison},
        {OpKind::kNotEq, ExprPrecedence::kComparison},
        {OpKind::kGt, ExprPrecedence::kComparison},
        {OpKind::kGtE, ExprPrecedence::kComparison},
        {OpKind::kAnd, ExprPrecedence::kBooleanAnd},
        {OpKind::kOr, ExprPrecedence::kBooleanOr},
        {OpKind::kMatMul, ExprPrecedence::kMult},
        {OpKind::kIfThenElse, ExprPrecedence::kIfThenElse},
    };
    int n = static_cast<int>(OpKind::kSpecialEnd);
    std::vector<ExprPrecedence> table(n + 1, ExprPrecedence::kUnkown);
    for (const auto& kv : raw_table) {
      table[static_cast<int>(kv.first)] = kv.second;
    }
    return table;
  }();

  // Key is the type index of Doc
  static const std::unordered_map<uint32_t, ExprPrecedence> doc_type_precedence = {
      {LiteralDocNode::RuntimeTypeIndex(), ExprPrecedence::kIdentity},
      {ExprStringDocNode::RuntimeTypeIndex(), ExprPrecedence::kIdentity},
      {IdDocNode::RuntimeTypeIndex(), ExprPrecedence::kIdentity},
      {AttrAccessDocNode::RuntimeTypeIndex(), ExprPrecedence::kIdentity},
      {IndexDocNode::RuntimeTypeIndex(), ExprPrecedence::kIdentity},
      {CallDocNode::RuntimeTypeIndex(), ExprPrecedence::kIdentity},
      {LambdaDocNode::RuntimeTypeIndex(), ExprPrecedence::kLambda},
      {TupleDocNode::RuntimeTypeIndex(), ExprPrecedence::kIdentity},
      {ListDocNode::RuntimeTypeIndex(), ExprPrecedence::kIdentity},
      {DictDocNode::RuntimeTypeIndex(), ExprPrecedence::kIdentity},
  };

  if (const auto* op_doc = doc.as<OperationDocNode>()) {
    size_t kind = static_cast<int>(op_doc->kind);
    TVM_FFI_CHECK_LT(kind, op_kind_precedence.size(), ValueError) << "Invalid operation: " << kind;
    ExprPrecedence precedence = op_kind_precedence[kind];
    TVM_FFI_ICHECK(precedence != ExprPrecedence::kUnkown)
        << "Precedence for operator " << static_cast<int>(op_doc->kind) << " is unknown";
    return precedence;
  }
  auto it = doc_type_precedence.find(doc->type_index());
  if (it != doc_type_precedence.end()) {
    return it->second;
  }
  TVM_FFI_ICHECK(false) << "Precedence for doc type " << doc->GetTypeKey() << " is unknown";
  throw;
}

class PythonDocPrinter : public DocPrinter {
 public:
  explicit PythonDocPrinter(const PrinterConfig& options) : DocPrinter(options) {}

 protected:
  using DocPrinter::PrintDoc;

  void PrintTypedDoc(const LiteralDoc& doc) final;
  void PrintTypedDoc(const ExprStringDoc& doc) final;
  void PrintTypedDoc(const IdDoc& doc) final;
  void PrintTypedDoc(const AttrAccessDoc& doc) final;
  void PrintTypedDoc(const IndexDoc& doc) final;
  void PrintTypedDoc(const OperationDoc& doc) final;
  void PrintTypedDoc(const CallDoc& doc) final;
  void PrintTypedDoc(const LambdaDoc& doc) final;
  void PrintTypedDoc(const ListDoc& doc) final;
  void PrintTypedDoc(const DictDoc& doc) final;
  void PrintTypedDoc(const TupleDoc& doc) final;
  void PrintTypedDoc(const SliceDoc& doc) final;
  void PrintTypedDoc(const StmtBlockDoc& doc) final;
  void PrintTypedDoc(const AssignDoc& doc) final;
  void PrintTypedDoc(const IfDoc& doc) final;
  void PrintTypedDoc(const WhileDoc& doc) final;
  void PrintTypedDoc(const BreakDoc& doc) final;
  void PrintTypedDoc(const ContinueDoc& doc) final;
  void PrintTypedDoc(const ForDoc& doc) final;
  void PrintTypedDoc(const ExprStmtDoc& doc) final;
  void PrintTypedDoc(const AssertDoc& doc) final;
  void PrintTypedDoc(const ReturnDoc& doc) final;
  void PrintTypedDoc(const ScopeDoc& doc) final;
  void PrintTypedDoc(const FunctionDoc& doc) final;
  void PrintTypedDoc(const ClassDoc& doc) final;
  void PrintTypedDoc(const CommentDoc& doc) final;
  void PrintTypedDoc(const DocStringDoc& doc) final;
  void PrintTypedDoc(const OpCallDoc& doc) final;

 private:
  void NewLineWithoutIndent() {
    size_t start_pos = output_.tellp();
    output_ << "\n";
    size_t end_pos = output_.tellp();
    underlines_exempted_.push_back({start_pos, end_pos});
  }

  template <typename DocType>
  void PrintJoinedDocs(const ffi::Array<DocType>& docs, const std::string& separator) {
    bool is_first = true;
    for (auto& doc : docs) {
      if (is_first) {
        is_first = false;
      } else {
        output_ << separator;
      }
      PrintDoc(doc);
    }
  }

  void PrintIndentedBlock(const ffi::Array<StmtDoc>& docs) {
    IncreaseIndent();
    for (const StmtDoc& d : docs) {
      NewLine();
      PrintDoc(d);
    }
    if (docs.empty()) {
      NewLine();
      output_ << "pass";
    }
    DecreaseIndent();
  }

  void PrintDecorators(const ffi::Array<ExprDoc>& decorators) {
    for (const ExprDoc& decorator : decorators) {
      output_ << "@";
      PrintDoc(decorator);
      NewLine();
    }
  }

  /*!
   * \brief Print expression and add parenthesis if needed.
   */
  void PrintChildExpr(const ExprDoc& doc, ExprPrecedence parent_precedence,
                      bool parenthesis_for_same_precedence = false) {
    ExprPrecedence doc_precedence = GetExprPrecedence(doc);
    if (doc_precedence < parent_precedence ||
        (parenthesis_for_same_precedence && doc_precedence == parent_precedence)) {
      output_ << "(";
      PrintDoc(doc);
      output_ << ")";
    } else {
      PrintDoc(doc);
    }
  }

  /*!
   * \brief Print expression and add parenthesis if doc has lower precedence than parent.
   */
  void PrintChildExpr(const ExprDoc& doc, const ExprDoc& parent,
                      bool parenthesis_for_same_precedence = false) {
    ExprPrecedence parent_precedence = GetExprPrecedence(parent);
    return PrintChildExpr(doc, parent_precedence, parenthesis_for_same_precedence);
  }

  /*!
   * \brief Print expression and add parenthesis if doc doesn't have higher precedence than parent.
   *
   * This function should be used to print an child expression that needs to be wrapped
   * by parenthesis even if it has the same precedence as its parent, e.g., the `b` in `a + b`
   * and the `b` and `c` in `a if b else c`.
   */
  void PrintChildExprConservatively(const ExprDoc& doc, const ExprDoc& parent) {
    PrintChildExpr(doc, parent, /*parenthesis_for_same_precedence=*/true);
  }

  void MaybePrintCommentInline(const StmtDoc& stmt) {
    if (stmt->comment.has_value()) {
      const std::string& comment = stmt->comment.value();
      bool has_newline = std::find(comment.begin(), comment.end(), '\n') != comment.end();
      TVM_FFI_CHECK(!has_newline, ValueError)
          << "the comment string of " << stmt->GetTypeKey() << " cannot have newline.";
      size_t start_pos = output_.tellp();
      output_ << "  # " << comment;
      size_t end_pos = output_.tellp();
      underlines_exempted_.push_back({start_pos, end_pos});
    }
  }

  void MaybePrintCommenMultiLines(const StmtDoc& stmt, bool new_line = false) {
    if (stmt->comment.has_value()) {
      std::vector<std::string> comment_lines = support::Split(stmt->comment.value(), '\n');
      bool first_line = true;
      size_t start_pos = output_.tellp();
      for (const std::string& line : comment_lines) {
        if (first_line) {
          output_ << "# " << line;
          first_line = false;
        } else {
          NewLine() << "# " << line;
        }
      }
      size_t end_pos = output_.tellp();
      underlines_exempted_.push_back({start_pos, end_pos});
      if (new_line) {
        NewLine();
      }
    }
  }

  void PrintDocString(const ffi::String& comment) {
    size_t start_pos = output_.tellp();
    output_ << "\"\"\"";

    std::vector<std::string> comment_lines = support::Split(comment, '\n');
    for (const std::string& line : comment_lines) {
      if (line.empty()) {
        // No indentation on empty line
        output_ << "\n";
      } else {
        NewLine() << line;
      }
    }

    NewLine() << "\"\"\"";
    size_t end_pos = output_.tellp();
    underlines_exempted_.push_back({start_pos, end_pos});
  }

  void PrintBlockComment(const ffi::String& comment) {
    IncreaseIndent();
    NewLine();
    PrintDocString(comment);
    DecreaseIndent();
  }
};

void PythonDocPrinter::PrintTypedDoc(const LiteralDoc& doc) {
  const ffi::Any& value = doc->value;
  if (value == nullptr) {
    output_ << "None";
  } else if (const auto* int_imm = value.as<IntImmNode>()) {
    PrimType int_ty = int_imm->ty.as_or_throw<PrimType>();
    if (int_ty.MatchesCode(DLDataTypeCode::kDLBool)) {
      output_ << (int_imm->value ? "True" : "False");
    } else {
      output_ << int_imm->value;
    }
  } else if (const auto* float_imm = value.as<FloatImmNode>()) {
    // TODO(yelite): Make float number printing roundtrippable
    if (std::isinf(float_imm->value) || std::isnan(float_imm->value)) {
      output_ << '"' << float_imm->value << '"';
    } else if (std::nearbyint(float_imm->value) == float_imm->value) {
      // Special case for floating-point values which would be
      // formatted using %g, are not displayed in scientific
      // notation, and whose fractional part is zero.
      //
      // By default, using `operator<<(std::ostream&, double)`
      // delegates to the %g printf formatter.  This strips off any
      // trailing zeros, and also strips the decimal point if no
      // trailing zeros are found.  When parsed in python, due to the
      // missing decimal point, this would incorrectly convert a float
      // to an integer.  Providing the `std::showpoint` modifier
      // instead delegates to the %#g printf formatter.  On its own,
      // this resolves the round-trip errors, but also prevents the
      // trailing zeros from being stripped off.
      std::showpoint(output_);
      std::fixed(output_);
      output_.precision(1);
      output_ << float_imm->value;
    } else {
      std::defaultfloat(output_);
      std::noshowpoint(output_);
      output_.precision(17);
      output_ << float_imm->value;
    }

  } else if (const auto opt_str = value.as<ffi::String>()) {
    output_ << "\"" << support::StrEscape((*opt_str).data(), (*opt_str).size()) << "\"";
  } else {
    TVM_FFI_THROW(TypeError) << "Unsupported literal value type: " << value.GetTypeKey();
  }
}

void PythonDocPrinter::PrintTypedDoc(const ExprStringDoc& doc) {
  this->output_ << '"';
  {
    ScopedExprStringEscapeBuf escaping_scope(&this->output_);
    this->PrintDoc(doc->value);
    TVM_FFI_ICHECK(!escaping_scope.saw_newline())
        << "An expression rendered inside a Python string literal must be one line";
  }
  this->output_ << '"';
}

void PythonDocPrinter::PrintTypedDoc(const IdDoc& doc) { output_ << doc->name; }

void PythonDocPrinter::PrintTypedDoc(const AttrAccessDoc& doc) {
  PrintChildExpr(doc->value, doc);
  output_ << "." << doc->name;
}

void PythonDocPrinter::PrintTypedDoc(const IndexDoc& doc) {
  PrintChildExpr(doc->value, doc);
  if (doc->indices.size() == 0) {
    output_ << "[()]";
  } else {
    output_ << "[";
    PrintJoinedDocs(doc->indices, ", ");
    output_ << "]";
  }
}

const std::string OperatorToString(OperationDocNode::Kind operation_kind) {
  static const std::vector<std::string> op_kind2str = []() {
    using OpKind = OperationDocNode::Kind;
    std::map<OpKind, std::string> raw_table = {
        {OpKind::kUSub, "-"},       //
        {OpKind::kInvert, "~"},     //
        {OpKind::kNot, "not "},     //
        {OpKind::kAdd, "+"},        //
        {OpKind::kSub, "-"},        //
        {OpKind::kMult, "*"},       //
        {OpKind::kDiv, "/"},        //
        {OpKind::kFloorDiv, "//"},  //
        {OpKind::kMod, "%"},        //
        {OpKind::kPow, "**"},       //
        {OpKind::kLShift, "<<"},    //
        {OpKind::kRShift, ">>"},    //
        {OpKind::kBitAnd, "&"},     //
        {OpKind::kBitOr, "|"},      //
        {OpKind::kBitXor, "^"},     //
        {OpKind::kLt, "<"},         //
        {OpKind::kLtE, "<="},       //
        {OpKind::kEq, "=="},        //
        {OpKind::kNotEq, "!="},     //
        {OpKind::kGt, ">"},         //
        {OpKind::kGtE, ">="},       //
        {OpKind::kAnd, "and"},      //
        {OpKind::kOr, "or"},        //
        {OpKind::kMatMul, "@"},     //
    };

    std::vector<std::string> table;
    table.resize(static_cast<int>(OperationDocNode::Kind::kSpecialEnd) + 1);

    for (const auto& kv : raw_table) {
      table[static_cast<int>(kv.first)] = kv.second;
    }

    return table;
  }();

  auto op_index = static_cast<int>(operation_kind);
  TVM_FFI_ICHECK_LT(op_index, op_kind2str.size());
  const std::string str = op_kind2str[op_index];
  TVM_FFI_ICHECK(!str.empty()) << "OperationDocNode::Kind " << static_cast<int>(operation_kind)
                               << " cannot be converted to operator token in Python directly.";
  return str;
}

void PythonDocPrinter::PrintTypedDoc(const OperationDoc& doc) {
  using OpKind = OperationDocNode::Kind;
  if (doc->kind < OpKind::kUnaryEnd) {
    // Unary Operators
    TVM_FFI_ICHECK_EQ(doc->operands.size(), 1);
    output_ << OperatorToString(doc->kind);
    PrintChildExpr(doc->operands[0], doc);
  } else if (doc->kind == OpKind::kPow) {
    // Power operator is different than other binary operators
    // It's right-associative and binds less tightly than unary operator on its right.
    // https://docs.python.org/3/reference/expressions.html#the-power-operator
    // https://docs.python.org/3/reference/expressions.html#operator-precedence
    TVM_FFI_ICHECK_EQ(doc->operands.size(), 2);
    PrintChildExprConservatively(doc->operands[0], doc);
    output_ << " ** ";
    PrintChildExpr(doc->operands[1], ExprPrecedence::kUnary);
  } else if (doc->kind < OpKind::kBinaryEnd) {
    // Binary Operator
    TVM_FFI_ICHECK_EQ(doc->operands.size(), 2);
    PrintChildExpr(doc->operands[0], doc);
    output_ << " " << OperatorToString(doc->kind) << " ";
    PrintChildExprConservatively(doc->operands[1], doc);
  } else if (doc->kind == OpKind::kIfThenElse) {
    TVM_FFI_CHECK_EQ(doc->operands.size(), 3, ValueError)
        << "IfThenElse requires 3 operands, but got " << doc->operands.size();
    PrintChildExpr(doc->operands[1], doc);
    output_ << " if ";
    PrintChildExprConservatively(doc->operands[0], doc);
    output_ << " else ";
    PrintChildExprConservatively(doc->operands[2], doc);
  } else {
    TVM_FFI_THROW(InternalError) << "Unknown OperationDocNode::Kind "
                                 << static_cast<int>(doc->kind);
    throw;
  }
}

void PythonDocPrinter::PrintTypedDoc(const CallDoc& doc) {
  PrintChildExpr(doc->callee, doc);

  output_ << "(";

  // Print positional args
  bool is_first = true;
  for (const ExprDoc& arg : doc->args) {
    if (is_first) {
      is_first = false;
    } else {
      output_ << ", ";
    }
    PrintDoc(arg);
  }

  // Print keyword args
  TVM_FFI_ICHECK_EQ(doc->kwargs_keys.size(), doc->kwargs_values.size())
      << "CallDoc should have equal number of elements in kwargs_keys and kwargs_values.";
  for (size_t i = 0; i < doc->kwargs_keys.size(); i++) {
    if (is_first) {
      is_first = false;
    } else {
      output_ << ", ";
    }
    const ffi::String& keyword = doc->kwargs_keys[i];
    output_ << keyword;
    output_ << "=";
    PrintDoc(doc->kwargs_values[i]);
  }

  output_ << ")";
}

void PythonDocPrinter::PrintTypedDoc(const LambdaDoc& doc) {
  output_ << "lambda ";
  PrintJoinedDocs(doc->args, ", ");
  output_ << ": ";
  PrintChildExpr(doc->body, doc);
}

void PythonDocPrinter::PrintTypedDoc(const ListDoc& doc) {
  output_ << "[";
  PrintJoinedDocs(doc->elements, ", ");
  output_ << "]";
}

void PythonDocPrinter::PrintTypedDoc(const TupleDoc& doc) {
  output_ << "(";
  if (doc->elements.size() == 1) {
    PrintDoc(doc->elements[0]);
    output_ << ",";
  } else {
    PrintJoinedDocs(doc->elements, ", ");
  }
  output_ << ")";
}

void PythonDocPrinter::PrintTypedDoc(const DictDoc& doc) {
  TVM_FFI_ICHECK_EQ(doc->keys.size(), doc->values.size())
      << "DictDoc should have equal number of elements in keys and values.";
  output_ << "{";
  size_t idx = 0;
  for (const ExprDoc& key : doc->keys) {
    if (idx > 0) {
      output_ << ", ";
    }
    PrintDoc(key);
    output_ << ": ";
    PrintDoc(doc->values[idx]);
    idx++;
  }
  output_ << "}";
}

void PythonDocPrinter::PrintTypedDoc(const SliceDoc& doc) {
  if (doc->start != nullptr) {
    PrintDoc(doc->start.value());
  }
  output_ << ":";
  if (doc->stop != nullptr) {
    PrintDoc(doc->stop.value());
  }
  if (doc->step != nullptr) {
    output_ << ":";
    PrintDoc(doc->step.value());
  }
}

void PythonDocPrinter::PrintTypedDoc(const StmtBlockDoc& doc) {
  for (const StmtDoc& stmt : doc->stmts) {
    PrintDoc(stmt);
    if (stmt != doc->stmts.back()) {
      NewLine();
    }
  }
}

void PythonDocPrinter::PrintTypedDoc(const AssignDoc& doc) {
  if (const auto* tuple_doc = doc->lhs.as<TupleDocNode>()) {
    PrintJoinedDocs(tuple_doc->elements, ", ");
  } else {
    PrintDoc(doc->lhs);
  }

  if (doc->annotation) {
    output_ << ": ";
    PrintDoc(doc->annotation.value());
  }
  if (doc->rhs) {
    output_ << " = ";
    if (const auto* tuple_doc = doc->rhs.as<TupleDocNode>()) {
      if (tuple_doc->elements.size() > 1) {
        PrintJoinedDocs(tuple_doc->elements, ", ");
      } else {
        PrintDoc(doc->rhs.value());
      }
    } else {
      PrintDoc(doc->rhs.value());
    }
  }
  MaybePrintCommentInline(doc);
}

void PythonDocPrinter::PrintTypedDoc(const IfDoc& doc) {
  MaybePrintCommenMultiLines(doc, true);
  output_ << "if ";
  PrintDoc(doc->predicate);
  output_ << ":";

  PrintIndentedBlock(doc->then_branch);

  if (!doc->else_branch.empty()) {
    NewLine();
    output_ << "else:";
    PrintIndentedBlock(doc->else_branch);
  }
}

void PythonDocPrinter::PrintTypedDoc(const WhileDoc& doc) {
  MaybePrintCommenMultiLines(doc, true);
  output_ << "while ";
  PrintDoc(doc->predicate);
  output_ << ":";

  PrintIndentedBlock(doc->body);
}

void PythonDocPrinter::PrintTypedDoc(const BreakDoc& doc) { output_ << "break"; }

void PythonDocPrinter::PrintTypedDoc(const ContinueDoc& doc) { output_ << "continue"; }

void PythonDocPrinter::PrintTypedDoc(const ForDoc& doc) {
  MaybePrintCommenMultiLines(doc, true);
  output_ << "for ";
  if (const auto* tuple = doc->lhs.as<TupleDocNode>()) {
    if (tuple->elements.size() == 1) {
      PrintDoc(tuple->elements[0]);
      output_ << ",";
    } else {
      PrintJoinedDocs(tuple->elements, ", ");
    }
  } else {
    PrintDoc(doc->lhs);
  }
  output_ << " in ";
  PrintDoc(doc->rhs);
  output_ << ":";

  PrintIndentedBlock(doc->body);
}

void PythonDocPrinter::PrintTypedDoc(const ScopeDoc& doc) {
  MaybePrintCommenMultiLines(doc, true);
  output_ << "with ";
  PrintDoc(doc->rhs);
  if (doc->lhs != nullptr) {
    output_ << " as ";
    PrintDoc(doc->lhs.value());
  }
  output_ << ":";

  PrintIndentedBlock(doc->body);
}

void PythonDocPrinter::PrintTypedDoc(const ExprStmtDoc& doc) {
  PrintDoc(doc->expr);
  MaybePrintCommentInline(doc);
}

void PythonDocPrinter::PrintTypedDoc(const AssertDoc& doc) {
  output_ << "assert ";
  PrintDoc(doc->test);
  if (doc->msg.has_value()) {
    output_ << ", ";
    PrintDoc(doc->msg.value());
  }
  MaybePrintCommentInline(doc);
}

void PythonDocPrinter::PrintTypedDoc(const ReturnDoc& doc) {
  output_ << "return ";
  PrintDoc(doc->value);
  MaybePrintCommentInline(doc);
}

void PythonDocPrinter::PrintTypedDoc(const FunctionDoc& doc) {
  for (const AssignDoc& arg_doc : doc->args) {
    TVM_FFI_ICHECK(!arg_doc->comment.has_value())
        << "Function arg cannot have comment attached to them.";
  }

  PrintDecorators(doc->decorators);

  output_ << "def ";
  PrintDoc(doc->name);

  output_ << "(";
  PrintJoinedDocs(doc->args, ", ");
  output_ << ")";

  if (doc->return_type.has_value()) {
    output_ << " -> ";
    PrintDoc(doc->return_type.value());
  }

  output_ << ":";

  if (doc->comment.has_value()) {
    PrintBlockComment(doc->comment.value());
  }
  PrintIndentedBlock(doc->body);
  NewLineWithoutIndent();
}

void PythonDocPrinter::PrintTypedDoc(const ClassDoc& doc) {
  PrintDecorators(doc->decorators);

  output_ << "class ";
  PrintDoc(doc->name);
  output_ << ":";

  if (doc->comment.has_value()) {
    PrintBlockComment(doc->comment.value());
  }
  PrintIndentedBlock(doc->body);
}

void PythonDocPrinter::PrintTypedDoc(const CommentDoc& doc) {
  if (doc->comment.has_value()) {
    MaybePrintCommenMultiLines(doc, false);
  }
}

void PythonDocPrinter::PrintTypedDoc(const DocStringDoc& doc) {
  if (doc->comment.has_value() && !doc->comment.value().empty()) {
    PrintDocString(doc->comment.value());
  }
}

void PythonDocPrinter::PrintTypedDoc(const OpCallDoc& doc) {
  PrintDoc(doc->callee);

  output_ << "(";

  // Print positional args
  bool wrote_any = false;
  for (const Doc& arg : doc->args) {
    if (wrote_any) {
      output_ << ", ";
    }
    wrote_any = true;
    PrintDoc(arg);
  }
  // workspace first (if present and non-empty)
  if (doc->workspace.has_value() && !doc->workspace.value()->keys.empty()) {
    if (wrote_any) output_ << ", ";
    wrote_any = true;
    output_ << "workspace=";
    PrintDoc(doc->workspace.value());
  }
  // dispatch next (if present)
  if (doc->dispatch.has_value()) {
    if (wrote_any) output_ << ", ";
    wrote_any = true;
    output_ << "dispatch=";
    PrintDoc(doc->dispatch.value());
  }
  // Flatten config as keyword args: key=value
  if (doc->config.has_value() && !doc->config.value()->keys.empty()) {
    const auto* dict = doc->config.value().as<DictDocNode>();
    // Only flatten if all keys are literal strings; otherwise, fallback to config={...}
    bool all_str_keys = true;
    for (const ExprDoc& k : dict->keys) {
      if (!k.as<LiteralDocNode>()) {
        all_str_keys = false;
        break;
      }
      const auto* lit = k.as<LiteralDocNode>();
      if (!lit->value.as<ffi::String>()) {
        all_str_keys = false;
        break;
      }
    }
    if (all_str_keys) {
      int n = dict->keys.size();
      for (int i = 0; i < n; ++i) {
        const auto* lit = dict->keys[i].as<LiteralDocNode>();
        std::string key = lit->value.as_or_throw<ffi::String>();
        if (wrote_any) output_ << ", ";
        wrote_any = true;
        output_ << key << "=";
        PrintDoc(dict->values[i]);
      }
    } else {
      if (wrote_any) output_ << ", ";
      wrote_any = true;
      output_ << "config=";
      PrintDoc(doc->config.value());
    }
  }
  output_ << ")";
}

ffi::String DocToPythonScript(Doc doc, const PrinterConfig& cfg) {
  if (cfg->num_context_lines < 0) {
    cfg->num_context_lines = std::numeric_limits<int32_t>::max();
  }
  PythonDocPrinter printer(cfg);
  printer.Append(doc, cfg);
  std::string script = printer.GetString();

  // GetString terminates non-empty output with one newline.  Preserve the
  // established DocToPythonScript result without normalizing any other
  // trailing whitespace.
  if (!script.empty()) {
    TVM_FFI_ICHECK_EQ(script.back(), '\n');
    script.pop_back();
  }

  if (!cfg->render_invisible_path_info) return script;
  return RenderInvisiblePathInfo(script, cfg->path_to_underline, printer.GetVisiblePaths());
}

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef().def("script.printer.DocToPythonScript", DocToPythonScript);
}

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