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

#include <vpux/compiler/dialect/VPUMI40XX/ops.hpp>
#include <vpux/compiler/dialect/VPURegMapped/types.hpp>
#include "vpux/compiler/utils/error.hpp"

using namespace vpux;
using namespace VPUMI40XX;

mlir::LogicalResult OpRanges::verify() {
    const auto op = getOperation();

    const auto types = getTypes();
    const auto begins = getBegins();
    const auto ends = getEnds();

    // expected to be generated by SameVariadicOperandSize trait
    // see https://discourse.llvm.org/t/mlir-samevariadicoperandsize-doesnt-generate-check-in-verifier/77091
    const auto typesSize = types.size();
    const auto beginsSize = begins.size();
    const auto endsSize = ends.size();
    if (typesSize != beginsSize || beginsSize != endsSize) {
        return errorAt(op,
                       "'VPUMI40XX.OpRanges' has sizes of arguments 'types' = {0}, 'begins' = {1} and 'ends' = {2} "
                       "that are expected to be the same",
                       typesSize, beginsSize, endsSize);
    }

    mlir::DenseSet<std::tuple<VPURegMapped::TaskType, uint32_t, uint32_t>> ranges;
    for (size_t i = 0; i < typesSize; ++i) {
        // consistency of arguments is already checked by built-in verifier
        const auto taskType = mlir::cast<VPURegMapped::TaskTypeAttr>(types[i]).getValue();
        const auto begin = mlir::cast<VPURegMapped::IndexType>(begins[i].getType());
        const auto end = mlir::cast<VPURegMapped::IndexType>(ends[i].getType());

        if (begin.getValue() != 0) {
            return errorAt(op,
                           "'VPUMI40XX.OpRanges' range at idx {0}, has begin with value component of index = {1}, "
                           "expected to be 0",
                           i, begin.getValue());
        }

        if (begin.getTileIdx() != end.getTileIdx()) {
            return errorAt(op,
                           "'VPUMI40XX.OpRanges' range at idx {0}, has begin with tile index = {1} and end with tile "
                           "index = {2}, that are expected to be the same",
                           i, begin.getTileIdx(), end.getTileIdx());
        }

        if (begin.getListIdx() != end.getListIdx()) {
            return errorAt(op,
                           "'VPUMI40XX.OpRanges' range at idx {0}, has begin with list index = {1} and end with list "
                           "index = {2}, that are expected to be the same",
                           i, begin.getListIdx(), end.getListIdx());
        }

        if (begin.getValue() > end.getValue()) {
            return errorAt(op,
                           "'VPUMI40XX.OpRanges' range at idx {0} doesn't monotonously non-decreasing from begin to "
                           "end, begin/end value component of index are {1} and {2}",
                           i, begin.getValue(), end.getValue());
        }

        const auto [_, isNewElement] = ranges.insert(std::make_tuple(taskType, begin.getTileIdx(), begin.getListIdx()));
        if (!isNewElement) {
            return errorAt(op, "'VPUMI40XX.OpRanges' range at idx {0} is duplicate of already present range", i);
        }
    }

    return mlir::success();
}

namespace {

template <class Range>
Range findRange(mlir::ArrayAttr types, mlir::Operation::operand_range values, VPURegMapped::TaskType type,
                VPURegMapped::IndexType index) {
    const auto taskTypes = types | transformed([](auto attr) {
                               return mlir::cast<VPURegMapped::TaskTypeAttr>(attr).getValue();
                           });
    const auto zipped = zip(taskTypes, values);

    const auto areFromTheSameRange = [](auto lhs, auto rhs) {
        return lhs.getTileIdx() == rhs.getTileIdx() && lhs.getListIdx() == rhs.getListIdx();
    };
    const auto isInRange = [&](auto entry) {
        const auto& [entryType, entryBegin] = entry;
        const auto beginIndex = mlir::cast<VPURegMapped::IndexType>(entryBegin.getType());
        return type == entryType && areFromTheSameRange(index, beginIndex);
    };
    const auto range = llvm::find_if(zipped, isInRange);
    return range == std::end(zipped) ? Range{} : Range{std::get<1>(*range)};
}

}  // namespace

TaskForwardRange OpRanges::getForwardRange(VPURegMapped::TaskType type, VPURegMapped::IndexType index) {
    return findRange<TaskForwardRange>(getTypes(), getBegins(), type, index);
}

TaskBackwardRange OpRanges::getBackwardRange(VPURegMapped::TaskType type, VPURegMapped::IndexType index) {
    return findRange<TaskBackwardRange>(getTypes(), getEnds(), type, index);
}
