#include "flang/Optimizer/Transforms/Passes.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
#include "mlir/IR/BuiltinDialect.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/SymbolTable.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Support/LLVM.h"
#include "llvm/ADT/SmallPtrSet.h"

namespace fir {
#define GEN_PASS_DEF_OMPMARKDECLARETARGETPASS
#include "flang/Optimizer/Transforms/Passes.h.inc"
} // namespace fir

namespace {
class OMPMarkDeclareTargetPass
    : public fir::impl::OMPMarkDeclareTargetPassBase<OMPMarkDeclareTargetPass> {

  void markNestedFuncs(mlir::omp::DeclareTargetDeviceType parentDevTy,
                       mlir::omp::DeclareTargetCaptureClause parentCapClause,
                       mlir::Operation *currOp,
                       llvm::SmallPtrSet<mlir::Operation *, 16> visited) {
    if (visited.contains(currOp))
      return;
    visited.insert(currOp);

    currOp->walk([&, this](mlir::Operation *op) {
      if (auto callOp = llvm::dyn_cast<mlir::CallOpInterface>(op)) {
        if (auto symRef = llvm::dyn_cast_if_present<mlir::SymbolRefAttr>(
                callOp.getCallableForCallee())) {
          if (auto currFOp =
                  getOperation().lookupSymbol<mlir::func::FuncOp>(symRef)) {
            auto current = llvm::dyn_cast<mlir::omp::DeclareTargetInterface>(
                currFOp.getOperation());

            if (current.isDeclareTarget()) {
              auto currentDt = current.getDeclareTargetDeviceType();

              // Found the same function twice, with different device_types,
              // mark as Any as it belongs to both
              if (currentDt != parentDevTy &&
                  currentDt != mlir::omp::DeclareTargetDeviceType::any) {
                current.setDeclareTarget(
                    mlir::omp::DeclareTargetDeviceType::any,
                    current.getDeclareTargetCaptureClause());
              }
            } else {
              current.setDeclareTarget(parentDevTy, parentCapClause);
            }

            markNestedFuncs(parentDevTy, parentCapClause, currFOp, visited);
          }
        }
      }
    });
  }

  // This pass executes on mlir::ModuleOp's marking functions contained within
  // as implicitly declare target if they are called from within an explicitly
  // marked declare target function or a target region (TargetOp)
  void runOnOperation() override {
    for (auto functionOp : getOperation().getOps<mlir::func::FuncOp>()) {
      auto declareTargetOp = llvm::dyn_cast<mlir::omp::DeclareTargetInterface>(
          functionOp.getOperation());
      if (declareTargetOp.isDeclareTarget()) {
        llvm::SmallPtrSet<mlir::Operation *, 16> visited;
        markNestedFuncs(declareTargetOp.getDeclareTargetDeviceType(),
                        declareTargetOp.getDeclareTargetCaptureClause(),
                        functionOp, visited);
      }
    }

    // TODO: Extend to work with reverse-offloading, this shouldn't
    // require too much effort, just need to check the device clause
    // when it's lowering has been implemented and change the
    // DeclareTargetDeviceType argument from nohost to host depending on
    // the contents of the device clause
    getOperation()->walk([&](mlir::omp::TargetOp tarOp) {
      llvm::SmallPtrSet<mlir::Operation *, 16> visited;
      markNestedFuncs(mlir::omp::DeclareTargetDeviceType::nohost,
                      mlir::omp::DeclareTargetCaptureClause::to, tarOp,
                      visited);
    });
  }
};

} // namespace

namespace fir {
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
createOMPMarkDeclareTargetPass() {
  return std::make_unique<OMPMarkDeclareTargetPass>();
}
} // namespace fir