atx_reader 0.1.0

//! Decoding pipeline: take parsed container bytes and produce RGBA pixels.

use std::io::Cursor;

use astc_decode::astc_decode;

use crate::error::{AtxError, Result};
use crate::format::AstcFootprint;
use crate::parser::{AtxContainer, AtxHeader, TexturePayload};

/// How the raw ASTC blocks are laid out inside the `astc` payload.
#[derive(Debug, Clone, Copy)]
pub enum PayloadLayout {
    /// Let [`decode_with`] choose the layout based on the chunk type:
    /// macro-tiled Morton for the `astc` chunk, linear for the `LZFS`
    /// chunk. Identical to [`PayloadLayout::MacroTiledMorton`] with
    /// `macro_blocks = 32` for the lower-level [`AtxDecoder`].
    Auto,

    /// Plain row-major order across the (padded) block grid. Use for the
    /// LZFSE-decompressed payload, or for new `.atx` variants that don't
    /// apply Apple's swizzling.
    Linear,

    /// Apple's iOS macro-tiled Morton layout: the padded block grid is
    /// divided into `macro_blocks × macro_blocks` squares, the squares are
    /// emitted in row-major order, and inside each square the blocks are in
    /// Morton (Z-order). `macro_blocks = 32` matches the verified sample.
    MacroTiledMorton { macro_blocks: u32 },
}

impl Default for PayloadLayout {
    fn default() -> Self {
        Self::Auto
    }
}

/// Options controlling the decode.
#[derive(Debug, Clone, Copy)]
pub struct DecodeOptions {
    /// ASTC block footprint. Defaults to `Astc4x4`.
    pub footprint: AstcFootprint,
    /// Block layout in the payload. When set to [`PayloadLayout::Auto`],
    /// [`decode_with`] picks the layout based on the chunk type — Morton
    /// macro-tiles for the `astc` chunk, linear for the `LZFS` chunk.
    pub layout: PayloadLayout,
    /// Padded image dimensions (in pixels). When `None`, the padding is
    /// derived from the layout: for `MacroTiledMorton { macro_blocks }`,
    /// each dimension is rounded up to a multiple of
    /// `macro_blocks * block_dim` pixels; for `Linear`, to a multiple of
    /// the block dimension.
    pub padded_size: Option<(u32, u32)>,
}

impl Default for DecodeOptions {
    fn default() -> Self {
        Self {
            footprint: AstcFootprint::Astc4x4,
            layout: PayloadLayout::Auto,
            padded_size: None,
        }
    }
}

impl DecodeOptions {
    fn layout_or(&self, fallback: PayloadLayout) -> PayloadLayout {
        match self.layout {
            PayloadLayout::Auto => fallback,
            other => other,
        }
    }
}

/// An RGBA8 image produced by [`decode`] / [`decode_with`].
#[derive(Debug, Clone)]
pub struct DecodedImage {
    pub width: u32,
    pub height: u32,
    /// Row-major, 4 bytes per pixel, top-left origin.
    pub pixels: Vec<u8>,
}

/// Convenience: parse + decode in one call, using default options
/// (ASTC 4×4, Apple macro-tiled Morton with `macro_blocks = 32`).
pub fn decode(bytes: &[u8]) -> Result<DecodedImage> {
    decode_with(bytes, &DecodeOptions::default())
}

/// Convenience: parse + decode using custom options. Auto-detects the
/// payload variant (`astc` chunk vs LZFSE-compressed `LZFS` chunk) and the
/// appropriate default layout when the user did not specify one.
pub fn decode_with(bytes: &[u8], opts: &DecodeOptions) -> Result<DecodedImage> {
    let container = AtxContainer::parse(bytes)?;
    let header = container.header()?;
    let payload = container.texture_payload()?;

    let (blocks, layout): (std::borrow::Cow<'_, [u8]>, PayloadLayout) = match payload {
        TexturePayload::Astc(p) => (
            std::borrow::Cow::Borrowed(p),
            opts.layout_or(PayloadLayout::MacroTiledMorton { macro_blocks: 32 }),
        ),
        TexturePayload::Lzfse(p) => {
            let decompressed = decompress_lzfse(p)?;
            (
                std::borrow::Cow::Owned(decompressed),
                opts.layout_or(PayloadLayout::Linear),
            )
        }
    };

    let mut effective = *opts;
    effective.layout = layout;
    AtxDecoder::new(header, blocks.as_ref())
        .with_options(effective)
        .decode()
}

#[cfg(feature = "lzfse")]
fn decompress_lzfse(stream: &[u8]) -> Result<Vec<u8>> {
    use lzfse_rust::LzfseRingDecoder;
    let mut dec = LzfseRingDecoder::default();
    let mut out = Vec::new();
    dec.decode(&mut Cursor::new(stream), &mut out)
        .map_err(|e| AtxError::LzfseDecode(e.to_string()))?;
    Ok(out)
}

#[cfg(not(feature = "lzfse"))]
fn decompress_lzfse(_stream: &[u8]) -> Result<Vec<u8>> {
    Err(AtxError::LzfseUnavailable)
}

/// Decode an `.atx` to an [`image::RgbaImage`].
#[cfg(feature = "image")]
pub fn decode_to_image(bytes: &[u8]) -> Result<image::RgbaImage> {
    let img = decode(bytes)?;
    image::RgbaImage::from_raw(img.width, img.height, img.pixels)
        .ok_or_else(|| AtxError::AstcDecode("RGBA buffer size did not match dimensions".into()))
}

/// Lower-level decoder. Useful when you already hold a parsed header and
/// payload, or want to tweak options between parsing and decoding.
#[derive(Debug, Clone)]
pub struct AtxDecoder<'a> {
    header: AtxHeader,
    payload: &'a [u8],
    opts: DecodeOptions,
}

impl<'a> AtxDecoder<'a> {
    pub fn new(header: AtxHeader, payload: &'a [u8]) -> Self {
        Self {
            header,
            payload,
            opts: DecodeOptions::default(),
        }
    }

    pub fn with_options(mut self, opts: DecodeOptions) -> Self {
        self.opts = opts;
        self
    }

    pub fn footprint(mut self, fp: AstcFootprint) -> Self {
        self.opts.footprint = fp;
        self
    }

    pub fn layout(mut self, layout: PayloadLayout) -> Self {
        self.opts.layout = layout;
        self
    }

    pub fn padded_size(mut self, w: u32, h: u32) -> Self {
        self.opts.padded_size = Some((w, h));
        self
    }

    pub fn decode(self) -> Result<DecodedImage> {
        let w = self.header.width;
        let h = self.header.height;
        if w == 0 || h == 0 {
            return Err(AtxError::AstcDecode(format!(
                "invalid header dimensions {w}x{h}"
            )));
        }
        let fp = self.opts.footprint;
        let bw = fp.block_width();
        let bh = fp.block_height();

        let layout = match self.opts.layout {
            PayloadLayout::Auto => PayloadLayout::MacroTiledMorton { macro_blocks: 32 },
            other => other,
        };

        let (pad_w, pad_h) = self
            .opts
            .padded_size
            .unwrap_or_else(|| derive_padding(w, h, bw, bh, &layout));

        if pad_w < w || pad_h < h {
            return Err(AtxError::AstcDecode(format!(
                "padded size {pad_w}x{pad_h} is smaller than image {w}x{h}"
            )));
        }

        let blocks_w = pad_w.div_ceil(bw);
        let blocks_h = pad_h.div_ceil(bh);
        let total_blocks = (blocks_w as usize) * (blocks_h as usize);
        let needed = total_blocks * 16;

        let linear = match layout {
            PayloadLayout::Linear | PayloadLayout::Auto => {
                if self.payload.len() < needed {
                    return Err(AtxError::TooShort {
                        needed,
                        got: self.payload.len(),
                    });
                }
                std::borrow::Cow::Borrowed(&self.payload[..needed])
            }
            PayloadLayout::MacroTiledMorton { macro_blocks } => {
                std::borrow::Cow::Owned(linearize_macro_morton(
                    self.payload,
                    blocks_w,
                    blocks_h,
                    macro_blocks,
                )?)
            }
        };

        let mut full = vec![0u8; (pad_w as usize) * (pad_h as usize) * 4];
        let stride = pad_w as usize * 4;
        astc_decode(
            Cursor::new(linear.as_ref()),
            pad_w,
            pad_h,
            fp.to_astc_decode(),
            |x, y, block| {
                if x < pad_w && y < pad_h {
                    let i = (y as usize) * stride + (x as usize) * 4;
                    full[i..i + 4].copy_from_slice(&block);
                }
            },
        )
        .map_err(|e| AtxError::AstcDecode(e.to_string()))?;

        // Crop pad_w x pad_h → w x h.
        let mut pixels = vec![0u8; (w as usize) * (h as usize) * 4];
        let dst_stride = (w as usize) * 4;
        let copy_bytes = dst_stride;
        for y in 0..h as usize {
            let src = y * stride;
            let dst = y * dst_stride;
            pixels[dst..dst + copy_bytes].copy_from_slice(&full[src..src + copy_bytes]);
        }

        Ok(DecodedImage {
            width: w,
            height: h,
            pixels,
        })
    }
}

fn derive_padding(w: u32, h: u32, bw: u32, bh: u32, layout: &PayloadLayout) -> (u32, u32) {
    let (mult_x, mult_y) = match *layout {
        PayloadLayout::Auto | PayloadLayout::Linear => (bw, bh),
        PayloadLayout::MacroTiledMorton { macro_blocks } => (bw * macro_blocks, bh * macro_blocks),
    };
    (round_up(w, mult_x), round_up(h, mult_y))
}

fn round_up(value: u32, multiple: u32) -> u32 {
    if multiple == 0 {
        return value;
    }
    value.div_ceil(multiple) * multiple
}

fn linearize_macro_morton(
    payload: &[u8],
    blocks_w: u32,
    blocks_h: u32,
    macro_blocks: u32,
) -> Result<Vec<u8>> {
    if macro_blocks == 0 {
        return Err(AtxError::AstcDecode("macro_blocks must be > 0".into()));
    }
    if blocks_w % macro_blocks != 0 || blocks_h % macro_blocks != 0 {
        return Err(AtxError::AstcDecode(format!(
            "padded block grid {blocks_w}x{blocks_h} is not divisible by macro_blocks={macro_blocks}",
        )));
    }
    let macros_x = blocks_w / macro_blocks;
    let macros_y = blocks_h / macro_blocks;
    let total_blocks = (blocks_w as usize) * (blocks_h as usize);
    let needed = total_blocks * 16;
    if payload.len() < needed {
        return Err(AtxError::TooShort {
            needed,
            got: payload.len(),
        });
    }

    let mut linear = vec![0u8; needed];
    let stride_blocks = blocks_w as usize;
    let mut src_idx = 0usize;
    for my in 0..macros_y {
        for mx in 0..macros_x {
            for i in 0..(macro_blocks * macro_blocks) {
                let (lx, ly) = morton_xy(i);
                let dst_bx = mx * macro_blocks + lx;
                let dst_by = my * macro_blocks + ly;
                let dst_idx = ((dst_by as usize) * stride_blocks + dst_bx as usize) * 16;
                linear[dst_idx..dst_idx + 16].copy_from_slice(&payload[src_idx..src_idx + 16]);
                src_idx += 16;
            }
        }
    }
    Ok(linear)
}

fn morton_xy(idx: u32) -> (u32, u32) {
    let mut x = 0u32;
    let mut y = 0u32;
    for i in 0..16 {
        x |= ((idx >> (2 * i)) & 1) << i;
        y |= ((idx >> (2 * i + 1)) & 1) << i;
    }
    (x, y)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn morton_first_block_is_origin() {
        assert_eq!(morton_xy(0), (0, 0));
    }

    #[test]
    fn morton_z_order_progression() {
        // Standard Z-order: indices 0..4 trace a 2x2 block.
        assert_eq!(morton_xy(1), (1, 0));
        assert_eq!(morton_xy(2), (0, 1));
        assert_eq!(morton_xy(3), (1, 1));
    }

    #[test]
    fn round_up_handles_exact_multiples() {
        assert_eq!(round_up(128, 128), 128);
        assert_eq!(round_up(0, 128), 0);
        assert_eq!(round_up(129, 128), 256);
        assert_eq!(round_up(1170, 128), 1280);
        assert_eq!(round_up(2532, 128), 2560);
    }

    #[test]
    fn padding_matches_verified_sample() {
        let (pw, ph) = derive_padding(
            1170,
            2532,
            4,
            4,
            &PayloadLayout::MacroTiledMorton { macro_blocks: 32 },
        );
        assert_eq!((pw, ph), (1280, 2560));
    }
}