//===----------------------------------------------------------------------===//
//
// This file implements patterns to convert operations in the LLVM dialect to
// operations in the Tessera dialect.
//
//===----------------------------------------------------------------------===//

#include "mlir/Bytecode/BytecodeOpInterface.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/LLVMIR/LLVMTypes.h"
#include "mlir/IR/BuiltinDialect.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/IRMapping.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Interfaces/CallInterfaces.h"
#include "mlir/Interfaces/FunctionInterfaces.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/DialectConversion.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "src/enzyme_ad/jax/Dialect/Tessera/Dialect.h"
#include "src/enzyme_ad/jax/Passes/Tessera/Passes.h"

namespace mlir {
namespace enzyme {
namespace tessera {
#define GEN_PASS_DEF_LLVMTOTESSERAPASS
#include "src/enzyme_ad/jax/Passes/Tessera/Passes.h.inc"
} // namespace tessera
} // namespace enzyme
} // namespace mlir

using namespace mlir;
using namespace mlir::enzyme;
using namespace mlir::enzyme::tessera;

//===----------------------------------------------------------------------===//
// Rewrite Patterns
//===----------------------------------------------------------------------===//

namespace {

// Rewrite 'llvm.func' -> 'tessera.define'
class FuncOpRewrite final : public OpRewritePattern<LLVM::LLVMFuncOp> {
public:
  FuncOpRewrite(MLIRContext *ctx,
                const llvm::DenseMap<unsigned, mlir::Type> &byRefTypesByIndex)
      : OpRewritePattern<LLVM::LLVMFuncOp>(ctx),
        byRefTypesByIndex(byRefTypesByIndex) {}

  LogicalResult matchAndRewrite(LLVM::LLVMFuncOp funcOp,
                                PatternRewriter &rewriter) const override {

    auto module = funcOp->getParentOfType<ModuleOp>();
    auto *ctx = funcOp->getContext();

    // Only rewrite if op has tessera_op or pure_tessera_op attribute
    StringAttr tesseraOpAttr;
    bool isPure = false;

    if (auto attr = funcOp->getAttrOfType<StringAttr>("tessera_op")) {
      tesseraOpAttr = attr;
    } else if (auto attr =
                   funcOp->getAttrOfType<StringAttr>("pure_tessera_op")) {
      tesseraOpAttr = attr;
      isPure = true;
    }

    if (!tesseraOpAttr)
      return failure();

    // Parse the tessera op attribute, which is expected to be in the format:
    // "tessera_op(arg1:byref, arg2, ...):globals=index1,..." or
    // "pure_tessera_op(arg1:byref, arg2, ...):globals=index1,..." where
    // index1,... corresponds positionally, in left-to-right order, to the byref
    // args in the arg list above, and are used to look up the types of the
    // byref args.
    StringRef raw = tesseraOpAttr.getValue();

    // Parse op name (everything before the '(')
    StringRef tesseraName = raw.take_while([](char c) { return c != '('; });

    // Parse args in parentheses
    StringRef argList = raw.slice(raw.find('(') + 1, raw.find(')'));

    // Parse indices after ":globals=" (order corresponds to order byref
    // args appear in argList)
    SmallVector<StringRef> indexList;
    StringRef indicesStr = raw.substr(raw.find(')') + 1);
    if (!indicesStr.empty() && indicesStr.consume_front(":globals=")) {
      indicesStr.split(indexList, ',');
    }

    // Identify which args are marked byref and look up their types in
    // byRefTypesByIndex map
    SmallVector<Attribute> byRefTypes;
    unsigned numByRefFound = 0;
    if (!argList.trim().empty()) {
      SmallVector<StringRef> argParts;
      argList.split(argParts, ',');
      for (unsigned i = 0, e = argParts.size(); i != e; ++i) {
        StringRef arg = argParts[i].trim();

        // Check if arg is marked as byref. If not, push a unit attribute so
        // that the byRefTypes array has the same size as the number of args.
        if (!arg.contains(":byref") && !arg.contains(": byref")) {
          byRefTypes.push_back(UnitAttr::get(ctx));
          continue;
        }

        if (numByRefFound >= indexList.size()) {
          funcOp->emitError("tessera: not enough byref indices for byref args");
          return failure();
        }

        // Find the types of the byref arguments by looking for global variables
        // with names that match the pattern "__tessera_byref_arg_type_<idx>"
        // where <idx> is a number parsed in the tessera_op attribute after
        // "globals=".
        auto byRefIndex = indexList[numByRefFound++];
        unsigned idx;
        if (byRefIndex.trim().getAsInteger(10, idx)) {
          funcOp->emitError("tessera: invalid byref type index: ")
              << byRefIndex;
          return failure();
        }
        auto it = byRefTypesByIndex.find(idx);
        if (it == byRefTypesByIndex.end()) {
          funcOp->emitError("tessera: no byref type found for index: ") << idx;
          return failure();
        }
        byRefTypes.push_back(TypeAttr::get(it->second));
      }
    }

    // The format guarantees one global index per byref arg.
    // A mismatch means the annotation string is malformed or the Clang
    // plugin's emission order has drifted from this parser's assumption.
    if (numByRefFound != indexList.size()) {
      funcOp->emitError(
          "tessera: mismatch between byref arg count and global index count");
      return failure();
    }

    auto funcName = funcOp.getName();
    auto llvmFuncType = funcOp.getFunctionType();
    auto params = llvmFuncType.getParams();
    auto retType = llvmFuncType.getReturnType();

    auto fnType = FunctionType::get(
        ctx, params,
        isa<LLVM::LLVMVoidType>(retType) ? TypeRange{} : TypeRange{retType});

    // Replace current function name with tessera name defined in
    // tessera_op / pure_tessera_op attribute
    if (failed(SymbolTable::replaceAllSymbolUses(
            funcOp.getSymNameAttr(), StringAttr::get(ctx, tesseraName),
            module)))
      return failure();

    // Create the tessera.define op with the new name, function type, byRef
    // args, sizes, and purity (side effect free) attribute
    auto tesseraDefineOp = tessera::DefineOp::create(
        rewriter, funcOp.getLoc(), tesseraName.str(), fnType,
        ArrayAttr::get(ctx, byRefTypes), isPure);

    // Copy over all attributes other than the function name and type
    // and tessera_op / pure_tessera_op attribute.
    for (const auto &namedAttr : funcOp->getAttrs()) {
      if (namedAttr.getName() != SymbolTable::getSymbolAttrName() &&
          namedAttr.getName() != funcOp.getFunctionTypeAttrName() &&
          namedAttr.getName() != "tessera_op" &&
          namedAttr.getName() != "pure_tessera_op")
        tesseraDefineOp->setAttr(namedAttr.getName(), namedAttr.getValue());
    }

    // Store the original function name so we can convert back to it later
    tesseraDefineOp->setAttr("tessera.original_name",
                             rewriter.getStringAttr(funcName));

    // Clone body of function
    if (!funcOp.isExternal()) {
      rewriter.inlineRegionBefore(funcOp.getBody(), tesseraDefineOp.getBody(),
                                  tesseraDefineOp.end());
    }

    rewriter.eraseOp(funcOp);
    return success();
  }

private:
  const llvm::DenseMap<unsigned, mlir::Type> &byRefTypesByIndex;
};

// Rewrite 'llvm.call' -> 'tessera.call'
class CallOpRewrite final : public OpRewritePattern<LLVM::CallOp> {
public:
  using OpRewritePattern<LLVM::CallOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(LLVM::CallOp callOp,
                                PatternRewriter &rewriter) const override {

    auto module = callOp->getParentOfType<ModuleOp>();

    auto calleeAttr = callOp.getCalleeAttr();
    if (!calleeAttr)
      return failure();

    auto callee = SymbolTable::lookupSymbolIn(module, calleeAttr);
    // Only rewrite if callee is a tessera.define op
    auto defineOp = dyn_cast_or_null<tessera::DefineOp>(callee);
    if (!defineOp)
      return failure();

    Value sretPtr;
    Type sretType;
    auto operands = callOp.getOperands();
    auto argAttrs = callOp.getArgAttrsAttr();
    SmallVector<Attribute> newArgAttrs;
    SmallVector<NamedAttribute> newAttrs;

    // Check if first operand has sret attribute. If so, use its pointed-to
    // type as the SSA return type, since tessera.call returns values directly
    // rather than writing through a pointer.
    if (!operands.empty() && argAttrs) {
      if (auto sretAttr = defineOp.getSretAttr()) {
        sretPtr = callOp.getOperand(0);
        sretType = cast<TypeAttr>(sretAttr).getValue();
        // Build arg attrs without first element
        for (int j = 1; j < argAttrs.size(); j++)
          newArgAttrs.push_back(argAttrs[j]);
        // Filter out arg_attrs from attributes
        for (auto attr : callOp->getAttrs()) {
          if (attr.getName() != callOp.getArgAttrsAttrName())
            newAttrs.push_back(attr);
        }
        newAttrs.push_back(rewriter.getNamedAttr(
            callOp.getArgAttrsAttrName(), rewriter.getArrayAttr(newArgAttrs)));
      }
    }

    SmallVector<Value> newOperands;
    SmallVector<int32_t> loadedOperands;

    // Build operands without first element if sretPtr is present.
    // If a pointer operand has a byVal attribute or was marked as
    // byRef by the user, load the value from the pointer and store
    // that as the new operand.
    int argOffset = sretPtr ? 1 : 0;
    for (unsigned i = 0; i < operands.size() - argOffset; i++) {
      auto operand = callOp.getOperand(i + argOffset);

      if (!isa<LLVM::LLVMPointerType>(operand.getType())) {
        newOperands.push_back(operand);
        continue;
      }

      // Determine whether to load pointer and what type to load based on
      // byVal attribute or byRef attribute on the tessera.define op.
      // If neither is present, just pass the pointer through.
      Type pointeeType;
      if (auto byValAttr =
              defineOp.getArgAttr(i, LLVM::LLVMDialect::getByValAttrName())) {
        pointeeType = cast<TypeAttr>(byValAttr).getValue();
      } else if (auto byRefType = defineOp.getByRefType(i)) {
        pointeeType = byRefType;
      }

      if (pointeeType) {
        auto loadedVal = LLVM::LoadOp::create(rewriter, callOp.getLoc(),
                                              pointeeType, operand);
        newOperands.push_back(loadedVal);
        loadedOperands.push_back(i);
      } else {
        newOperands.push_back(operand);
      }
    }

    newAttrs.push_back(
        rewriter.getNamedAttr("tessera.loaded_operands",
                              rewriter.getDenseI32ArrayAttr(loadedOperands)));

    // Create tessera.call op with SSA return type
    if (sretPtr) {
      auto newCall =
          tessera::CallOp::create(rewriter, callOp.getLoc(),
                                  TypeRange{sretType}, newOperands, newAttrs);
      LLVM::StoreOp::create(rewriter, callOp.getLoc(), newCall.getResult(0),
                            sretPtr);
      rewriter.eraseOp(callOp);
    } else {
      rewriter.replaceOpWithNewOp<tessera::CallOp>(
          callOp, callOp.getResultTypes(), newOperands, callOp->getAttrs());
    }

    return success();
  }
};

// Rewrite 'llvm.return' -> 'tessera.return'
class ReturnOpRewrite final : public OpRewritePattern<LLVM::ReturnOp> {
public:
  using OpRewritePattern<LLVM::ReturnOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(LLVM::ReturnOp returnOp,
                                PatternRewriter &rewriter) const override {

    // Only rewrite if parent op is a tessera.define op
    if (!isa<tessera::DefineOp>(returnOp->getParentOp()))
      return failure();

    rewriter.replaceOpWithNewOp<tessera::ReturnOp>(returnOp,
                                                   returnOp.getOperands());
    return success();
  }
};

//===----------------------------------------------------------------------===//
// Pass to convert Func operations into Tessera operations
//===----------------------------------------------------------------------===//

struct LLVMToTesseraPass
    : public enzyme::tessera::impl::LLVMToTesseraPassBase<LLVMToTesseraPass> {
  using LLVMToTesseraPassBase::LLVMToTesseraPassBase;

  void runOnOperation() override {
    MLIRContext *ctx = &getContext();
    RewritePatternSet patterns(ctx);
    auto module = cast<ModuleOp>(getOperation());

    // Build a lookup of byref argument index -> resolved type by scanning the
    // module once for globals named "__tessera_byref_arg_type_<index>" (names
    // may be mangled with compiler-added prefixes, e.g.
    // "_ZL26__tessera_byref_arg_type_0", so we search for the marker rather
    // than requiring it at the start). Each tessera.define op's own byref index
    // list (from its "globals=" attribute) is later resolved against this map,
    // avoiding a full module rescan per op.
    llvm::DenseMap<unsigned, mlir::Type> byRefTypesByIndex;
    StringRef prefix = "__tessera_byref_arg_type_";
    for (auto global : module.getOps<mlir::LLVM::GlobalOp>()) {
      StringRef name = global.getSymName();
      size_t pos = name.find(prefix);
      if (pos == StringRef::npos)
        continue;
      StringRef idxStr = name.drop_front(pos + prefix.size());
      unsigned idx;
      if (idxStr.getAsInteger(10, idx))
        continue; // not a well-formed index suffix, skip

      auto [it, inserted] =
          byRefTypesByIndex.try_emplace(idx, global.getType());
      if (!inserted) {
        llvm::errs()
            << "Tessera: found multiple globals matching byref type index "
            << idx << ":\n";
        return signalPassFailure();
      }
    }

    patterns.add<FuncOpRewrite>(ctx, byRefTypesByIndex);
    patterns.add<CallOpRewrite, ReturnOpRewrite>(ctx);

    GreedyRewriteConfig config;
    config.setUseTopDownTraversal(true);
    if (failed(applyPatternsGreedily(getOperation(), std::move(patterns),
                                     config))) {
      llvm::errs() << "Failed to convert LLVM dialect operations to tessera "
                      "dialect operations\n";
      return signalPassFailure();
    }

    // Clean up llvm.global.annotations after conversion if it still exists
    if (auto annotations = module.lookupSymbol("llvm.global.annotations")) {
      annotations->erase();
    }
  }
};
} // namespace
