forge-ibl-bake 0.1.0

// SPDX-License-Identifier: MIT OR Apache-2.0
// forge-ibl-bake · deveraux.dev
//! Hand-rolled KTX2 container writer (Khronos KTX 2.0 spec, no supercompression).
//!
//! Emits spec-conformant KTX2 files for 2D textures and cubemaps with any number
//! of mip levels. Supports VkFormat 83 (R16G16_SFLOAT) and 97 (R16G16B16A16_SFLOAT).
//! Level data is stored coarsest-first per spec §4.7; DFD describes each format.
//! DESIGN_TARGET: validated by the independent `ktx2` crate reader in test #6.

use std::path::Path;

use crate::error::BakeError;

/// KTX2 file identifier (12 bytes): «KTX 20»\r\n\x1A\n
pub const KTX2_MAGIC: [u8; 12] =
    [0xAB, 0x4B, 0x54, 0x58, 0x20, 0x32, 0x30, 0xBB, 0x0D, 0x0A, 0x1A, 0x0A];

pub const VK_FORMAT_R16G16_SFLOAT: u32 = 83;
pub const VK_FORMAT_R16G16B16A16_SFLOAT: u32 = 97;

/// Write a 2D texture with one mip level.
pub fn write_ktx2_2d(
    path: &Path,
    width: u32,
    height: u32,
    vk_format: u32,
    pixel_data: &[u8],
) -> Result<(), BakeError> {
    let bpp = format_bytes_per_pixel(vk_format)?;
    let expected = (width as usize) * (height as usize) * bpp;
    if pixel_data.len() != expected {
        return Err(BakeError::InvalidSize(format!(
            "2D: expected {expected} bytes, got {}",
            pixel_data.len()
        )));
    }

    let dfd = build_dfd(vk_format)?;
    let buf = build_ktx2(
        vk_format,
        width,
        height,
        0,  // pixelDepth
        0,  // layerCount (= 1 layer)
        1,  // faceCount
        &dfd,
        &[pixel_data],
    );
    std::fs::write(path, &buf)?;
    Ok(())
}

/// Write a cubemap (faceCount=6) with `mip_data.len()` mip levels.
/// `mip_data[0]` = finest mip (6 faces × face_size² × bpp bytes).
/// `mip_data[N-1]` = coarsest mip (6 faces × (face_size/2^(N-1))² × bpp).
pub fn write_ktx2_cubemap(
    path: &Path,
    face_size: u32,
    vk_format: u32,
    mip_data: &[Vec<u8>],
) -> Result<(), BakeError> {
    if mip_data.is_empty() {
        return Err(BakeError::InvalidSize("cubemap requires ≥1 mip".into()));
    }
    let bpp = format_bytes_per_pixel(vk_format)?;
    for (i, mip) in mip_data.iter().enumerate() {
        let dim = ((face_size >> i) as usize).max(1);
        let expected = 6 * dim * dim * bpp;
        if mip.len() != expected {
            return Err(BakeError::InvalidSize(format!(
                "cubemap mip {i}: expected {expected} bytes (6 × {dim}² × {bpp}), got {}",
                mip.len()
            )));
        }
    }
    let dfd = build_dfd(vk_format)?;
    let slices: Vec<&[u8]> = mip_data.iter().map(|v| v.as_slice()).collect();
    let buf = build_ktx2(
        vk_format,
        face_size,
        face_size,
        0,
        0,
        6,
        &dfd,
        &slices,
    );
    std::fs::write(path, &buf)?;
    Ok(())
}

// ── Core builder ─────────────────────────────────────────────────────────────

/// Layout:
/// [Header 80B][LevelIndex 24B×L][DFD][padding][mip data coarsest→finest]
fn build_ktx2(
    vk_format: u32,
    pixel_width: u32,
    pixel_height: u32,
    pixel_depth: u32,
    layer_count: u32,
    face_count: u32,
    dfd: &[u8],
    mip_slices: &[&[u8]],
) -> Vec<u8> {
    let level_count = mip_slices.len() as u32;
    let bpp = format_bytes_per_pixel(vk_format).unwrap_or(1);

    // ── Compute section offsets ──
    let header_size = 80usize;
    let index_size = 24 * level_count as usize;
    let dfd_offset = header_size + index_size;
    let dfd_len = dfd.len();
    let kvd_offset = dfd_offset + dfd_len;
    let kvd_len = 0usize;

    // Image data alignment: max(bytes_per_texel, 4).
    let alignment = bpp.max(4) as u64;
    let data_start = align_up((kvd_offset + kvd_len) as u64, alignment) as usize;

    // ── Compute per-level offsets (coarsest → finest in file) ──
    // Spec §4.7: data is stored from mip[N-1] (smallest) to mip[0] (largest).
    let mut level_offsets = vec![0u64; level_count as usize];
    let mut cursor = data_start as u64;
    for mip in (0..level_count as usize).rev() {
        level_offsets[mip] = cursor;
        cursor += mip_slices[mip].len() as u64;
        cursor = align_up(cursor, alignment); // KTX2 §4.7: each level starts on alignment boundary
    }

    // ── Assemble ──
    let total = cursor as usize;
    let mut buf = Vec::with_capacity(total.max(data_start));

    // Header (80 bytes).
    buf.extend_from_slice(&KTX2_MAGIC);
    put_u32(&mut buf, vk_format);
    put_u32(&mut buf, 2); // typeSize: 2 bytes per component (SFLOAT16)
    put_u32(&mut buf, pixel_width);
    put_u32(&mut buf, pixel_height);
    put_u32(&mut buf, pixel_depth);
    put_u32(&mut buf, layer_count);
    put_u32(&mut buf, face_count);
    put_u32(&mut buf, level_count);
    put_u32(&mut buf, 0); // supercompressionScheme = none
    put_u32(&mut buf, dfd_offset as u32);
    put_u32(&mut buf, dfd_len as u32);
    put_u32(&mut buf, kvd_offset as u32);
    put_u32(&mut buf, kvd_len as u32);
    put_u64(&mut buf, 0); // sgdByteOffset = 0
    put_u64(&mut buf, 0); // sgdByteLength = 0
    debug_assert_eq!(buf.len(), 80, "header must be exactly 80 bytes");

    // Level index (24 bytes × levelCount).
    for i in 0..level_count as usize {
        put_u64(&mut buf, level_offsets[i]);
        put_u64(&mut buf, mip_slices[i].len() as u64); // byteLength
        put_u64(&mut buf, mip_slices[i].len() as u64); // uncompressedByteLength
    }
    debug_assert_eq!(buf.len(), header_size + index_size);

    // DFD.
    buf.extend_from_slice(dfd);

    // Alignment padding before data.
    buf.resize(data_start, 0u8);

    // Image data: coarsest first (mip[N-1]), finest last (mip[0]).
    for mip in (0..mip_slices.len()).rev() {
        buf.extend_from_slice(mip_slices[mip]);
        let aligned = align_up(buf.len() as u64, alignment) as usize;
        buf.resize(aligned, 0u8);
    }

    buf
}

// ── Data Format Descriptor ────────────────────────────────────────────────────
//
// Minimal KDF basic block for SFLOAT16 channels.
// Ref: Khronos Data Format Specification 1.3, §13 & §B.4.

fn build_dfd(vk_format: u32) -> Result<Vec<u8>, BakeError> {
    match vk_format {
        VK_FORMAT_R16G16_SFLOAT => Ok(dfd_sfloat(4, &[
            (0, 0xC0),   // R: bit 0, signed float
            (16, 0xC1),  // G: bit 16, signed float
        ])),
        VK_FORMAT_R16G16B16A16_SFLOAT => Ok(dfd_sfloat(8, &[
            (0, 0xC0),   // R
            (16, 0xC1),  // G
            (32, 0xC2),  // B
            (48, 0xCF),  // A (channel index 15 = KHR_DF_CHANNEL_RGBSDA_ALPHA)
        ])),
        other => Err(BakeError::UnsupportedFormat(other)),
    }
}

/// Build a KDF basic descriptor block for SFLOAT16 channels.
/// `samples`: (bitOffset, channelType byte) per channel.
/// channelType = channelIndex | KHR_DF_SAMPLE_DATATYPE_SIGNED(0x40) | FLOAT(0x80).
fn dfd_sfloat(bytes_per_texel: u8, samples: &[(u16, u8)]) -> Vec<u8> {
    let n = samples.len() as u16;
    let block_size = 24u16 + n * 16; // header + sample descriptors
    let dfd_total = 4u32 + block_size as u32;

    let mut b = Vec::with_capacity(dfd_total as usize);
    put_u32(&mut b, dfd_total); // dfdTotalSize (includes this u32)

    // Basic descriptor block header (24 bytes).
    put_u16(&mut b, 0); // vendorId = 0 (Khronos)
    put_u16(&mut b, 0); // descriptorType = 0 (basic)
    put_u16(&mut b, 2); // versionNumber = 2
    put_u16(&mut b, block_size);
    b.push(1); // colorModel = 1 (KHR_DF_MODEL_RGBSDA)
    b.push(1); // colorPrimaries = 1 (BT709)
    b.push(1); // transferFunction = 1 (LINEAR)
    b.push(0); // flags = 0
    b.extend_from_slice(&[0u8; 4]); // texelBlockDimension[4] = 0 (1×1×1×1 block)
    b.push(bytes_per_texel); // bytesPlane0
    b.extend_from_slice(&[0u8; 7]); // bytesPlane1..7

    // Sample descriptors (16 bytes each).
    for &(bit_offset, ch_type) in samples {
        put_u16(&mut b, bit_offset);
        b.push(15); // bitLength - 1: 16-bit channel → 15
        b.push(ch_type);
        put_u32(&mut b, f32::NEG_INFINITY.to_bits()); // lower bound
        put_u32(&mut b, f32::INFINITY.to_bits()); // upper bound
    }

    b
}

// ── Helpers ───────────────────────────────────────────────────────────────────

pub fn format_bytes_per_pixel(vk_format: u32) -> Result<usize, BakeError> {
    match vk_format {
        VK_FORMAT_R16G16_SFLOAT => Ok(4),
        VK_FORMAT_R16G16B16A16_SFLOAT => Ok(8),
        other => Err(BakeError::UnsupportedFormat(other)),
    }
}

fn align_up(n: u64, align: u64) -> u64 {
    (n + align - 1) / align * align
}

fn put_u16(buf: &mut Vec<u8>, v: u16) {
    buf.extend_from_slice(&v.to_le_bytes());
}

fn put_u32(buf: &mut Vec<u8>, v: u32) {
    buf.extend_from_slice(&v.to_le_bytes());
}

fn put_u64(buf: &mut Vec<u8>, v: u64) {
    buf.extend_from_slice(&v.to_le_bytes());
}