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

#include "vpux/compiler/dialect/VPU/IR/dialect.hpp"
#include "vpux/compiler/dialect/VPU/IR/ops/internal.hpp"
#include "vpux/compiler/dialect/VPU/transforms/passes.hpp"
#include "vpux/compiler/dialect/VPU/utils/vertical_fusion/v1/merge_vf_region_rewriter.hpp"
#include "vpux/compiler/dialect/VPU/utils/vertical_fusion/v2/merge_vf_region_rewriter.hpp"
#include "vpux/compiler/dialect/config/IR/utils.hpp"
#include "vpux/compiler/utils/attributes.hpp"
#include "vpux/compiler/utils/rewriter.hpp"

namespace vpux::VPU {
#define GEN_PASS_DECL_MERGEVFSUBGRAPHS
#define GEN_PASS_DEF_MERGEVFSUBGRAPHS
#include "vpux/compiler/dialect/VPU/passes.hpp.inc"
}  // namespace vpux::VPU

using namespace vpux;
using namespace VPU;

namespace {

//
// MergeVfSubgraphsPass
//

class MergeVfSubgraphsPass final : public VPU::impl::MergeVfSubgraphsBase<MergeVfSubgraphsPass> {
public:
    explicit MergeVfSubgraphsPass(bool enableVerticalFusionPipelining, bool enablePrefetchTiling,
                                  const WorkloadManagementMode workloadManagementMode, Logger log)
            : _enableVerticalFusionPipelining(enableVerticalFusionPipelining),
              _enablePrefetchTiling(enablePrefetchTiling),
              _workloadManagementMode(workloadManagementMode) {
        Base::initLogger(log, Base::getArgumentName());
    }

    mlir::LogicalResult initialize(mlir::MLIRContext* ctx) final;

private:
    void safeRunOnFunc() final;
    bool _enableVerticalFusionPipelining = false;
    bool _enablePrefetchTiling = true;
    WorkloadManagementMode _workloadManagementMode = WorkloadManagementMode::PWLM_V0_1_PAGES;
};

mlir::LogicalResult MergeVfSubgraphsPass::initialize(mlir::MLIRContext* ctx) {
    if (mlir::failed(Base::initialize(ctx))) {
        return mlir::failure();
    }
    if (enableVerticalFusionPipelining.hasValue()) {
        _log.trace("Overloading MergeVfSubgraphsPass argument by MLIR variable");
        _enableVerticalFusionPipelining = enableVerticalFusionPipelining;
    }
    if (tilingMode.hasValue()) {
        _log.trace("Overloading the default value {0} of the '_enablePrefetchTiling' field to the value {1} of the "
                   "pass option 'tilingMode' generated by MLIR",
                   _enablePrefetchTiling, tilingMode.getValue());
        _enablePrefetchTiling = tilingMode.getValue() == "PREFETCH";
    }
    if (workloadManagementModeOpt.hasValue()) {
        _workloadManagementMode = workloadManagementModeOpt.getValue();
    }
    return mlir::success();
}

//
// safeRunOnModule
//

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

    auto module = func->getParentOfType<mlir::ModuleOp>();
    auto maybeLayerCostModelAnalysis = getCachedParentAnalysis<VPU::LayerCostModelAnalysis>(module);
    auto layerCostModel =
            VPU::LayerCostModelAnalysis::getOrCreateLayerCostModel(maybeLayerCostModelAnalysis, &ctx, _log);

    const auto costFunction = std::make_unique<VPU::LayerVPUNNCost>(func, layerCostModel, _log);

    // TODO rewrite with largest spill detection on each step
    SmallVector<int64_t> maxTiling;
    func->walk([&](VPU::VerticalFusionOp op) {
        const auto tiling = parseIntArrayAttr<int64_t>(op.getTilingStrategy());
        auto dim = getVFTilingDim(tiling);
        if (dim.has_value()) {
            maxTiling.emplace_back(tiling[dim.value().ind()]);
        }
    });

    auto maxTilingOp = std::max_element(maxTiling.begin(), maxTiling.end());

    mlir::RewritePatternSet patterns(&ctx);
    if (_workloadManagementMode <= WorkloadManagementMode::PWLM_V0_1_PAGES) {
        patterns.add<VPU::VF::v1::MergeVFRegionRewriter>(&ctx, _enableVerticalFusionPipelining, _enablePrefetchTiling,
                                                         costFunction, _log);
    } else {
        patterns.add<VPU::VF::v2::MergeVFRegionRewriter>(&ctx, _enableVerticalFusionPipelining, _enablePrefetchTiling,
                                                         costFunction, _log);
    }

    auto config = getDefaultGreedyRewriteConfig();
    // VF is a cost-based pass. Whether some VF blocks can be merged may depend on preceding blocks.
    // However, if a block is scheduled before its previous block, additional iterations are needed.
    // In the worst case, this can exceed the default maxIterations limit.
    // Therefore, we increase maxIterations to ensure all possible merges are performed.
    // TODO: #192550 Refactor the pass to avoid iteration-related issues.
    config.setMaxIterations(config.getMaxIterations() * 10);
    config.setUseTopDownTraversal(maxTilingOp == maxTiling.end() || maxTilingOp != std::prev(maxTiling.end()));
    if (mlir::failed(mlir::applyPatternsGreedily(func, std::move(patterns), config))) {
        signalPassFailure();
    }
}

}  // namespace

//
// createMergeVfSubgraphsPass
//

std::unique_ptr<mlir::Pass> VPU::createMergeVfSubgraphsPass(bool enableVerticalFusionPipelining,
                                                            bool enablePrefetchTiling,
                                                            const WorkloadManagementMode workloadManagementMode,
                                                            Logger log) {
    return std::make_unique<MergeVfSubgraphsPass>(enableVerticalFusionPipelining, enablePrefetchTiling,
                                                  workloadManagementMode, log);
}
