/*!
 * \file target/codegen_cutedsl.h
 * \brief Utility to generate CuTeDSL code
 */
#ifndef TVM_TL_TARGET_CODEGEN_CUTEDSL_H_
#define TVM_TL_TARGET_CODEGEN_CUTEDSL_H_

#include "support/check.h"
#include <tvm/target/codegen.h>
#include <tvm/tirx/expr.h>
#include <tvm/tirx/op.h>

#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>

#include "codegen_py.h"

namespace tvm {
namespace codegen {

class CodeGenTileLangCuTeDSL final : public CodeGenTileLangPY {
public:
  CodeGenTileLangCuTeDSL();

protected:
  void InitFuncState_(const PrimFunc &f) override;

  void PrintFuncDecorator_(std::ostream &os) override; // NOLINT(*)
  void PreFunctionBody_(const PrimFunc &f) override;

protected:
  void PrintType(DataType t, std::ostream &os) override; // NOLINT(*)

  void VisitExpr_(const BroadcastNode *op,
                  std::ostream &os) override; // NOLINT(*)
  void VisitExpr_(const FloatImmNode *op,
                  std::ostream &os) override;                       // NOLINT(*)
  void VisitExpr_(const IntImmNode *op, std::ostream &os) override; // NOLINT(*)
  void VisitExpr_(const CastNode *op, std::ostream &os) override;   // NOLINT(*)
  void VisitExpr_(const DivNode *op, std::ostream &os) override;    // NOLINT(*)
  void VisitExpr_(const MinNode *op, std::ostream &os) override;    // NOLINT(*)
  void VisitExpr_(const MaxNode *op, std::ostream &os) override;    // NOLINT(*)
  void VisitExpr_(const CallNode *op, std::ostream &os) override;   // NOLINT(*)
  void VisitExpr_(const SelectNode *op, std::ostream &os) override; // NOLINT(*)
  void VisitExpr_(const LetNode *op, std::ostream &os) override;    // NOLINT(*)
  void VisitExpr_(const ShuffleNode *op,
                  std::ostream &os) override; // NOLINT(*)
  void VisitExpr_(const BufferLoadNode *op,
                  std::ostream &os) override; // NOLINT(*)

  void VisitStmt_(const BufferStoreNode *op) override;
  void VisitStmt_(const BindNode *op) override;
  void VisitStmt_(const AllocBufferNode *op) override;
  void VisitStmt_(const AttrStmtNode *op) override;
  void VisitStmt_(const ForNode *op) override;
  void VisitStmt_(const IfThenElseNode *op) override;
  void VisitStmt_(const EvaluateNode *op) override;
  void VisitStmt_(const SeqStmtNode *op) override;

protected:
  virtual void PrintVecElemLoad_(const std::string &vec, DataType t, int i,
                                 std::ostream &os); // NOLINT(*)
  virtual void PrintVecElemStore_(const std::string &vec, DataType t, int i,
                                  const std::string &value);
  virtual void PrintVecStore_(const BufferNode *buffer, DataType t,
                              PrimExpr base, const std::string &value);
  void PrintVecBinaryOp_(const std::string &opstr, DataType dtype, PrimExpr lhs,
                         PrimExpr rhs,
                         std::ostream &os); // NOLINT(*)
  void PrintBinaryExpr_(const std::string &opstr, DataType dtype, PrimExpr lhs,
                        PrimExpr rhs,
                        std::ostream &os) override; // NOLINT(*)
  void PrintBinaryIntrinsic_(const CallNode *op, const char *opstr,
                             std::ostream &os) override; // NOLINT(*)

  void PrintCallExtern_(Type ret_type, ffi::String global_symbol,
                        const ffi::Array<PrimExpr> &args, bool skip_first_arg,
                        std::ostream &os) override; // NOLINT(*)

  std::string GetBufferPtr_(const BufferNode *buffer, PrimExpr index);
  std::string GetBufferRef_(DataType t, const BufferNode *buffer,
                            PrimExpr index) override;
  PrimExpr LinearizeBufferIndices_(const BufferNode *buffer,
                                   const ffi::Array<PrimExpr> &indices);

  // Get pointer string from a Var expression (local buffer -> vid.iterator)
  std::string GetVarPtr_(const PrimExpr &expr);

  /*!
   * \brief Print expr representing the thread tag
   * \param IterVar iv The thread index to be binded;
   */
  virtual void BindThreadIndex_(const IterVar &iv); // NOLINT(*)

  virtual void PrintStorageSync_(const CallNode *op);

  std::string
  CanonicalizeFastmathFunctionName_(const std::string &func_name) const;

private:
  // The name of the mbarrier array in shared memory
  const std::string mbarrier_name_ = "mbarrier";

  std::unordered_map<const VarNode *, IntImm> unroll_factor_;

  std::unordered_set<const VarNode *> raw_pointer_vars_;
  std::unordered_set<const VarNode *> zero_like_vars_;

  std::vector<std::string> eviction_policy_names_ = {
      "EVICT_NORMAL", "EVICT_FIRST", "EVICT_LAST"};

  // Fastmath configuration (read from PassContext)
  bool enable_fastmath_ = false;

  // Loop-break guard transformation state
  // When a for-loop contains loop_break(), we replace `break` with a guard
  // variable pattern since CuTeDSL doesn't support early exit (break).
  bool in_break_loop_ = false;
  int loop_break_counter_ = 0;
  int current_break_id_ = -1;
  // Set to true when loop_break() replacement is emitted within current SeqStmt
  bool break_emitted_in_seq_ = false;
};

} // namespace codegen
} // namespace tvm

#endif // TVM_TL_TARGET_CODEGEN_CUTEDSL_H_
