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

#include "vpux/compiler/dialect/VPUIP/IR/dialect.hpp"
#include "vpux/compiler/dialect/IE/IR/dialect.hpp"
#include "vpux/compiler/dialect/VPU/IR/dialect.hpp"
#include "vpux/compiler/dialect/VPU/IR/types.hpp"
#include "vpux/compiler/dialect/VPUIP/IR/dialect_interfaces.hpp"
#include "vpux/compiler/dialect/VPUIP/IR/ops.hpp"
#include "vpux/compiler/dialect/VPUIP/IR/types.hpp"
#include "vpux/compiler/dialect/VPUIP/interfaces/strategies.hpp"
#include "vpux/compiler/dialect/const/attributes/content.hpp"
#include "vpux/compiler/dialect/const/dialect.hpp"
#include "vpux/compiler/dialect/const/ops.hpp"
#include "vpux/compiler/dialect/core/IR/dialect.hpp"
#include "vpux/compiler/dialect/core/IR/unified_func_inliner_interface.hpp"
#include "vpux/compiler/dialect/net/IR/dialect.hpp"
#include "vpux/compiler/utils/error.hpp"
#include "vpux/compiler/utils/rewriter.hpp"

#include <llvm/ADT/TypeSwitch.h>
#include <mlir/Dialect/MemRef/IR/MemRef.h>
#include <mlir/Dialect/Quant/IR/Quant.h>
#include <mlir/Transforms/InliningUtils.h>

using namespace vpux;

namespace {

//
// VPUIPInlinerInterface
//

struct VPUIPInlinerInterface : public mlir::DialectInlinerInterface {
    using DialectInlinerInterface::DialectInlinerInterface;

    bool isLegalToInline(mlir::Operation*, mlir::Operation*, bool) const final {
        return true;
    }

    bool isLegalToInline(mlir::Operation*, mlir::Region*, bool, mlir::IRMapping&) const final {
        return true;
    }

    bool isLegalToInline(mlir::Region*, mlir::Region*, bool, mlir::IRMapping&) const final {
        return true;
    }
};

}  // namespace

//
// initialize
//

void vpux::VPUIP::VPUIPDialect::initialize() {
    addOperations<
#define GET_OP_LIST
#include <vpux/compiler/dialect/VPUIP/ops.cpp.inc>
            >();

    registerAttributes();
    registerTypes();

    addInterfaces<VPUIPInlinerInterface>();
    addInterfaces<VPUIP::StrategyFactoryCache>();

    Core::registerDispatchedInlinerInterface<VPUIP::VPUIPInlinerDispatchAttr, VPUIP::FuncInlinerInterface>(
            getContext());
}

//
// Operation executor attributes
//

namespace {

constexpr StringLiteral executorAttrName = "VPUIP.executor";
constexpr StringLiteral executorInstanceMaskAttrName = "VPUIP.executorIdx";

}  // namespace

void vpux::VPUIP::VPUIPDialect::setExecutor(mlir::async::ExecuteOp execOp, IndexedSymbolAttr executor) {
    VPUX_THROW_UNLESS(executor != nullptr, "Got an empty executor");
    execOp->setAttr(executorAttrName, executor);
}

void vpux::VPUIP::VPUIPDialect::setExecutorInstanceMask(mlir::async::ExecuteOp execOp,
                                                        mlir::ArrayAttr executorInstanceMask) {
    VPUX_THROW_UNLESS(executorInstanceMask != nullptr, "Got an empty executor instance");
    execOp->setAttr(executorInstanceMaskAttrName, executorInstanceMask);

    // In case of some DMA tasks port is also maintained as part of operation iself
    if (getExecutorKind(execOp) == config::ExecutorKind::DMA_NN) {
        if (executorInstanceMask.getValue().size() != 1) {
            // In case there are multiple executors subsequent passes should take care
            // of handling this
            return;
        }

        auto portIdxAttr = mlir::cast<mlir::IntegerAttr>(executorInstanceMask[0]);

        auto* bodyBlock = execOp.getBody();
        for (auto& op : bodyBlock->getOperations()) {
            auto opToCheck = &op;

            if (auto dmaOp = mlir::dyn_cast<VPUIP::DMATypeOpInterface>(opToCheck)) {
                dmaOp.setPortAttribute(portIdxAttr);
            }
        }
    }
}

llvm::StringLiteral vpux::VPUIP::VPUIPDialect::getExecutorAttrName() {
    return executorAttrName;
}

IndexedSymbolAttr vpux::VPUIP::VPUIPDialect::getExecutor(mlir::async::ExecuteOp execOp) {
    const auto executor = execOp->getAttr(executorAttrName);
    VPUX_THROW_UNLESS(executor != nullptr, "Can't find Executor attributes for Operation at '{0}'", execOp->getLoc());

    const auto executorSymbol = mlir::dyn_cast<vpux::IndexedSymbolAttr>(executor);
    VPUX_THROW_UNLESS(executorSymbol != nullptr, "Unsupported Executor attribute '{0}'", executorSymbol);

    return executorSymbol;
}

config::ExecutorKind vpux::VPUIP::VPUIPDialect::getExecutorKind(mlir::async::ExecuteOp execOp) {
    const auto executor = getExecutor(execOp);
    const auto maybeExecutorKind = vpux::config::symbolizeExecutorKind(executor.getLeafNameAttr());
    VPUX_THROW_WHEN(!maybeExecutorKind.has_value(), "Unsupported Executor Kind attribute '{0}'",
                    executor.getLeafNameAttr());

    return maybeExecutorKind.value();
}

bool vpux::VPUIP::VPUIPDialect::hasExecutorInstanceMask(mlir::async::ExecuteOp execOp) {
    return (execOp->getAttr(executorInstanceMaskAttrName) != nullptr);
}

bool vpux::VPUIP::VPUIPDialect::isComputeExecutorKind(config::ExecutorKind executorKind) {
    static const llvm::DenseSet<config::ExecutorKind> computeExecutors = {
            config::ExecutorKind::DPU, config::ExecutorKind::SHAVE_ACT, config::ExecutorKind::M2I};
    return computeExecutors.find(executorKind) != computeExecutors.end();
}

mlir::ArrayAttr vpux::VPUIP::VPUIPDialect::getExecutorInstanceMask(mlir::async::ExecuteOp execOp) {
    const auto executorInstanceMask = execOp->getAttr(executorInstanceMaskAttrName);
    VPUX_THROW_UNLESS(executorInstanceMask != nullptr, "Can't find Executor Instance attributes for Operation at '{0}'",
                      execOp->getLoc());

    const auto executorInstanceMaskSymbol = mlir::dyn_cast<mlir::ArrayAttr>(executorInstanceMask);
    VPUX_THROW_UNLESS(executorInstanceMaskSymbol != nullptr, "Unsupported Executor Instance attribute '{0}'",
                      executorInstanceMaskSymbol);

    return executorInstanceMaskSymbol;
}

//
// materializeConstant
//

mlir::Operation* vpux::VPUIP::VPUIPDialect::materializeConstant(mlir::OpBuilder& builder, mlir::Attribute value,
                                                                mlir::Type type, mlir::Location loc) {
    if (!mlir::isa<Const::ContentAttr>(value)) {
        (void)errorAt(loc, "Can't materialize VPUIP Constant from Attribute '{0}'", value);
        return nullptr;
    }

    if (!mlir::isa<mlir::MemRefType>(type)) {
        (void)errorAt(loc, "Can't materialize VPUIP Constant for Type '{0}'", type);
        return nullptr;
    }

    return builder.create<Const::DeclareOp>(loc, type, mlir::cast<Const::ContentAttr>(value));
}

//
// Generated
//

#include <vpux/compiler/dialect/VPUIP/dialect.cpp.inc>
