/*!
 * \file inject_assumes.cc
 * \brief Inject assumes on buffer's shape boundary check. Also convert
 * user-authored assumes to optimizer and runtime-check AttrNodes.
 */

#include "common/assume.h"
#include "common/attr.h"
#include "support/check.h"
#include "tvm/arith/analyzer.h"
#include "tvm/ir/expr.h"
#include "tvm/ir/transform.h"
#include <tvm/ffi/extra/structural_equal.h>
#include <tvm/ffi/extra/structural_hash.h>
#include <tvm/runtime/logging.h>
#include <tvm/tirx/builtin.h>
#include <tvm/tirx/expr.h>
#include <tvm/tirx/op.h>
#include <tvm/tirx/stmt.h>
#include <tvm/tirx/stmt_functor.h>
#include <tvm/tirx/transform.h>

#include <algorithm>
#include <sstream>
#include <unordered_map>
#include <vector>

namespace tvm::tl {
using namespace tirx;

class AssumeInjector : public tvm::tirx::StmtExprMutator {
  using Base = tvm::tirx::StmtExprMutator;

public:
  AssumeInjector(PrimFunc f) : f(f) {}
  static PrimFunc Substitute(PrimFunc f) {
    auto injector = AssumeInjector(f);
    f.CopyOnWrite()->body = injector(f->body);
    return f;
  }

private:
  struct AssumeCreator {
    struct Item {
      PrimExpr expr;
      std::vector<Buffer> buffers;
    };

    ffi::StructuralHash sh;
    ffi::StructuralEqual se;
    // grouped by expr, since the amount of variadic shape symbols is usually
    // much smaller than buffer
    std::vector<Item> items;
    // hash => index in items
    std::unordered_map<size_t, std::vector<size_t>> buckets;
    void addExpr(PrimExpr e, Buffer buffer) {
      size_t h = sh(e);
      auto &bucket = buckets[h];
      auto it = std::find_if(bucket.begin(), bucket.end(),
                             [&](size_t y) { return se(e, items[y].expr); });
      if (it == bucket.end()) {
        auto index = items.size();
        items.push_back({e, {buffer}});
        bucket.push_back(index);
      } else {
        items[*it].buffers.push_back(buffer);
      }
    }

    void addBuffer(Buffer buf) {
      for (auto shape : buf->shape) {
        if (shape->IsInstance<IntImmNode>())
          continue;
        addExpr(shape, buf);
      }
    }

    // --- Stride divisibility for sub-byte dtypes ---
    struct StrideDivisibilityItem {
      PrimExpr stride;
      int pack_factor;
      std::vector<Buffer> buffers;
    };
    std::vector<StrideDivisibilityItem> stride_div_items;
    std::unordered_map<size_t, std::vector<size_t>> stride_div_buckets;

    void addStrideExpr(PrimExpr stride, int pack_factor, Buffer buffer) {
      size_t h = sh(stride);
      auto &bucket = stride_div_buckets[h];
      auto it = std::find_if(bucket.begin(), bucket.end(), [&](size_t y) {
        return se(stride, stride_div_items[y].stride);
      });
      if (it == bucket.end()) {
        auto index = stride_div_items.size();
        stride_div_items.push_back({stride, pack_factor, {buffer}});
        bucket.push_back(index);
      } else {
        auto &item = stride_div_items[*it];
        item.buffers.push_back(buffer);
        // Use the largest pack_factor (strongest constraint)
        item.pack_factor = std::max(item.pack_factor, pack_factor);
      }
    }

    void addBufferStrides(Buffer buf) {
      int element_bits = buf->dtype.bits() * buf->dtype.lanes();
      if (element_bits >= 8 || buf->strides.empty())
        return;
      int pack_factor = 8 / element_bits;
      for (size_t k = 0; k + 1 < buf->strides.size(); ++k) {
        auto stride = buf->strides[k];
        if (stride->IsInstance<IntImmNode>())
          continue;
        addStrideExpr(stride, pack_factor, buf);
      }
    }

    Stmt build(Stmt body) {
      auto analyzer = arith::Analyzer{};
      for (const auto &e : items) {
        auto simplified =
            analyzer.Simplify(GT(e.expr, make_zero(e.expr->dtype)));
        std::stringstream ss;
        ss << "Buffer shape should be greater than 0: shape `" << e.expr
           << "` from buffer ";
        for (size_t i = 0; i < e.buffers.size(); i++) {
          if (i)
            ss << ", ";
          ss << "`" << e.buffers[i]->name << "`";
        }
        body = AttrStmt(simplified, tirx::attr::tilelang_assume,
                        StringImm(ss.str()), body);
      }
      // Inject stride divisibility assumes for sub-byte dtypes.
      // E.g. for fp4 (pack_factor=2), non-last-dim strides must be even.
      for (const auto &e : stride_div_items) {
        auto cond =
            EQ(floormod(e.stride, make_const(e.stride.dtype(), e.pack_factor)),
               make_zero(e.stride.dtype()));
        std::stringstream ss;
        ss << "Sub-byte buffer stride must be divisible by " << e.pack_factor
           << ": stride `" << e.stride << "` from buffer ";
        for (size_t i = 0; i < e.buffers.size(); i++) {
          if (i)
            ss << ", ";
          ss << "`" << e.buffers[i]->name << "`";
        }
        body = AttrStmt(cond, tirx::attr::tilelang_assume, StringImm(ss.str()),
                        body);
      }
      return body;
    }
  };

  Stmt VisitStmt_(const DeclBufferNode *op) final {
    AssumeCreator c;
    c.addBuffer(op->buffer);
    c.addBufferStrides(op->buffer);
    return SeqStmt({DeclBuffer(op->buffer, op->span), c.build(Evaluate(0))});
  }

  Stmt VisitStmt_(const SeqStmtNode *op) final {
    struct AssumeGroup {
      std::optional<PrimExpr> e;
      std::vector<Stmt> stmts;
    };
    std::vector<AssumeGroup> groups = {AssumeGroup{std::nullopt, {}}};
    for (size_t i = 0; i < op->seq.size(); i++) {
      auto stmt = VisitStmt(op->seq[i]);
      // Convert assume in evaluate form to assume attribute.
      // By default, we have the following IR:
      //    T.assume(cond1)
      //    Stmt1
      //    Stmt2
      //    T.assume(cond2)
      // Each user-authored assume is converted into two nested attributes:
      // the outer runtime-check marker is consumed by MakePackedAPI, while the
      // inner tilelang_assume remains available to optimization passes. Buffer
      // shape/stride assumptions created by AssumeCreator intentionally do not
      // receive the runtime-check marker because nullable buffers bind missing
      // symbolic shapes to zero.
      if (auto e = GetAssumeExprInEvaluateForm(stmt)) {
        groups.push_back(AssumeGroup{*e, {}});
      } else {
        groups.back().stmts.push_back(stmt);
      }
    }
    for (size_t i = groups.size(); i--;) {
      auto &g = groups[i];
      if (g.e) {
        Stmt body = g.stmts.size() == 1 ? g.stmts[0] : SeqStmt(g.stmts);
        std::stringstream ss;
        ss << "Assume: " << *(g.e);
        StringImm message(ss.str());
        body = AttrStmt(*g.e, tirx::attr::tilelang_assume, message,
                        std::move(body));
        body = AttrStmt(*g.e, tl::attr::kAssumeRequiresRuntimeCheck, message,
                        std::move(body));
        groups[i - 1].stmts.push_back(std::move(body));
      } else {
        ICHECK(i == 0) << "only the first group can have no assume";
      }
    }
    return groups[0].stmts.size() == 1 ? groups[0].stmts[0]
                                       : SeqStmt(groups[0].stmts);
    // return SeqStmt(groups[0].stmts);
  }

  Stmt VisitStmt_(const SBlockNode *op) final {
    auto body = VisitStmt(op->body);
    AssumeCreator c;

    // NOTE(chaofan): We only inject assumes from function arguments in the
    // root block.
    if (IsHostMainBlock(op)) {
      for (auto item : f->buffer_map) {
        c.addBuffer(item.second);
        c.addBufferStrides(item.second);
      }
    }
    for (auto item : op->alloc_buffers) {
      c.addBuffer(item);
      c.addBufferStrides(item);
    }
    for (auto item : op->match_buffers) {
      c.addBuffer(item->buffer);
      c.addBufferStrides(item->buffer);
    }

    return SBlock(op->iter_vars, op->reads, op->writes, op->name_hint,
                  c.build(body), op->init, op->alloc_buffers, op->match_buffers,
                  op->annotations, op->span);
  }

  PrimFunc f;
};

using namespace tirx::transform;

tvm::transform::Pass InjectAssumes() {
  auto pass_func = [=](PrimFunc f, IRModule m, PassContext ctx) {
    return AssumeInjector::Substitute(f);
  };
  return CreatePrimFuncPass(pass_func, 0, "tl.InjectAssumes", {});
}

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef().def("tl.transform.InjectAssumes", InjectAssumes);
}

} // namespace tvm::tl
