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

#include "vpux/compiler/dialect/VPUMI40XX/utils.hpp"
#include "vpux/compiler/dialect/VPUASM/ops.hpp"
#include "vpux/compiler/dialect/VPUASM/types.hpp"
#include "vpux/compiler/dialect/VPUASM/utils.hpp"
#include "vpux/compiler/dialect/VPUIP/utils/utils.hpp"
#include "vpux/compiler/dialect/VPUIPDPU/ops.hpp"
#include "vpux/compiler/dialect/config/IR/resources.hpp"
#include "vpux/compiler/utils/platform_resources.hpp"

#include "vpux/compiler/dialect/ELF/IR/attributes.hpp"
#include "vpux/compiler/dialect/ELF/IR/ops.hpp"

namespace vpux {
namespace VPUASM {

vpux::VPURT::BufferSection getBufferLocation(mlir::Operation* symTableOp, mlir::SymbolRefAttr symRef) {
    VPUX_THROW_UNLESS(symTableOp->hasTrait<mlir::OpTrait::SymbolTable>(),
                      "The symTableOp parameter must have the SymbolTable trait");
    auto symTable = mlir::SymbolTable(symTableOp);

    auto referencedOp = symTable.lookupSymbolIn(symTableOp, symRef);
    if (auto logicalSec = referencedOp->getParentOfType<ELF::LogicalSectionOp>()) {
        return logicalSec.getSecLocation();
    } else if (auto dataSec = referencedOp->getParentOfType<ELF::DataSectionOp>()) {
        return dataSec.getSecLocation();
    }
    VPUX_THROW("BufferLocation can not be retrieved!");
}

vpux::VPURT::BufferSection getBufferLocation(ELF::SymbolReferenceMap& symRefMap, mlir::SymbolRefAttr symRef) {
    auto referencedOp = symRefMap.lookupSymbol(symRef);

    if (auto logicalSec = referencedOp->getParentOfType<ELF::LogicalSectionOp>()) {
        return logicalSec.getSecLocation();
    } else if (auto dataSec = referencedOp->getParentOfType<ELF::DataSectionOp>()) {
        return dataSec.getSecLocation();
    }
    VPUX_THROW("BufferLocation can not be retrieved!");
}

BufferType getBufferType(ELF::SymbolReferenceMap& symRefMap, mlir::SymbolRefAttr symRef) {
    auto referencedOp = symRefMap.lookupSymbol(symRef);

    if (auto bufferOp = mlir::dyn_cast<VPUASM::DeclareBufferOp>(referencedOp)) {
        return bufferOp.getBufferType();
    } else if (auto constantOp = mlir::dyn_cast<VPUASM::ConstBufferOp>(referencedOp)) {
        return constantOp.getBufferType();
    }
    VPUX_THROW("SymRef {0} does not point to a VPUASM::BufferType buffer", symRef.getLeafReference().getValue());
}

bool isWorkLoadManagementDMA(mlir::Operation* op) {
    return mlir::isa<VPUASM::DPUInvariantOp, VPUASM::DPUVariantOp, VPUIPDPU::DPUInvariantOp, VPUIPDPU::DPUVariantOp,
                     VPUASM::ActKernelInvocationOp, VPUASM::ActKernelRangeOp, VPUASM::DeclareTaskBufferOp>(op);
}

uint32_t getTileSelectMaskForBuffer(VPUASM::DeclareBufferOp buffer) {
    auto bufferLocation = buffer.getBufferType().getLocation();
    if (bufferLocation.getSection() != VPURT::BufferSection::CMX_NN) {
        return 0;
    }

    return VPUMI40XX::generateTileMask({static_cast<uint32_t>(bufferLocation.getSectionIndex())});
}

uint32_t getTileSelectMaskForBuffer(VPUASM::DeclareTaskBufferOp taskBuffer) {
    return VPUMI40XX::generateTileMask({static_cast<uint32_t>(taskBuffer.getTileIndex())});
}

uint32_t getActCompressionEntryTileMask(VPUASM::NNDMAOp dmaOp, ELF::SymbolReferenceMap& symRefMap) {
    auto actCompressionSizeEntry = dmaOp.getActCompressionSizeEntry();
    if (actCompressionSizeEntry.has_value()) {
        auto actCompBufferRef = symRefMap.lookupSymbol(actCompressionSizeEntry.value());
        VPUX_THROW_UNLESS(actCompBufferRef, "Could not find symbol name entry for {0} of {1}",
                          actCompressionSizeEntry.value(), dmaOp);

        if (mlir::isa<VPUASM::DeclareBufferOp>(actCompBufferRef)) {
            auto actCompBuffer = mlir::cast<VPUASM::DeclareBufferOp>(actCompBufferRef);
            return getTileSelectMaskForBuffer(actCompBuffer);
        }
    }
    return 0;
}

SparsityMap getSparsityMapBuffTileMask(VPUASM::NNDMAOp dmaOp, ELF::SymbolReferenceMap& symRefMap) {
    auto sparsityMapBuffer = dmaOp.getActCompressionSparsityMap();
    SparsityMap sparsityMap{};

    if (sparsityMapBuffer.has_value()) {
        auto sparsityMapBufferRef = symRefMap.lookupSymbol(sparsityMapBuffer.value());
        VPUX_THROW_UNLESS(sparsityMapBufferRef, "Could not find symbol name entry for {0} of {1}",
                          sparsityMapBuffer.value(), dmaOp);

        if (auto buffer = mlir::dyn_cast_if_present<VPUASM::DeclareBufferOp>(sparsityMapBufferRef)) {
            sparsityMap.tileSelectMaskForBuffer = getTileSelectMaskForBuffer(buffer);
            sparsityMap.size = buffer.getBinarySize(config::getArch(dmaOp));
        }
    }
    return sparsityMap;
}

void setResourceRequirement(mlir::ModuleOp moduleOp, elf::NetworkMetadata& metadata) {
    metadata.mResourceRequirements.nn_slice_count_ = VPUIP::getNumTilesUsed(moduleOp);
    uint32_t workspace_offset = 0;
    // E#179925 Compiler workspace is to be extended to include stacks and metadata
    if (config::getArch(moduleOp) != config::ArchKind::NPU40XX) {
        workspace_offset = CMX_WORKSPACE_OFFSET;
    }
    metadata.mResourceRequirements.nn_slice_length_ =
            workspace_offset +
            checked_cast<uint32_t>(config::getAvailableMemory(moduleOp, vpux::VPU::MemoryKind::CMX_NN).getByteSize());
}

SmallVector<uint32_t> getCMXStackFrames(mlir::ModuleOp moduleOp) {
    auto tileOp = config::getTileExecutor(moduleOp);
    auto tileCount = checked_cast<uint32_t>(tileOp.getCount());
    auto shvPerTile = checked_cast<uint32_t>(tileOp.getSubExecutor(config::ExecutorKind::SHAVE_ACT).getCount());

    SmallVector<uint32_t> stacksOffsets(shvPerTile);
    // SHAVE stacks grows backwards!
    // Set the address to the end of the allocated section so it does not override
    // outside of its buffer
    auto stackSize = static_cast<uint32_t>(CMX_SHAVE_STACK_SIZE.count());
    // First two stacks reserved at the beginning of the CMX space
    stacksOffsets[0] = stackSize;
    stacksOffsets[1] = stacksOffsets[0] + stackSize;

    const size_t defaultStacksNum = 2;
    // Check if additional stack frames are needed
    if (auto extraStacks = shvPerTile - defaultStacksNum; extraStacks > 0) {
        auto shaveStacksMem = config::getShaveStacksReservedMemory(moduleOp, VPU::MemoryKind::CMX_NN);
        VPUX_THROW_WHEN(shaveStacksMem == nullptr, "Missing reserved CMX memory for additional shave stack frames");
        auto shaveStacksMemOffset = shaveStacksMem.getOffset();
        VPUX_THROW_WHEN(shaveStacksMemOffset == std::nullopt,
                        "No address allocated for additional shave stack frames in CMX");
        auto shaveStacksMemSize = checked_cast<uint32_t>(shaveStacksMem.getByteSize());
        VPUX_THROW_WHEN(shaveStacksMemSize < extraStacks * stackSize,
                        "Insufficient memory allocated for additional shave stack frames in CMX");

        for (auto extraStackIdx : irange(extraStacks)) {
            // Additional stacks reserved after CMX workspace
            stacksOffsets[defaultStacksNum + extraStackIdx] =
                    CMX_WORKSPACE_OFFSET + shaveStacksMemOffset.value() + (extraStackIdx + 1) * stackSize;
        }
    }

    SmallVector<uint32_t> stackFrameAddrs(tileCount * shvPerTile);
    for (auto tileIdx : irange(tileCount)) {
        for (auto offset : llvm::enumerate(stacksOffsets)) {
            // Combine base address with offset to point inside reserved CMX memory
            stackFrameAddrs[tileIdx * shvPerTile + offset.index()] = offset.value() | CMX_BASE_ADDR;
        }
    }
    return stackFrameAddrs;
}

void insertBinaryDimsIntoVector(SmallVector<uint8_t>& dimsVector, vpux::NDTypeInterface ndType) {
    auto shape = ndType.getShape();
    const auto dimsOrder = ndType.getDimsOrder();
    const auto memShape = dimsOrder.toMemoryOrder(shape);

    for (auto& memDim : memShape | reversed) {
        auto dim = checked_cast<int32_t>(memDim);
        ArrayRef<uint8_t> valueAsArray(reinterpret_cast<const uint8_t*>(&dim), sizeof(dim));
        dimsVector.insert(dimsVector.end(), valueAsArray.begin(), valueAsArray.end());
    }
}

void insertBinaryStridesIntoVector(SmallVector<uint8_t>& stridesVector, vpux::NDTypeInterface ndType) {
    auto strides = ndType.getMemStrides();
    for (auto&& stride : strides | reversed) {
        ArrayRef<uint8_t> valueAsArray(reinterpret_cast<const uint8_t*>(&stride), sizeof(stride));
        stridesVector.insert(stridesVector.end(), valueAsArray.begin(), valueAsArray.end());
    }
}

}  // namespace VPUASM
}  // namespace vpux
