/*
 * Copyright 2026 Rive
 */

// Parse-only KTX2 reader for Rive runtime.
//
// Extracts vkFormat, logical pixel dims, and mip-level block payloads from a
// KTX2 container. Does NOT decompress — block data is returned verbatim so
// the renderer can hand it straight to the GPU's compressed texture upload
// path. Today only BC7 (UNORM / SRGB) is recognized; other vkFormat values
// are rejected.
//
// Spec: https://registry.khronos.org/KTX/specs/2.0/ktxspec.v2.html

#include "rive/decoders/decode_ktx2.hpp"
#include "rive/decoders/astc_footprints.hpp"
#include "rive/decoders/texture_decoder.hpp"

#include <algorithm>
#include <cstdio>
#include <cstring>

namespace rive
{
namespace
{
constexpr uint8_t Ktx2Identifier[12] = {
    0xAB,
    0x4B,
    0x54,
    0x58,
    0x20,
    0x32,
    0x30,
    0xBB,
    0x0D,
    0x0A,
    0x1A,
    0x0A,
};

constexpr uint32_t VK_FORMAT_BC7_UNORM_BLOCK = 145;
constexpr uint32_t VK_FORMAT_BC7_SRGB_BLOCK = 146;
constexpr uint32_t VK_FORMAT_ETC2_R8G8B8A8_UNORM_BLOCK = 151;
constexpr uint32_t VK_FORMAT_ETC2_R8G8B8A8_SRGB_BLOCK = 152;
// ASTC block range endpoints; the body maps the family by index between them.
constexpr uint32_t VK_FORMAT_ASTC_4x4_UNORM_BLOCK = 157;
constexpr uint32_t VK_FORMAT_ASTC_12x12_SRGB_BLOCK = 184;
constexpr uint32_t SupercompressionNone = 0;

#pragma pack(push, 1)
struct Ktx2Header
{
    uint32_t vkFormat;
    uint32_t typeSize;
    uint32_t pixelWidth;
    uint32_t pixelHeight;
    uint32_t pixelDepth;
    uint32_t layerCount;
    uint32_t faceCount;
    uint32_t levelCount;
    uint32_t supercompressionScheme;
    uint32_t dfdByteOffset;
    uint32_t dfdByteLength;
    uint32_t kvdByteOffset;
    uint32_t kvdByteLength;
    uint64_t sgdByteOffset;
    uint64_t sgdByteLength;
};
struct Ktx2LevelIndex
{
    uint64_t byteOffset;
    uint64_t byteLength;
    uint64_t uncompressedByteLength;
};
#pragma pack(pop)
static_assert(sizeof(Ktx2Header) == 68, "KTX2 header must be 68 bytes");
static_assert(sizeof(Ktx2LevelIndex) == 24,
              "KTX2 level index entry must be 24 bytes");

constexpr uint32_t MaxDimension = 16384;
constexpr uint32_t MaxLevels = 16;

inline uint64_t expectedBlockBytes(uint32_t pixelWidth,
                                   uint32_t pixelHeight,
                                   uint32_t blockWidth,
                                   uint32_t blockHeight,
                                   uint32_t bytesPerBlock)
{
    const uint64_t blocksX = (pixelWidth + blockWidth - 1u) / blockWidth;
    const uint64_t blocksY = (pixelHeight + blockHeight - 1u) / blockHeight;
    return blocksX * blocksY * bytesPerBlock;
}
} // namespace

bool DecodeKtx2(const uint8_t* bytes,
                size_t byteCount,
                Ktx2DecodeResult& out,
                const Ktx2HwSupport& hwSupport)
{
    if (byteCount < sizeof(Ktx2Identifier) + sizeof(Ktx2Header))
    {
        std::fprintf(stderr, "DecodeKtx2: file too small\n");
        return false;
    }
    if (std::memcmp(bytes, Ktx2Identifier, sizeof(Ktx2Identifier)) != 0)
    {
        std::fprintf(stderr, "DecodeKtx2: bad magic\n");
        return false;
    }

    Ktx2Header header;
    std::memcpy(&header, bytes + sizeof(Ktx2Identifier), sizeof(header));

    // Map vkFormat → (GPUTextureFormat, blockWidth, blockHeight,
    // bytesPerBlock, srgb).
    GPUTextureFormat outFormat;
    uint8_t blockWidth = 4;
    uint8_t blockHeight = 4;
    uint32_t bytesPerBlock = 16;
    bool srgb = false;
    if (header.vkFormat == VK_FORMAT_BC7_UNORM_BLOCK ||
        header.vkFormat == VK_FORMAT_BC7_SRGB_BLOCK)
    {
        outFormat = GPUTextureFormat::bc7;
        srgb = (header.vkFormat == VK_FORMAT_BC7_SRGB_BLOCK);
    }

    else if (header.vkFormat == VK_FORMAT_ETC2_R8G8B8A8_UNORM_BLOCK ||
             header.vkFormat == VK_FORMAT_ETC2_R8G8B8A8_SRGB_BLOCK)
    {
        outFormat = GPUTextureFormat::etc2;
        // ETC2 RGBA8 = 8 bytes EAC alpha + 8 bytes ETC2 RGB = 16 per block.
        bytesPerBlock = 16;
        srgb = (header.vkFormat == VK_FORMAT_ETC2_R8G8B8A8_SRGB_BLOCK);
    }

    else if (header.vkFormat >= VK_FORMAT_ASTC_4x4_UNORM_BLOCK &&
             header.vkFormat <= VK_FORMAT_ASTC_12x12_SRGB_BLOCK)
    {
        const uint32_t idx =
            (header.vkFormat - VK_FORMAT_ASTC_4x4_UNORM_BLOCK) / 2u;
        outFormat = GPUTextureFormat::astc;

        blockWidth = AstcFootprints[idx].width;
        blockHeight = AstcFootprints[idx].height;
        srgb = (header.vkFormat % 2u) == 0u; // SRGB at UNORM+1.
    }
    else
    {
        std::fprintf(stderr,
                     "DecodeKtx2: unsupported vkFormat %u\n",
                     header.vkFormat);
        return false;
    }

    if (header.supercompressionScheme != SupercompressionNone)
    {
        std::fprintf(stderr,
                     "DecodeKtx2: supercompressionScheme %u not supported\n",
                     header.supercompressionScheme);
        return false;
    }
    if (header.faceCount != 1 || header.layerCount != 0)
    {
        std::fprintf(
            stderr,
            "DecodeKtx2: cubemaps/arrays not supported (faces=%u layers=%u)\n",
            header.faceCount,
            header.layerCount);
        return false;
    }
    if (header.pixelWidth == 0 || header.pixelWidth > MaxDimension ||
        header.pixelHeight == 0 || header.pixelHeight > MaxDimension)
    {
        std::fprintf(stderr,
                     "DecodeKtx2: dimensions out of range (%ux%u)\n",
                     header.pixelWidth,
                     header.pixelHeight);
        return false;
    }

    const uint32_t levelCount = header.levelCount == 0 ? 1u : header.levelCount;
    if (levelCount > MaxLevels)
    {
        std::fprintf(stderr,
                     "DecodeKtx2: levelCount %u exceeds cap %u\n",
                     levelCount,
                     MaxLevels);
        return false;
    }

    const size_t levelIndexOffset = sizeof(Ktx2Identifier) + sizeof(Ktx2Header);
    const size_t levelIndexBytes =
        static_cast<size_t>(levelCount) * sizeof(Ktx2LevelIndex);
    if (byteCount < levelIndexOffset + levelIndexBytes)
    {
        std::fprintf(stderr, "DecodeKtx2: truncated level index\n");
        return false;
    }

    // Read level entries (level 0 first in the index = largest mip).
    std::vector<Ktx2LevelIndex> entries(levelCount);
    std::memcpy(entries.data(), bytes + levelIndexOffset, levelIndexBytes);

    // Validate each level + sum total payload size for the output buffer.
    uint64_t totalBytes = 0;
    for (uint32_t i = 0; i < levelCount; ++i)
    {
        const Ktx2LevelIndex& e = entries[i];
        const uint32_t logW = std::max<uint32_t>(1u, header.pixelWidth >> i);
        const uint32_t logH = std::max<uint32_t>(1u, header.pixelHeight >> i);
        const uint64_t expected = expectedBlockBytes(logW,
                                                     logH,
                                                     blockWidth,
                                                     blockHeight,
                                                     bytesPerBlock);
        if (e.byteLength != expected)
        {
            std::fprintf(
                stderr,
                "DecodeKtx2: level %u byteLength %llu != block grid %llu "
                "for %ux%u\n",
                i,
                static_cast<unsigned long long>(e.byteLength),
                static_cast<unsigned long long>(expected),
                logW,
                logH);
            return false;
        }
        if (e.byteOffset > byteCount || e.byteLength > byteCount - e.byteOffset)
        {
            std::fprintf(stderr, "DecodeKtx2: level %u out of buffer\n", i);
            return false;
        }
        totalBytes += e.byteLength;
    }

    // Concatenate level 0 .. N-1 into one contiguous buffer (largest first).
    out.format = outFormat;
    out.pixelWidth = header.pixelWidth;
    out.pixelHeight = header.pixelHeight;
    out.levelCount = levelCount;
    out.blockWidth = blockWidth;
    out.blockHeight = blockHeight;
    out.srgb = srgb;
    out.softwareDecoded = false;
    out.blocks.resize(static_cast<size_t>(totalBytes));
    size_t writeOffset = 0;
    for (uint32_t i = 0; i < levelCount; ++i)
    {
        const Ktx2LevelIndex& e = entries[i];
        std::memcpy(out.blocks.data() + writeOffset,
                    bytes + e.byteOffset,
                    static_cast<size_t>(e.byteLength));
        writeOffset += static_cast<size_t>(e.byteLength);
    }

    // HW-cap fallback: if the backend can't sample this format directly,
    // software-decode mip 0 to RGBA8 in place. Caller uploads as rgba32.
    bool needFallback = false;
    switch (out.format)
    {
        case GPUTextureFormat::bc1:
        case GPUTextureFormat::bc2:
        case GPUTextureFormat::bc3:
        case GPUTextureFormat::bc7:
            needFallback = !hwSupport.bc;
            break;
        case GPUTextureFormat::astc:
            needFallback = !hwSupport.astc;
            break;
        case GPUTextureFormat::etc2:
            needFallback = !hwSupport.etc2;
            break;
        default:
            break;
    }
    if (needFallback)
    {
        // Decode every level. Source layout: level 0 first, levels tight
        // (matches how we just wrote `out.blocks`). Output layout: each
        // level's logical width * height * 4 bytes, also tight.
        //
        // We allocate the decoded chain into a temp buffer first so the
        // original block bytes remain valid for the per-level decode calls.
        std::vector<uint8_t> decoded;
        size_t totalRgba = 0;
        for (uint32_t i = 0; i < levelCount; ++i)
        {

            const uint32_t logW = std::max<uint32_t>(1u, out.pixelWidth >> i);
            const uint32_t logH = std::max<uint32_t>(1u, out.pixelHeight >> i);

            totalRgba += static_cast<size_t>(logW) * logH * 4;
        }
        decoded.reserve(totalRgba);

        size_t srcOffset = 0;
        for (uint32_t i = 0; i < levelCount; ++i)
        {

            const uint32_t logW = std::max<uint32_t>(1u, out.pixelWidth >> i);
            const uint32_t logH = std::max<uint32_t>(1u, out.pixelHeight >> i);
            const size_t levelBytes =
                static_cast<size_t>(entries[i].byteLength);
            auto bmp = decode_texture(out.blocks.data() + srcOffset,
                                      levelBytes,
                                      logW,
                                      logH,
                                      out.format,
                                      out.blockWidth,
                                      out.blockHeight);
            if (!bmp)
            {
                std::fprintf(stderr,
                             "DecodeKtx2: HW lacks support for format %u "
                             "and software decoder unavailable (level %u)\n",
                             static_cast<unsigned>(out.format),
                             i);
                return false;
            }
            // Match the PNG runtime path: premultiplied texels.
            bmp->pixelFormat(Bitmap::PixelFormat::RGBAPremul);
            decoded.insert(decoded.end(),
                           bmp->bytes(),
                           bmp->bytes() + bmp->numBytes());
            srcOffset += levelBytes;
        }

        out.blocks = std::move(decoded);
        out.format = GPUTextureFormat::rgba32;

        out.blockWidth = 1;
        out.blockHeight = 1;
        out.srgb = false;
        out.softwareDecoded = true;
        // `out.levelCount` already matches the KTX2 level count.
    }

    return true;
}

} // namespace rive