/* Copyright (c) V-Nova International Limited 2025-2026. All rights reserved.
 * This software is licensed under the BSD-3-Clause-Clear License by V-Nova Limited.
 * No patent licenses are granted under this license. For enquiries about patent licenses,
 * please contact legal@v-nova.com.
 * The LCEVCdec software is a stand-alone project and is NOT A CONTRIBUTION to any other project.
 * If the software is incorporated into another project, THE TERMS OF THE BSD-3-CLAUSE-CLEAR LICENSE
 * AND THE ADDITIONAL LICENSING INFORMATION CONTAINED IN THIS FILE MUST BE MAINTAINED, AND THE
 * SOFTWARE DOES NOT AND MUST NOT ADOPT THE LICENSE OF THE INCORPORATING PROJECT. However, the
 * software may be incorporated into a project under a compatible license provided the requirements
 * of the BSD-3-Clause-Clear license are respected, and V-Nova Limited remains
 * licensor of the software ONLY UNDER the BSD-3-Clause-Clear license (not the compatible license).
 * ANY ONWARD DISTRIBUTION, WHETHER STAND-ALONE OR AS PART OF ANY OTHER PROJECT, REMAINS SUBJECT TO
 * THE EXCLUSION OF PATENT LICENSES PROVISION OF THE BSD-3-CLAUSE-CLEAR LICENSE. */

// Vulkan Apply shader

#version 450

#include "specialization_ids.h"

// Specialized workgropup size
layout(local_size_x_id = SPECID_LOCAL_SIZE_X,
       local_size_y_id = SPECID_LOCAL_SIZE_Y,
       local_size_z_id = SPECID_LOCAL_SIZE_Z) in;

layout(std430, set = 0, binding = 0) buffer restrict readonly CommandBuffer {
    uint samples[];
} cmd;

layout(std430, set = 0, binding = 1) buffer restrict PlaneBuffer {
    uint samples[];
} dst;

layout(push_constant) uniform consts {
    uint srcWidth;
    uint srcHeight;
    uint residualOffset;
    uint containerStride;
    uint containerOffset;
    uint saturate;
    uint tileX;
    uint tileY;
    uint tileWidth;
    uint commandOffset; // uint32 offset into merged command buffer for this tile
    uint commandCount;   // number of commands in this tile
    uint blocksInRowMagic; // ceil(2^32 / blocksInRow); 0 means blocksInRow == 1
    uint dds;
    uint rasterOrder;
};

const uint ADD = 0;
const uint SET = 1;
const uint SETZERO = 2;
const uint CLEARSET = 3;
const int SATURATION = 32767;

#include "common.glsl"

uint applyCounter = 0;

void applyTu(uint firstContainerInTu, uint firstResidual, uint operation, uint tuWidth)
{
    for (uint j = 0; j < tuWidth; ++j) {
        for (uint i = 0; i < (tuWidth >> 1); ++i) {
            const uint containerIndex = firstContainerInTu + i + j * containerStride;
            const uint residualPairIndex = commandOffset + residualOffset + (firstResidual >> 1) +
                                           applyCounter * tuWidth * (tuWidth >> 1) + i + j * (tuWidth >> 1);

            if (containerIndex >= containerOffset + containerStride * srcHeight) {
                return;
            }

            ivec2 pixelPair = ivec2(SATURATION);
            if (saturate != 1) {
                ivec2 residualPair;
                if (operation == ADD || operation == SET) {
                    residualPair = unpack16bitInternals(cmd.samples[residualPairIndex]);
                }

                switch(operation) {
                    case ADD : pixelPair = saturateS16(unpack16bitInternals(dst.samples[containerIndex]) + residualPair); break;
                    case SET : pixelPair = residualPair; break;
                    case SETZERO : pixelPair = ivec2(0); break;
                    default : break;
                }
            }
            dst.samples[containerIndex] = pack16bitInternals(pixelPair);
        }
    }
    if (operation != SETZERO) {
        applyCounter++;
    }
}

void processSubBlock(uint bitmaskMost, uint bitmaskLeast, uint operation,
                      uint startingContainerInSubBlock, uint blockX, uint blockY,
                      uint subBlockIdx, uint tu, uint tuWidth, uint tusInBlockWidth)
{
    const uint containerWidth = (srcWidth + 1) >> 1;
    for (uint i = 0; i < 64; ++i) {
        const uint bitIndex = 31 - (i & 31u); // same as 31 - (i % 32)
        const uint bitSelector = 1u << bitIndex;
        const uint currentBitmask = (i < 32u) ? bitmaskMost : bitmaskLeast;
        // tusInBlockWidth is 8 (dds) or 16 (non-dds) — both power of 2.
        const uint ix = i & (tusInBlockWidth - 1u);
        const uint iy = i / tusInBlockWidth;
        const uint x = blockX * 16 + ix * (tuWidth >> 1);
        const uint y = blockY * 32 + subBlockIdx * 8 + iy * tuWidth;
        if(x >= containerWidth || y >= srcHeight) {
            continue;
        }

        const uint index = startingContainerInSubBlock + ix * (tuWidth >> 1) + iy * tuWidth * containerStride;
        if ((currentBitmask & bitSelector) != 0u) {
            switch(operation) {
                case ADD : applyTu(index, tu, ADD, tuWidth); break;
                case SETZERO : applyTu(index, tu, SETZERO, tuWidth); break;
                default : applyTu(index, tu, SET, tuWidth); // SET or CLEARSET
            }
        } else if (operation == CLEARSET) {
            applyTu(index, 0, SETZERO, tuWidth);
        }
    }
}

void processBlock(const uint bitmask[8], uint operation, uint block, uint tu, uint tuWidth)
{
    const uint pixelsInBlockLength = 32;
    const uint tusInBlockWidth = (dds != 0u) ? 8 : 16;
    const uint width = (tileWidth != 0) ? tileWidth : srcWidth;
    const uint blocksInRow = (width + (pixelsInBlockLength - 1u)) / pixelsInBlockLength;
    // Recover blockY = block / blocksInRow via a 32x32->64 multiply by the
    // CPU-supplied reciprocal, avoiding an integer divide. blocksInRowMagic == 0
    // is the sentinel for blocksInRow == 1 (the magic 2^32 wouldn't fit in u32).
    uint blockX;
    uint blockY;
    if (blocksInRowMagic != 0u) {
        uint hi;
        uint lo;
        umulExtended(block, blocksInRowMagic, hi, lo);
        blockY = hi;
        blockX = block - blockY * blocksInRow;
    } else {
        blockY = block;
        blockX = 0u;
    }
    const uint firstContainerInBLock = containerOffset + (tileX >> 1) + blockX * (pixelsInBlockLength >> 1) + tileY * containerStride + blockY * containerStride * pixelsInBlockLength;

    for(int i = 0; i<((dds != 0u) ? 1 : 4); ++i) {
        processSubBlock(bitmask[2*i], bitmask[2*i + 1], operation, firstContainerInBLock + i * 8 * containerStride, blockX, blockY, uint(i), tu, tuWidth, tusInBlockWidth);
    }
}

void processRaster(const uint bitmask[8], uint operation, uint block, uint tu, uint tuWidth)
{
    // containerWidth is the logical image width in uint32 units (no stride padding).
    // TU indices from the decoder are laid out using the image width, so the
    // coordinate mapping must use containerWidth rather than containerStride.
    const uint containerWidth = (srcWidth + 1) >> 1;
    const uint containersInRaster = (256 >> 1) * tuWidth;
    const uint containerStartIndex = block * containersInRaster;
    const uint step = tuWidth >> 1;
    // Track (x, y) = (containerIndex % containerWidth, containerIndex / containerWidth)
    // incrementally so the per-bit loop body avoids an integer divmod against
    // the runtime containerWidth. The starting divide runs once per command.
    uint y = containerStartIndex / containerWidth;
    uint x = containerStartIndex - y * containerWidth;
    for (uint i = 0; i < 256; ++i) {
        const uint segment = i >> 5; // same as i / 32
        const uint bitIndex = 31 - (i & 31u); // same as 31 - (i % 32)
        const uint bitSelector = 1u << bitIndex;
        if ((bitmask[segment] & bitSelector) != 0u
            && x + step <= containerWidth && y * tuWidth + tuWidth <= srcHeight) {
            // Memory address uses containerStride (includes any row padding).
            const uint index = containerOffset + x + y * tuWidth * containerStride;
            applyTu(index, tu, ADD, tuWidth);
        }
        x += step;
        // while (not if) for the degenerate step > containerWidth case.
        while (x >= containerWidth) {
            x -= containerWidth;
            y += 1u;
        }
    }
}

void main()
{
    if (gl_GlobalInvocationID.x >= commandCount) {
        return;
    }

    const uint CMDSIZE = 10; // 64 bit command + four 64 bit bitmasks
    const uint tuWidth = (dds != 0u) ? 4 : 2;
    const uint commandIndex = uint(commandOffset) + gl_GlobalInvocationID.x * CMDSIZE;
    const ivec4 command = ivec4(cmd.samples[commandIndex], cmd.samples[commandIndex + 1],
                          cmd.samples[commandIndex + 2], cmd.samples[commandIndex + 3]);

    const uint operation = command.x & 0x3;        // One of GpuApplyOperation.
    const uint block = (command.x >> 2) & 0x3FFFF; // Number of block within temporal buffer to operate on.
    const uint tu = ((command.x >> 20) & 0xFFF) |
              ((command.y & 0x7FFF) << 12); // Offset into CommandBuffer of first residual for operation.
    const uint bitStart = (command.y >> 15) & 0xFF;  // First set bit in bitmask.
    const uint bitCount = (command.y >> 23) & 0x1FF; // Number of set bits in bitmask.

    uint bitmask[8];
    bitmask[1] = command.z;
    bitmask[0] = command.w;
    if (dds == 0u || rasterOrder != 0u) { // when in raster order, we make use of all 4 bitmasks for DDS
        const int map[] = {3, 2, 5, 4, 7, 6};
        for (int i = 0; i < 6; ++i) {
            bitmask[map[i]] = cmd.samples[commandIndex + 4 + i];
        }
    }

    if (rasterOrder != 0u) {
        // temporal off - not tiled
        processRaster(bitmask, operation, block, tu, tuWidth);
    }  else {
        // temporal on or tile-mode block order
        processBlock(bitmask, operation, block, tu, tuWidth);
    }
}
