//
// Copyright (C) 2023-2026 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//

#include "vpux/compiler/dialect/IE/IR/dialect.hpp"
#include "vpux/compiler/dialect/IE/IR/ops_interfaces.hpp"
#include "vpux/compiler/dialect/IE/transforms/passes.hpp"
#include "vpux/compiler/utils/rewriter.hpp"
#include "vpux/compiler/utils/walk_utils.hpp"

namespace vpux::IE {
#define GEN_PASS_DECL_REMOVEVIEWLIKEOPSCHAINPASS
#define GEN_PASS_DEF_REMOVEVIEWLIKEOPSCHAINPASS
#include "vpux/compiler/dialect/IE/passes.hpp.inc"
}  // namespace vpux::IE

using namespace vpux;

namespace {

//
// ViewLikeOpsChainRewriter
//

class ViewLikeOpsChainRewriter final : public mlir::OpInterfaceRewritePattern<IE::ViewLikeOpInterface> {
public:
    ViewLikeOpsChainRewriter(mlir::MLIRContext* ctx, Logger log)
            : mlir::OpInterfaceRewritePattern<IE::ViewLikeOpInterface>(ctx), _log(log) {
    }

private:
    mlir::LogicalResult matchAndRewrite(IE::ViewLikeOpInterface origOp, mlir::PatternRewriter& rewriter) const final;

private:
    Logger _log;
};

/*
case 1:
     LayerOp                                LayerOp
        |                                      |
    ViewLikeOp1                             LayerOp
        |
       ...                =>
        |
    ViewLikeOpX
        |
    LayerOp

case 2:
     LayerOp                             LayerOp
        |                                   |
    ViewLikeOp1                         ViewLikeOpY
        |                                   |
       ...                =>               ...
        |                                LayerOp
    ViewLikeOpX
        |
    ViewLikeOpY
       ...
        |
     LayerOp
*/
mlir::LogicalResult ViewLikeOpsChainRewriter::matchAndRewrite(IE::ViewLikeOpInterface origOp,
                                                              mlir::PatternRewriter& rewriter) const {
    _log.trace("[{0}] Got '{1}' at '{2}'", getDebugName(), origOp->getName(), origOp->getLoc());

    if (!IE::isPureViewOp(origOp) || !origOp->hasOneUse()) {
        return mlir::failure();
    }

    auto inputType = mlir::cast<vpux::NDTypeInterface>(origOp->getOperand(0).getType());

    SmallVector<mlir::Operation*> viewLikeOps{origOp};
    auto* currentOp = *origOp->getUsers().begin();

    // Traverse the chain of ViewLikeOps forward to find the longest sub-chain where the final
    // output type matches the input type of the original op. When found, the entire sub-chain
    // can be bypassed by replacing the last ViewLikeOp with the original input.
    while (IE::isPureViewOp(currentOp) && currentOp->hasOneUse()) {
        viewLikeOps.push_back(currentOp);

        auto outputType = mlir::cast<vpux::NDTypeInterface>(currentOp->getResult(0).getType());
        if (inputType == outputType) {
            break;
        }

        currentOp = *currentOp->getUsers().begin();
    }

    auto* lastViewLikeOp = viewLikeOps.back();
    auto outputType = mlir::cast<vpux::NDTypeInterface>(lastViewLikeOp->getResult(0).getType());

    if (inputType != outputType) {
        return mlir::failure();
    }
    rewriter.replaceAllOpUsesWith(lastViewLikeOp, origOp->getOperand(0));

    return mlir::success();
}

//
// RemoveViewLikeOpsChainPass
//

class RemoveViewLikeOpsChainPass final : public IE::impl::RemoveViewLikeOpsChainPassBase<RemoveViewLikeOpsChainPass> {
public:
    explicit RemoveViewLikeOpsChainPass(Logger log) {
        Base::initLogger(log, Base::getArgumentName());
    }

private:
    void safeRunOnFunc() final;
};

void RemoveViewLikeOpsChainPass::safeRunOnFunc() {
    auto func = getOperation();
    auto& ctx = getContext();

    mlir::RewritePatternSet patterns(&ctx);
    patterns.add<ViewLikeOpsChainRewriter>(&ctx, _log);

    collectOpsAndApplyPatterns(func, std::move(patterns));
}

}  // namespace

//
// createRemoveViewLikeOpsChainPass
//

std::unique_ptr<mlir::Pass> vpux::IE::createRemoveViewLikeOpsChainPass(Logger log) {
    return std::make_unique<RemoveViewLikeOpsChainPass>(log);
}
