veridian_kernel/media/image_codecs/

png.rs

//! PNG decoder with full DEFLATE/zlib decompression.
//!
//! Supports all critical chunks (IHDR, PLTE, IDAT, IEND), ancillary chunks
//! (tRNS, gAMA), all 5 filter types, Adam7 interlacing, color types
//! 0/2/3/4/6, and bit depths 1/2/4/8/16.

#![allow(dead_code)]

use alloc::vec::Vec;

use super::{read_be_u16, read_be_u32, DecodedImage, ImageCodecError};

// ============================================================================
// PNG DECODER
// ============================================================================

/// PNG 8-byte signature.
pub(crate) const PNG_SIGNATURE: [u8; 8] = [137, 80, 78, 71, 13, 10, 26, 10];

/// PNG color types.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PngColorType {
    Grayscale = 0,
    Rgb = 2,
    Indexed = 3,
    GrayscaleAlpha = 4,
    Rgba = 6,
}

impl PngColorType {
    fn from_u8(v: u8) -> Result<Self, ImageCodecError> {
        match v {
            0 => Ok(Self::Grayscale),
            2 => Ok(Self::Rgb),
            3 => Ok(Self::Indexed),
            4 => Ok(Self::GrayscaleAlpha),
            6 => Ok(Self::Rgba),
            _ => Err(ImageCodecError::Unsupported),
        }
    }

    /// Number of channels (samples per pixel).
    fn channels(self) -> usize {
        match self {
            Self::Grayscale => 1,
            Self::Rgb => 3,
            Self::Indexed => 1,
            Self::GrayscaleAlpha => 2,
            Self::Rgba => 4,
        }
    }
}

/// Parsed IHDR data.
#[derive(Debug, Clone, Copy)]
struct PngIhdr {
    width: u32,
    height: u32,
    bit_depth: u8,
    color_type: PngColorType,
    interlace: u8,
}

/// Decode a PNG image from raw file data.
pub fn decode_png(data: &[u8]) -> Result<DecodedImage, ImageCodecError> {
    // Verify signature
    if data.len() < 8 || data[..8] != PNG_SIGNATURE {
        return Err(ImageCodecError::InvalidSignature);
    }

    // Parse chunks
    let mut pos: usize = 8;
    let mut ihdr: Option<PngIhdr> = None;
    let mut palette: Vec<(u8, u8, u8)> = Vec::new();
    let mut trns: Vec<u8> = Vec::new();
    let mut idat_data: Vec<u8> = Vec::new();
    let mut _gamma: u32 = 0; // stored as gamma * 100000

    while pos + 8 <= data.len() {
        let chunk_len = read_be_u32(data, pos) as usize;
        let chunk_type = &data[pos + 4..pos + 8];
        let chunk_data_start = pos + 8;
        let chunk_end = chunk_data_start + chunk_len;

        if chunk_end + 4 > data.len() {
            // Not enough data for chunk + CRC
            break;
        }

        match chunk_type {
            b"IHDR" => {
                if chunk_len < 13 {
                    return Err(ImageCodecError::CorruptData);
                }
                let cd = &data[chunk_data_start..chunk_end];
                ihdr = Some(PngIhdr {
                    width: read_be_u32(cd, 0),
                    height: read_be_u32(cd, 4),
                    bit_depth: cd[8],
                    color_type: PngColorType::from_u8(cd[9])?,
                    interlace: cd[12],
                });
            }
            b"PLTE" => {
                if !chunk_len.is_multiple_of(3) {
                    return Err(ImageCodecError::CorruptData);
                }
                let cd = &data[chunk_data_start..chunk_end];
                palette.clear();
                let mut i = 0;
                while i + 2 < cd.len() {
                    palette.push((cd[i], cd[i + 1], cd[i + 2]));
                    i += 3;
                }
            }
            b"tRNS" => {
                trns = data[chunk_data_start..chunk_end].to_vec();
            }
            b"gAMA" => {
                if chunk_len >= 4 {
                    _gamma = read_be_u32(data, chunk_data_start);
                }
            }
            b"IDAT" => {
                idat_data.extend_from_slice(&data[chunk_data_start..chunk_end]);
            }
            b"IEND" => {
                break;
            }
            _ => {
                // Skip unknown/ancillary chunks
            }
        }

        pos = chunk_end + 4; // skip CRC
    }

    let ihdr = ihdr.ok_or(ImageCodecError::CorruptData)?;
    if ihdr.width == 0 || ihdr.height == 0 {
        return Err(ImageCodecError::InvalidDimensions);
    }

    // Validate bit depth for color type
    match ihdr.color_type {
        PngColorType::Grayscale => {
            if !matches!(ihdr.bit_depth, 1 | 2 | 4 | 8 | 16) {
                return Err(ImageCodecError::Unsupported);
            }
        }
        PngColorType::Rgb | PngColorType::GrayscaleAlpha | PngColorType::Rgba => {
            if !matches!(ihdr.bit_depth, 8 | 16) {
                return Err(ImageCodecError::Unsupported);
            }
        }
        PngColorType::Indexed => {
            if !matches!(ihdr.bit_depth, 1 | 2 | 4 | 8) {
                return Err(ImageCodecError::Unsupported);
            }
        }
    }

    // Decompress zlib-wrapped IDAT data
    let raw_data = zlib_decompress(&idat_data)?;

    // Unfilter and produce RGBA output
    if ihdr.interlace == 1 {
        decode_png_interlaced(&ihdr, &raw_data, &palette, &trns)
    } else {
        decode_png_non_interlaced(&ihdr, &raw_data, &palette, &trns)
    }
}

/// Decode non-interlaced PNG scanlines.
fn decode_png_non_interlaced(
    ihdr: &PngIhdr,
    raw: &[u8],
    palette: &[(u8, u8, u8)],
    trns: &[u8],
) -> Result<DecodedImage, ImageCodecError> {
    let w = ihdr.width as usize;
    let h = ihdr.height as usize;
    let channels = ihdr.color_type.channels();
    let bits_per_pixel = channels * (ihdr.bit_depth as usize);
    let bytes_per_row = (w * bits_per_pixel).div_ceil(8);
    let bpp_bytes = bits_per_pixel.div_ceil(8); // filter byte stride

    let mut img = DecodedImage::new(ihdr.width, ihdr.height);
    let mut prev_row: Vec<u8> = alloc::vec![0u8; bytes_per_row];
    let mut pos: usize = 0;

    for y in 0..h {
        if pos >= raw.len() {
            return Err(ImageCodecError::TruncatedData);
        }
        let filter = raw[pos];
        pos += 1;

        if pos + bytes_per_row > raw.len() {
            return Err(ImageCodecError::TruncatedData);
        }

        let mut current_row: Vec<u8> = raw[pos..pos + bytes_per_row].to_vec();
        pos += bytes_per_row;

        // Apply PNG filter reconstruction
        png_unfilter(filter, &mut current_row, &prev_row, bpp_bytes)?;

        // Convert scanline to RGBA pixels
        png_scanline_to_rgba(
            &current_row,
            ihdr,
            palette,
            trns,
            &mut img,
            y as u32,
            0,
            w as u32,
        );

        prev_row = current_row;
    }

    Ok(img)
}

/// Decode Adam7 interlaced PNG.
fn decode_png_interlaced(
    ihdr: &PngIhdr,
    raw: &[u8],
    palette: &[(u8, u8, u8)],
    trns: &[u8],
) -> Result<DecodedImage, ImageCodecError> {
    let w = ihdr.width as usize;
    let h = ihdr.height as usize;
    let channels = ihdr.color_type.channels();
    let bits_per_pixel = channels * (ihdr.bit_depth as usize);

    // Adam7 pass parameters: (x_start, y_start, x_step, y_step)
    const ADAM7: [(usize, usize, usize, usize); 7] = [
        (0, 0, 8, 8),
        (4, 0, 8, 8),
        (0, 4, 4, 8),
        (2, 0, 4, 4),
        (0, 2, 2, 4),
        (1, 0, 2, 2),
        (0, 1, 1, 2),
    ];

    let mut img = DecodedImage::new(ihdr.width, ihdr.height);
    let mut pos: usize = 0;

    for &(x_start, y_start, x_step, y_step) in &ADAM7 {
        let pass_w = if x_start >= w {
            0
        } else {
            (w - x_start).div_ceil(x_step)
        };
        let pass_h = if y_start >= h {
            0
        } else {
            (h - y_start).div_ceil(y_step)
        };

        if pass_w == 0 || pass_h == 0 {
            continue;
        }

        let bytes_per_row = (pass_w * bits_per_pixel).div_ceil(8);
        let bpp_bytes = bits_per_pixel.div_ceil(8);
        let mut prev_row: Vec<u8> = alloc::vec![0u8; bytes_per_row];

        for pass_y in 0..pass_h {
            if pos >= raw.len() {
                return Err(ImageCodecError::TruncatedData);
            }
            let filter = raw[pos];
            pos += 1;

            if pos + bytes_per_row > raw.len() {
                return Err(ImageCodecError::TruncatedData);
            }

            let mut current_row: Vec<u8> = raw[pos..pos + bytes_per_row].to_vec();
            pos += bytes_per_row;

            png_unfilter(filter, &mut current_row, &prev_row, bpp_bytes)?;

            // Place pixels at correct interlaced positions
            let out_y = y_start + pass_y * y_step;
            for pass_x in 0..pass_w {
                let out_x = x_start + pass_x * x_step;
                let pixel = png_extract_pixel(&current_row, pass_x, ihdr, palette, trns);
                img.set_pixel(
                    out_x as u32,
                    out_y as u32,
                    pixel.0,
                    pixel.1,
                    pixel.2,
                    pixel.3,
                );
            }

            prev_row = current_row;
        }
    }

    Ok(img)
}

/// PNG filter reconstruction (RFC 2083 Section 9).
fn png_unfilter(
    filter: u8,
    current: &mut [u8],
    prev: &[u8],
    bpp: usize,
) -> Result<(), ImageCodecError> {
    let len = current.len();
    match filter {
        0 => {} // None
        1 => {
            // Sub
            for i in bpp..len {
                current[i] = current[i].wrapping_add(current[i - bpp]);
            }
        }
        2 => {
            // Up
            for i in 0..len {
                current[i] = current[i].wrapping_add(prev[i]);
            }
        }
        3 => {
            // Average
            for i in 0..len {
                let a = if i >= bpp { current[i - bpp] as u16 } else { 0 };
                let b = prev[i] as u16;
                current[i] = current[i].wrapping_add(((a + b) / 2) as u8);
            }
        }
        4 => {
            // Paeth
            for i in 0..len {
                let a = if i >= bpp { current[i - bpp] as i32 } else { 0 };
                let b = prev[i] as i32;
                let c = if i >= bpp { prev[i - bpp] as i32 } else { 0 };
                current[i] = current[i].wrapping_add(paeth_predictor(a, b, c) as u8);
            }
        }
        _ => return Err(ImageCodecError::CorruptData),
    }
    Ok(())
}

/// Paeth predictor function (integer only).
#[inline]
fn paeth_predictor(a: i32, b: i32, c: i32) -> i32 {
    let p = a + b - c;
    let pa = (p - a).abs();
    let pb = (p - b).abs();
    let pc = (p - c).abs();
    if pa <= pb && pa <= pc {
        a
    } else if pb <= pc {
        b
    } else {
        c
    }
}

/// Convert a full PNG scanline to RGBA pixels in the output image.
fn png_scanline_to_rgba(
    row: &[u8],
    ihdr: &PngIhdr,
    palette: &[(u8, u8, u8)],
    trns: &[u8],
    img: &mut DecodedImage,
    y: u32,
    x_start: u32,
    count: u32,
) {
    for x in 0..count {
        let pixel = png_extract_pixel(row, x as usize, ihdr, palette, trns);
        img.set_pixel(x_start + x, y, pixel.0, pixel.1, pixel.2, pixel.3);
    }
}

/// Extract a single pixel from a PNG scanline, returning (R, G, B, A).
fn png_extract_pixel(
    row: &[u8],
    x: usize,
    ihdr: &PngIhdr,
    palette: &[(u8, u8, u8)],
    trns: &[u8],
) -> (u8, u8, u8, u8) {
    let bd = ihdr.bit_depth as usize;

    match ihdr.color_type {
        PngColorType::Grayscale => {
            let v = extract_sample(row, x, bd);
            let v8 = scale_to_8bit(v, bd);
            let a = if trns.len() >= 2 {
                let trns_val = read_be_u16(trns, 0) as usize;
                if v == trns_val {
                    0
                } else {
                    255
                }
            } else {
                255
            };
            (v8, v8, v8, a)
        }
        PngColorType::Rgb => {
            let bytes_per_sample = if bd == 16 { 2 } else { 1 };
            let off = x * 3 * bytes_per_sample;
            let (r, g, b) = if bd == 16 {
                if off + 5 < row.len() {
                    (row[off], row[off + 2], row[off + 4])
                } else {
                    (0, 0, 0)
                }
            } else if off + 2 < row.len() {
                (row[off], row[off + 1], row[off + 2])
            } else {
                (0, 0, 0)
            };
            let a = if trns.len() >= 6 {
                let tr = read_be_u16(trns, 0);
                let tg = read_be_u16(trns, 2);
                let tb = read_be_u16(trns, 4);
                let (cr, cg, cb) = if bd == 16 {
                    (
                        read_be_u16(row, off),
                        read_be_u16(row, off + 2),
                        read_be_u16(row, off + 4),
                    )
                } else {
                    (r as u16, g as u16, b as u16)
                };
                if cr == tr && cg == tg && cb == tb {
                    0
                } else {
                    255
                }
            } else {
                255
            };
            (r, g, b, a)
        }
        PngColorType::Indexed => {
            let idx = extract_sample(row, x, bd);
            if idx < palette.len() {
                let (r, g, b) = palette[idx];
                let a = if idx < trns.len() { trns[idx] } else { 255 };
                (r, g, b, a)
            } else {
                (0, 0, 0, 255)
            }
        }
        PngColorType::GrayscaleAlpha => {
            let bytes_per_sample = if bd == 16 { 2 } else { 1 };
            let off = x * 2 * bytes_per_sample;
            let (v, a) = if bd == 16 {
                if off + 3 < row.len() {
                    (row[off], row[off + 2])
                } else {
                    (0, 0)
                }
            } else if off + 1 < row.len() {
                (row[off], row[off + 1])
            } else {
                (0, 0)
            };
            (v, v, v, a)
        }
        PngColorType::Rgba => {
            let bytes_per_sample = if bd == 16 { 2 } else { 1 };
            let off = x * 4 * bytes_per_sample;
            if bd == 16 {
                if off + 7 < row.len() {
                    (row[off], row[off + 2], row[off + 4], row[off + 6])
                } else {
                    (0, 0, 0, 0)
                }
            } else if off + 3 < row.len() {
                (row[off], row[off + 1], row[off + 2], row[off + 3])
            } else {
                (0, 0, 0, 0)
            }
        }
    }
}

/// Extract a sub-byte sample value from packed scanline data.
fn extract_sample(row: &[u8], index: usize, bit_depth: usize) -> usize {
    match bit_depth {
        1 => {
            let byte_idx = index / 8;
            let bit_idx = 7 - (index % 8);
            if byte_idx < row.len() {
                ((row[byte_idx] >> bit_idx) & 1) as usize
            } else {
                0
            }
        }
        2 => {
            let byte_idx = index / 4;
            let shift = 6 - (index % 4) * 2;
            if byte_idx < row.len() {
                ((row[byte_idx] >> shift) & 3) as usize
            } else {
                0
            }
        }
        4 => {
            let byte_idx = index / 2;
            let shift = if index.is_multiple_of(2) { 4 } else { 0 };
            if byte_idx < row.len() {
                ((row[byte_idx] >> shift) & 0xF) as usize
            } else {
                0
            }
        }
        8 => {
            if index < row.len() {
                row[index] as usize
            } else {
                0
            }
        }
        16 => {
            let off = index * 2;
            if off + 1 < row.len() {
                read_be_u16(row, off) as usize
            } else {
                0
            }
        }
        _ => 0,
    }
}

/// Scale a sample value from its native bit depth to 8-bit.
fn scale_to_8bit(val: usize, bit_depth: usize) -> u8 {
    match bit_depth {
        1 => {
            if val != 0 {
                255
            } else {
                0
            }
        }
        2 => (val * 85) as u8,         // 0->0, 1->85, 2->170, 3->255
        4 => ((val * 255) / 15) as u8, // scale 0-15 to 0-255
        8 => val as u8,
        16 => (val >> 8) as u8,
        _ => val as u8,
    }
}

// ============================================================================
// DEFLATE / ZLIB DECOMPRESSION (RFC 1950 / 1951)
// ============================================================================

/// Decompress zlib-wrapped data (CMF + FLG + compressed blocks + Adler-32).
fn zlib_decompress(data: &[u8]) -> Result<Vec<u8>, ImageCodecError> {
    if data.len() < 6 {
        return Err(ImageCodecError::TruncatedData);
    }

    let cmf = data[0];
    let _flg = data[1];

    // CMF: bits 0-3 = CM (must be 8 for deflate), bits 4-7 = CINFO
    if cmf & 0x0F != 8 {
        return Err(ImageCodecError::Unsupported);
    }

    // Verify CMF/FLG check
    let check = (cmf as u16) * 256 + (_flg as u16);
    if !check.is_multiple_of(31) {
        return Err(ImageCodecError::CorruptData);
    }

    // Check FDICT flag (bit 5 of FLG) -- we don't support preset dictionaries
    if _flg & 0x20 != 0 {
        return Err(ImageCodecError::Unsupported);
    }

    // Decompress DEFLATE stream starting at offset 2
    let compressed = &data[2..];
    let output = deflate_decompress(compressed)?;

    // Verify Adler-32 checksum (last 4 bytes of zlib stream)
    if data.len() >= 6 {
        let adler_offset = data.len() - 4;
        let stored_adler = read_be_u32(data, adler_offset);
        let computed_adler = adler32(&output);
        if stored_adler != computed_adler {
            // Some PNG encoders produce valid images with wrong checksums;
            // we log but don't fail for robustness.
            // return Err(ImageCodecError::ChecksumMismatch);
        }
    }

    Ok(output)
}

/// Compute Adler-32 checksum.
fn adler32(data: &[u8]) -> u32 {
    let mut a: u32 = 1;
    let mut b: u32 = 0;

    for &byte in data {
        a = (a + byte as u32) % 65521;
        b = (b + a) % 65521;
    }

    (b << 16) | a
}

/// Bit reader for DEFLATE stream.
struct BitReader<'a> {
    data: &'a [u8],
    byte_pos: usize,
    bit_pos: u8, // 0-7, bits consumed in current byte
}

impl<'a> BitReader<'a> {
    fn new(data: &'a [u8]) -> Self {
        Self {
            data,
            byte_pos: 0,
            bit_pos: 0,
        }
    }

    /// Read `n` bits (up to 25), LSB first.
    fn read_bits(&mut self, n: u8) -> Result<u32, ImageCodecError> {
        let mut result: u32 = 0;
        let mut bits_read: u8 = 0;

        while bits_read < n {
            if self.byte_pos >= self.data.len() {
                return Err(ImageCodecError::TruncatedData);
            }

            let available = 8 - self.bit_pos;
            let needed = n - bits_read;
            let take = if available < needed {
                available
            } else {
                needed
            };
            let mask = (1u32 << take) - 1;
            let bits = ((self.data[self.byte_pos] >> self.bit_pos) as u32) & mask;
            result |= bits << bits_read;
            bits_read += take;
            self.bit_pos += take;

            if self.bit_pos >= 8 {
                self.bit_pos = 0;
                self.byte_pos += 1;
            }
        }

        Ok(result)
    }

    /// Read a single bit.
    fn read_bit(&mut self) -> Result<u32, ImageCodecError> {
        self.read_bits(1)
    }

    /// Align to byte boundary.
    fn align(&mut self) {
        if self.bit_pos > 0 {
            self.bit_pos = 0;
            self.byte_pos += 1;
        }
    }

    /// Read a byte (must be byte-aligned).
    fn read_byte(&mut self) -> Result<u8, ImageCodecError> {
        if self.byte_pos >= self.data.len() {
            return Err(ImageCodecError::TruncatedData);
        }
        let b = self.data[self.byte_pos];
        self.byte_pos += 1;
        Ok(b)
    }
}

/// DEFLATE decompression (RFC 1951).
fn deflate_decompress(data: &[u8]) -> Result<Vec<u8>, ImageCodecError> {
    let mut reader = BitReader::new(data);
    let mut output: Vec<u8> = Vec::new();

    loop {
        let bfinal = reader.read_bit()?;
        let btype = reader.read_bits(2)?;

        match btype {
            0 => {
                // Stored (uncompressed) block
                reader.align();
                let len = reader.read_byte()? as u16 | ((reader.read_byte()? as u16) << 8);
                let _nlen = reader.read_byte()? as u16 | ((reader.read_byte()? as u16) << 8);
                for _ in 0..len {
                    output.push(reader.read_byte()?);
                }
            }
            1 => {
                // Fixed Huffman codes
                let (lit_tree, dist_tree) = build_fixed_huffman_trees();
                inflate_block(&mut reader, &lit_tree, &dist_tree, &mut output)?;
            }
            2 => {
                // Dynamic Huffman codes
                let (lit_tree, dist_tree) = decode_dynamic_huffman(&mut reader)?;
                inflate_block(&mut reader, &lit_tree, &dist_tree, &mut output)?;
            }
            _ => return Err(ImageCodecError::CorruptData),
        }

        if bfinal != 0 {
            break;
        }
    }

    Ok(output)
}

/// A Huffman tree for DEFLATE decoding, stored as a lookup table.
/// Maximum code length in DEFLATE is 15 bits.
struct HuffmanTree {
    /// For each code length, the number of codes and starting values.
    /// Stored as (min_code, symbols) per bit length.
    /// We use a simple linear decode approach.
    counts: [u16; 16],
    symbols: Vec<u16>,
    min_codes: [u32; 16],
    max_codes: [i32; 16],
    offsets: [u16; 16],
}

impl HuffmanTree {
    /// Build a Huffman tree from a list of code lengths.
    fn from_lengths(lengths: &[u8]) -> Result<Self, ImageCodecError> {
        let mut counts = [0u16; 16];
        let mut max_len: usize = 0;

        // Count occurrences of each code length
        for &len in lengths {
            let l = len as usize;
            if l > 0 {
                if l > 15 {
                    return Err(ImageCodecError::InvalidHuffmanTable);
                }
                counts[l] += 1;
                if l > max_len {
                    max_len = l;
                }
            }
        }

        // Compute starting codes for each length
        let mut code: u32 = 0;
        let mut next_code = [0u32; 16];
        let mut min_codes = [0u32; 16];
        let mut max_codes = [-1i32; 16];
        let mut offsets = [0u16; 16];

        let mut offset: u16 = 0;
        for bits in 1..=max_len {
            code = (code + counts[bits - 1] as u32) << 1;
            next_code[bits] = code;
            min_codes[bits] = code;
            offsets[bits] = offset;
            if counts[bits] > 0 {
                max_codes[bits] = (code + counts[bits] as u32 - 1) as i32;
            }
            offset += counts[bits];
        }

        // Assign symbols
        let total_symbols = offset as usize;
        let mut symbols = alloc::vec![0u16; total_symbols];
        let mut symbol_idx = [0u16; 16];
        symbol_idx[1..16].copy_from_slice(&offsets[1..16]);

        for (sym, &len) in lengths.iter().enumerate() {
            let l = len as usize;
            if l > 0 && l < 16 {
                let idx = symbol_idx[l] as usize;
                if idx < symbols.len() {
                    symbols[idx] = sym as u16;
                    symbol_idx[l] += 1;
                }
            }
        }

        Ok(Self {
            counts,
            symbols,
            min_codes,
            max_codes,
            offsets,
        })
    }

    /// Decode one symbol from the bit stream.
    fn decode(&self, reader: &mut BitReader) -> Result<u16, ImageCodecError> {
        let mut code: u32 = 0;

        for bits in 1..16u8 {
            code = (code << 1) | reader.read_bit()?;
            let b = bits as usize;
            if self.max_codes[b] >= 0 && code <= self.max_codes[b] as u32 {
                let idx = self.offsets[b] as usize + (code - self.min_codes[b]) as usize;
                if idx < self.symbols.len() {
                    return Ok(self.symbols[idx]);
                }
            }
        }

        Err(ImageCodecError::InvalidHuffmanTable)
    }
}

/// Build the fixed Huffman trees for DEFLATE block type 1.
fn build_fixed_huffman_trees() -> (HuffmanTree, HuffmanTree) {
    // Literal/length: 0-143 => 8 bits, 144-255 => 9 bits, 256-279 => 7 bits,
    // 280-287 => 8 bits
    let mut lit_lengths = [0u8; 288];
    lit_lengths[0..=143].fill(8);
    lit_lengths[144..=255].fill(9);
    lit_lengths[256..=279].fill(7);
    lit_lengths[280..=287].fill(8);

    // Distance: all 32 codes are 5 bits
    let dist_lengths = [5u8; 32];

    (
        HuffmanTree::from_lengths(&lit_lengths).unwrap_or_else(|_| HuffmanTree {
            counts: [0; 16],
            symbols: Vec::new(),
            min_codes: [0; 16],
            max_codes: [-1; 16],
            offsets: [0; 16],
        }),
        HuffmanTree::from_lengths(&dist_lengths).unwrap_or_else(|_| HuffmanTree {
            counts: [0; 16],
            symbols: Vec::new(),
            min_codes: [0; 16],
            max_codes: [-1; 16],
            offsets: [0; 16],
        }),
    )
}

/// Decode dynamic Huffman trees from DEFLATE block type 2 header.
fn decode_dynamic_huffman(
    reader: &mut BitReader,
) -> Result<(HuffmanTree, HuffmanTree), ImageCodecError> {
    let hlit = reader.read_bits(5)? as usize + 257;
    let hdist = reader.read_bits(5)? as usize + 1;
    let hclen = reader.read_bits(4)? as usize + 4;

    // Code length alphabet order
    const CL_ORDER: [usize; 19] = [
        16, 17, 18, 0, 8, 7, 9, 6, 10, 5, 11, 4, 12, 3, 13, 2, 14, 1, 15,
    ];

    let mut cl_lengths = [0u8; 19];
    for i in 0..hclen {
        cl_lengths[CL_ORDER[i]] = reader.read_bits(3)? as u8;
    }

    let cl_tree = HuffmanTree::from_lengths(&cl_lengths)?;

    // Decode literal/length + distance code lengths
    let total = hlit + hdist;
    let mut lengths: Vec<u8> = Vec::with_capacity(total);

    while lengths.len() < total {
        let sym = cl_tree.decode(reader)?;
        match sym {
            0..=15 => {
                lengths.push(sym as u8);
            }
            16 => {
                // Repeat previous length 3-6 times
                let extra = reader.read_bits(2)? as usize + 3;
                let prev = if let Some(&last) = lengths.last() {
                    last
                } else {
                    0
                };
                for _ in 0..extra {
                    if lengths.len() < total {
                        lengths.push(prev);
                    }
                }
            }
            17 => {
                // Repeat 0 for 3-10 times
                let extra = reader.read_bits(3)? as usize + 3;
                for _ in 0..extra {
                    if lengths.len() < total {
                        lengths.push(0);
                    }
                }
            }
            18 => {
                // Repeat 0 for 11-138 times
                let extra = reader.read_bits(7)? as usize + 11;
                for _ in 0..extra {
                    if lengths.len() < total {
                        lengths.push(0);
                    }
                }
            }
            _ => return Err(ImageCodecError::InvalidHuffmanTable),
        }
    }

    let lit_tree = HuffmanTree::from_lengths(&lengths[..hlit])?;
    let dist_tree = HuffmanTree::from_lengths(&lengths[hlit..hlit + hdist])?;

    Ok((lit_tree, dist_tree))
}

/// Length base values and extra bits for codes 257-285.
const LENGTH_BASE: [u16; 29] = [
    3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 15, 17, 19, 23, 27, 31, 35, 43, 51, 59, 67, 83, 99, 115, 131,
    163, 195, 227, 258,
];
const LENGTH_EXTRA: [u8; 29] = [
    0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4, 5, 5, 5, 5, 0,
];

/// Distance base values and extra bits for codes 0-29.
const DIST_BASE: [u16; 30] = [
    1, 2, 3, 4, 5, 7, 9, 13, 17, 25, 33, 49, 65, 97, 129, 193, 257, 385, 513, 769, 1025, 1537,
    2049, 3073, 4097, 6145, 8193, 12289, 16385, 24577,
];
const DIST_EXTRA: [u8; 30] = [
    0, 0, 0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8, 9, 9, 10, 10, 11, 11, 12, 12, 13,
    13,
];

/// Inflate one DEFLATE block using the given Huffman trees.
fn inflate_block(
    reader: &mut BitReader,
    lit_tree: &HuffmanTree,
    dist_tree: &HuffmanTree,
    output: &mut Vec<u8>,
) -> Result<(), ImageCodecError> {
    loop {
        let sym = lit_tree.decode(reader)?;

        if sym < 256 {
            // Literal byte
            output.push(sym as u8);
        } else if sym == 256 {
            // End of block
            break;
        } else {
            // Length/distance pair
            let len_idx = (sym - 257) as usize;
            if len_idx >= LENGTH_BASE.len() {
                return Err(ImageCodecError::CorruptData);
            }
            let length =
                LENGTH_BASE[len_idx] as usize + reader.read_bits(LENGTH_EXTRA[len_idx])? as usize;

            let dist_sym = dist_tree.decode(reader)? as usize;
            if dist_sym >= DIST_BASE.len() {
                return Err(ImageCodecError::CorruptData);
            }
            let distance =
                DIST_BASE[dist_sym] as usize + reader.read_bits(DIST_EXTRA[dist_sym])? as usize;

            if distance > output.len() {
                return Err(ImageCodecError::CorruptData);
            }

            // Copy from back-reference (byte-by-byte for overlapping copies)
            let start = output.len() - distance;
            for i in 0..length {
                let b = output[start + (i % distance)];
                output.push(b);
            }
        }
    }

    Ok(())
}

// ============================================================================
// TESTS
// ============================================================================

#[cfg(test)]
mod tests {
    #[allow(unused_imports)]
    use alloc::vec;

    use super::*;

    #[test]
    fn test_png_signature_check() {
        let bad = [0u8; 8];
        assert_eq!(decode_png(&bad), Err(ImageCodecError::InvalidSignature));
    }

    #[test]
    fn test_png_too_short() {
        let data = [137, 80, 78, 71];
        assert_eq!(decode_png(&data), Err(ImageCodecError::InvalidSignature));
    }

    #[test]
    fn test_png_color_type_from_u8() {
        assert_eq!(PngColorType::from_u8(0), Ok(PngColorType::Grayscale));
        assert_eq!(PngColorType::from_u8(2), Ok(PngColorType::Rgb));
        assert_eq!(PngColorType::from_u8(3), Ok(PngColorType::Indexed));
        assert_eq!(PngColorType::from_u8(4), Ok(PngColorType::GrayscaleAlpha));
        assert_eq!(PngColorType::from_u8(6), Ok(PngColorType::Rgba));
        assert_eq!(PngColorType::from_u8(7), Err(ImageCodecError::Unsupported));
    }

    #[test]
    fn test_png_color_type_channels() {
        assert_eq!(PngColorType::Grayscale.channels(), 1);
        assert_eq!(PngColorType::Rgb.channels(), 3);
        assert_eq!(PngColorType::Indexed.channels(), 1);
        assert_eq!(PngColorType::GrayscaleAlpha.channels(), 2);
        assert_eq!(PngColorType::Rgba.channels(), 4);
    }

    #[test]
    fn test_png_unfilter_none() {
        let mut row = vec![1, 2, 3, 4];
        let prev = vec![0, 0, 0, 0];
        png_unfilter(0, &mut row, &prev, 1).unwrap();
        assert_eq!(row, vec![1, 2, 3, 4]);
    }

    #[test]
    fn test_png_unfilter_sub() {
        let mut row = vec![1, 2, 3, 4];
        let prev = vec![0, 0, 0, 0];
        png_unfilter(1, &mut row, &prev, 1).unwrap();
        // Sub: each byte += previous byte in row
        // [1, 1+2=3, 3+3=6, 6+4=10]
        assert_eq!(row, vec![1, 3, 6, 10]);
    }

    #[test]
    fn test_png_unfilter_up() {
        let mut row = vec![1, 2, 3, 4];
        let prev = vec![10, 20, 30, 40];
        png_unfilter(2, &mut row, &prev, 1).unwrap();
        assert_eq!(row, vec![11, 22, 33, 44]);
    }

    #[test]
    fn test_png_unfilter_average() {
        let mut row = vec![0, 0, 0, 0];
        let prev = vec![10, 20, 30, 40];
        png_unfilter(3, &mut row, &prev, 1).unwrap();
        // Average: byte += floor((a + b) / 2) where a=left, b=above
        // [0+floor(0+10)/2=5, 0+floor(5+20)/2=12, 0+floor(12+30)/2=21,
        // 0+floor(21+40)/2=30]
        assert_eq!(row, vec![5, 12, 21, 30]);
    }

    #[test]
    fn test_png_unfilter_paeth() {
        let mut row = vec![10, 20, 30, 40];
        let prev = vec![0, 0, 0, 0];
        png_unfilter(4, &mut row, &prev, 1).unwrap();
        // Paeth with all-zero prev: a=left, b=0, c=0 => paeth=a (except first where
        // a=0) [10+paeth(0,0,0)=10, 20+paeth(10,0,0)=30, 30+paeth(30,0,0)=60,
        // 40+paeth(60,0,0)=100]
        assert_eq!(row, vec![10, 30, 60, 100]);
    }

    #[test]
    fn test_paeth_predictor_basic() {
        // When a=0, b=0, c=0 => p=0, pa=0, pb=0, pc=0 => returns a=0
        assert_eq!(paeth_predictor(0, 0, 0), 0);
        // When a=10, b=20, c=5 => p=25, pa=15, pb=5, pc=20 => returns b=20
        assert_eq!(paeth_predictor(10, 20, 5), 20);
    }

    #[test]
    fn test_extract_sample_1bit() {
        let row = vec![0b10110100];
        assert_eq!(extract_sample(&row, 0, 1), 1); // bit 7
        assert_eq!(extract_sample(&row, 1, 1), 0); // bit 6
        assert_eq!(extract_sample(&row, 2, 1), 1); // bit 5
        assert_eq!(extract_sample(&row, 3, 1), 1); // bit 4
        assert_eq!(extract_sample(&row, 4, 1), 0); // bit 3
        assert_eq!(extract_sample(&row, 5, 1), 1); // bit 2
    }

    #[test]
    fn test_extract_sample_4bit() {
        let row = vec![0xAB, 0xCD];
        assert_eq!(extract_sample(&row, 0, 4), 0xA);
        assert_eq!(extract_sample(&row, 1, 4), 0xB);
        assert_eq!(extract_sample(&row, 2, 4), 0xC);
        assert_eq!(extract_sample(&row, 3, 4), 0xD);
    }

    #[test]
    fn test_scale_to_8bit() {
        assert_eq!(scale_to_8bit(0, 1), 0);
        assert_eq!(scale_to_8bit(1, 1), 255);
        assert_eq!(scale_to_8bit(0, 2), 0);
        assert_eq!(scale_to_8bit(3, 2), 255);
        assert_eq!(scale_to_8bit(128, 8), 128);
        assert_eq!(scale_to_8bit(0xFF00, 16), 255);
    }

    #[test]
    fn test_png_unfilter_invalid_filter() {
        let mut row = vec![1, 2, 3];
        let prev = vec![0, 0, 0];
        assert_eq!(
            png_unfilter(5, &mut row, &prev, 1),
            Err(ImageCodecError::CorruptData)
        );
    }

    // -----------------------------------------------------------------------
    // Adler-32 / zlib tests
    // -----------------------------------------------------------------------

    #[test]
    fn test_adler32_empty() {
        assert_eq!(adler32(&[]), 0x00000001);
    }

    #[test]
    fn test_adler32_known() {
        // adler32("Wikipedia") = 0x11E60398
        let data = b"Wikipedia";
        assert_eq!(adler32(data), 0x11E60398);
    }

    // -----------------------------------------------------------------------
    // DEFLATE / Huffman tests
    // -----------------------------------------------------------------------

    #[test]
    fn test_huffman_tree_from_lengths() {
        // Simple tree: symbols 0,1 with lengths 1,1 => codes 0,1
        let lengths = [1u8, 1];
        let tree = HuffmanTree::from_lengths(&lengths).unwrap();
        assert_eq!(tree.counts[1], 2);
        assert_eq!(tree.symbols.len(), 2);
    }

    #[test]
    fn test_deflate_stored_block() {
        // Construct a minimal stored block: BFINAL=1, BTYPE=00, LEN=3, NLEN=~3,
        // data="abc"
        let mut block = Vec::new();
        block.push(0x01); // BFINAL=1, BTYPE=00 (stored) => bits: 1 00 = 0b001
                          // LEN = 3 (LE)
        block.push(0x03);
        block.push(0x00);
        // NLEN = !3 = 0xFFFC (LE)
        block.push(0xFC);
        block.push(0xFF);
        // Data
        block.push(b'a');
        block.push(b'b');
        block.push(b'c');

        let result = deflate_decompress(&block).unwrap();
        assert_eq!(result, b"abc");
    }

    #[test]
    fn test_bit_reader_basic() {
        let data = [0b10110100, 0xFF];
        let mut reader = BitReader::new(&data);
        // Read 4 bits LSB first from 0b10110100 => bits: 0,0,1,0 => 0b0100 = 4
        let v = reader.read_bits(4).unwrap();
        assert_eq!(v, 0b0100);
    }
}