//! JPEG decoder (JFIF baseline DCT, ITU-T T.81).
//!
//! Decodes baseline JPEG images (SOF0) into RGBA8 pixel data. Supports
//! 4:4:4, 4:2:2, and 4:2:0 chroma subsampling, Huffman entropy coding,
//! and restart markers (DRI/RST).

use crate::pixel::{Image, ImageError};

// ---------------------------------------------------------------------------
// JPEG marker constants (second byte after 0xFF prefix)
// ---------------------------------------------------------------------------

const MARKER_SOI: u8 = 0xD8;
const MARKER_EOI: u8 = 0xD9;
const MARKER_SOS: u8 = 0xDA;
const MARKER_DQT: u8 = 0xDB;
const MARKER_DHT: u8 = 0xC4;
const MARKER_SOF0: u8 = 0xC0;
const MARKER_DRI: u8 = 0xDD;

fn decode_err(msg: &str) -> ImageError {
    ImageError::Decode(msg.to_string())
}

// ---------------------------------------------------------------------------
// Zigzag order table
// ---------------------------------------------------------------------------

/// Maps zigzag index (0..63) to natural 8x8 row-major index.
const ZIGZAG: [usize; 64] = [
    0, 1, 8, 16, 9, 2, 3, 10, 17, 24, 32, 25, 18, 11, 4, 5, 12, 19, 26, 33, 40, 48, 41, 34, 27, 20,
    13, 6, 7, 14, 21, 28, 35, 42, 49, 56, 57, 50, 43, 36, 29, 22, 15, 23, 30, 37, 44, 51, 58, 59,
    52, 45, 38, 31, 39, 46, 53, 60, 61, 54, 47, 55, 62, 63,
];

// ---------------------------------------------------------------------------
// Byte reader
// ---------------------------------------------------------------------------

struct JpegReader<'a> {
    data: &'a [u8],
    pos: usize,
}

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

    fn remaining(&self) -> usize {
        self.data.len().saturating_sub(self.pos)
    }

    fn read_byte(&mut self) -> Result<u8, ImageError> {
        if self.pos >= self.data.len() {
            return Err(decode_err("unexpected end of JPEG data"));
        }
        let b = self.data[self.pos];
        self.pos += 1;
        Ok(b)
    }

    fn read_u16_be(&mut self) -> Result<u16, ImageError> {
        let hi = self.read_byte()? as u16;
        let lo = self.read_byte()? as u16;
        Ok((hi << 8) | lo)
    }

    fn read_bytes(&mut self, n: usize) -> Result<&'a [u8], ImageError> {
        if self.pos + n > self.data.len() {
            return Err(decode_err("unexpected end of JPEG data"));
        }
        let slice = &self.data[self.pos..self.pos + n];
        self.pos += n;
        Ok(slice)
    }

    fn skip(&mut self, n: usize) -> Result<(), ImageError> {
        if self.pos + n > self.data.len() {
            return Err(decode_err("unexpected end of JPEG data"));
        }
        self.pos += n;
        Ok(())
    }
}

// ---------------------------------------------------------------------------
// Quantization table
// ---------------------------------------------------------------------------

struct QuantTable {
    /// 64 quantization values in zigzag order.
    values: [u16; 64],
}

fn parse_dqt(
    reader: &mut JpegReader,
    tables: &mut [Option<QuantTable>; 4],
) -> Result<(), ImageError> {
    let length = reader.read_u16_be()? as usize;
    if length < 2 {
        return Err(decode_err("DQT: invalid length"));
    }
    let mut remaining = length - 2;

    while remaining > 0 {
        let pq_tq = reader.read_byte()?;
        remaining -= 1;
        let precision = pq_tq >> 4;
        let table_id = (pq_tq & 0x0F) as usize;
        if table_id >= 4 {
            return Err(decode_err("DQT: table id out of range"));
        }

        let mut values = [0u16; 64];
        if precision == 0 {
            // 8-bit values
            if remaining < 64 {
                return Err(decode_err("DQT: truncated 8-bit table"));
            }
            for v in &mut values {
                *v = reader.read_byte()? as u16;
            }
            remaining -= 64;
        } else if precision == 1 {
            // 16-bit values
            if remaining < 128 {
                return Err(decode_err("DQT: truncated 16-bit table"));
            }
            for v in &mut values {
                *v = reader.read_u16_be()?;
            }
            remaining -= 128;
        } else {
            return Err(decode_err("DQT: unsupported precision"));
        }

        tables[table_id] = Some(QuantTable { values });
    }
    Ok(())
}

// ---------------------------------------------------------------------------
// Huffman table
// ---------------------------------------------------------------------------

struct HuffTable {
    /// Number of codes of each length (1..=16).
    counts: [u8; 16],
    /// Symbol values in order of increasing code length.
    symbols: Vec<u8>,
    /// Index into symbols for the first code of each length.
    val_offset: [i32; 16],
    /// Maximum code value for each length (-1 if no codes).
    max_code: [i32; 16],
}

impl HuffTable {
    fn build(counts: [u8; 16], symbols: Vec<u8>) -> Self {
        let mut max_code = [-1i32; 16];
        let mut val_offset = [0i32; 16];

        let mut code = 0i32;
        let mut si = 0i32;
        for i in 0..16 {
            if counts[i] > 0 {
                val_offset[i] = si - code;
                code += counts[i] as i32;
                max_code[i] = code - 1;
                si += counts[i] as i32;
            }
            code <<= 1;
        }

        Self {
            counts,
            symbols,
            val_offset,
            max_code,
        }
    }
}

fn parse_dht(
    reader: &mut JpegReader,
    dc_tables: &mut [Option<HuffTable>; 4],
    ac_tables: &mut [Option<HuffTable>; 4],
) -> Result<(), ImageError> {
    let length = reader.read_u16_be()? as usize;
    if length < 2 {
        return Err(decode_err("DHT: invalid length"));
    }
    let mut remaining = length - 2;

    while remaining > 0 {
        let tc_th = reader.read_byte()?;
        remaining -= 1;
        let table_class = tc_th >> 4; // 0=DC, 1=AC
        let table_id = (tc_th & 0x0F) as usize;
        if table_class > 1 || table_id >= 4 {
            return Err(decode_err("DHT: invalid class or table id"));
        }

        if remaining < 16 {
            return Err(decode_err("DHT: truncated counts"));
        }
        let mut counts = [0u8; 16];
        for c in &mut counts {
            *c = reader.read_byte()?;
        }
        remaining -= 16;

        let total: usize = counts.iter().map(|&c| c as usize).sum();
        if remaining < total {
            return Err(decode_err("DHT: truncated symbols"));
        }
        let symbols = reader.read_bytes(total)?.to_vec();
        remaining -= total;

        let table = HuffTable::build(counts, symbols);
        if table_class == 0 {
            dc_tables[table_id] = Some(table);
        } else {
            ac_tables[table_id] = Some(table);
        }
    }
    Ok(())
}

// ---------------------------------------------------------------------------
// Frame header (SOF0)
// ---------------------------------------------------------------------------

struct ComponentInfo {
    id: u8,
    h_sample: u8,
    v_sample: u8,
    quant_table_id: u8,
}

struct FrameHeader {
    height: u16,
    width: u16,
    components: Vec<ComponentInfo>,
    h_max: u8,
    v_max: u8,
}

fn parse_sof0(reader: &mut JpegReader) -> Result<FrameHeader, ImageError> {
    let length = reader.read_u16_be()? as usize;
    if length < 8 {
        return Err(decode_err("SOF0: invalid length"));
    }
    let precision = reader.read_byte()?;
    if precision != 8 {
        return Err(decode_err("SOF0: only 8-bit precision supported"));
    }
    let height = reader.read_u16_be()?;
    let width = reader.read_u16_be()?;
    let num_components = reader.read_byte()? as usize;
    if num_components == 0 || num_components > 4 {
        return Err(decode_err("SOF0: invalid number of components"));
    }
    if length != 8 + 3 * num_components {
        return Err(decode_err("SOF0: length mismatch"));
    }

    let mut components = Vec::with_capacity(num_components);
    let mut h_max = 1u8;
    let mut v_max = 1u8;
    for _ in 0..num_components {
        let id = reader.read_byte()?;
        let hv = reader.read_byte()?;
        let h_sample = hv >> 4;
        let v_sample = hv & 0x0F;
        if h_sample == 0 || h_sample > 4 || v_sample == 0 || v_sample > 4 {
            return Err(decode_err("SOF0: invalid sampling factor"));
        }
        let quant_table_id = reader.read_byte()?;
        h_max = h_max.max(h_sample);
        v_max = v_max.max(v_sample);
        components.push(ComponentInfo {
            id,
            h_sample,
            v_sample,
            quant_table_id,
        });
    }

    if width == 0 || height == 0 {
        return Err(decode_err("SOF0: zero dimension"));
    }

    Ok(FrameHeader {
        height,
        width,
        components,
        h_max,
        v_max,
    })
}

// ---------------------------------------------------------------------------
// Scan header (SOS)
// ---------------------------------------------------------------------------

struct ScanComponentSelector {
    component_index: usize,
    dc_table_id: u8,
    ac_table_id: u8,
}

struct ScanHeader {
    components: Vec<ScanComponentSelector>,
}

fn parse_sos(reader: &mut JpegReader, frame: &FrameHeader) -> Result<ScanHeader, ImageError> {
    let length = reader.read_u16_be()? as usize;
    let num_components = reader.read_byte()? as usize;
    if num_components == 0 || num_components > 4 {
        return Err(decode_err("SOS: invalid number of components"));
    }
    if length != 6 + 2 * num_components {
        return Err(decode_err("SOS: length mismatch"));
    }

    let mut components = Vec::with_capacity(num_components);
    for _ in 0..num_components {
        let cs = reader.read_byte()?;
        let td_ta = reader.read_byte()?;
        let dc_table_id = td_ta >> 4;
        let ac_table_id = td_ta & 0x0F;

        // Find component index by id
        let component_index = frame
            .components
            .iter()
            .position(|c| c.id == cs)
            .ok_or_else(|| decode_err("SOS: unknown component id"))?;

        components.push(ScanComponentSelector {
            component_index,
            dc_table_id,
            ac_table_id,
        });
    }

    // Spectral selection and successive approximation (must be 0, 63, 0 for baseline)
    let _ss = reader.read_byte()?;
    let _se = reader.read_byte()?;
    let _ah_al = reader.read_byte()?;

    Ok(ScanHeader { components })
}

// ---------------------------------------------------------------------------
// Bit reader for entropy-coded data (MSB-first, with byte stuffing)
// ---------------------------------------------------------------------------

struct BitReader<'a> {
    data: &'a [u8],
    pos: usize,
    bit_buf: u32,
    bits_in_buf: u8,
    /// Set when we encounter a real marker during reading.
    marker_found: Option<u8>,
}

impl<'a> BitReader<'a> {
    fn new(data: &'a [u8], start: usize) -> Self {
        Self {
            data,
            pos: start,
            bit_buf: 0,
            bits_in_buf: 0,
            marker_found: None,
        }
    }

    /// Fill the bit buffer with at least `need` bits.
    fn fill_bits(&mut self, need: u8) -> Result<(), ImageError> {
        while self.bits_in_buf < need {
            if self.pos >= self.data.len() {
                return Err(decode_err("JPEG: unexpected end in entropy data"));
            }
            let b = self.data[self.pos];
            self.pos += 1;

            if b == 0xFF {
                if self.pos >= self.data.len() {
                    return Err(decode_err("JPEG: unexpected end after 0xFF"));
                }
                let next = self.data[self.pos];
                self.pos += 1;
                if next == 0x00 {
                    // Byte stuffing: literal 0xFF
                    self.bit_buf = (self.bit_buf << 8) | 0xFF;
                    self.bits_in_buf += 8;
                } else {
                    // Real marker found
                    self.marker_found = Some(next);
                    // Pad with zeros
                    self.bit_buf <<= 8;
                    self.bits_in_buf += 8;
                }
            } else {
                self.bit_buf = (self.bit_buf << 8) | (b as u32);
                self.bits_in_buf += 8;
            }
        }
        Ok(())
    }

    fn read_bits(&mut self, count: u8) -> Result<u16, ImageError> {
        if count == 0 {
            return Ok(0);
        }
        self.fill_bits(count)?;
        self.bits_in_buf -= count;
        let val = (self.bit_buf >> self.bits_in_buf) & ((1 << count) - 1);
        Ok(val as u16)
    }

    fn align_to_byte(&mut self) {
        self.bits_in_buf = 0;
        self.bit_buf = 0;
    }

    fn position(&self) -> usize {
        self.pos
    }
}

// ---------------------------------------------------------------------------
// Huffman decoding
// ---------------------------------------------------------------------------

fn huff_decode(reader: &mut BitReader, table: &HuffTable) -> Result<u8, ImageError> {
    let mut code = 0i32;
    for i in 0..16 {
        code = (code << 1) | reader.read_bits(1)? as i32;
        if table.counts[i] > 0 && code <= table.max_code[i] {
            let idx = (table.val_offset[i] + code) as usize;
            if idx >= table.symbols.len() {
                return Err(decode_err("Huffman: symbol index out of range"));
            }
            return Ok(table.symbols[idx]);
        }
    }
    Err(decode_err("Huffman: invalid code"))
}

/// Extend a value to a signed integer based on the JPEG sign convention.
fn extend(value: u16, bits: u8) -> i32 {
    if bits == 0 {
        return 0;
    }
    let vt = 1i32 << (bits - 1);
    let v = value as i32;
    if v < vt {
        v + (-1 << bits) + 1
    } else {
        v
    }
}

fn decode_dc(reader: &mut BitReader, table: &HuffTable) -> Result<i32, ImageError> {
    let category = huff_decode(reader, table)?;
    if category == 0 {
        return Ok(0);
    }
    if category > 15 {
        return Err(decode_err("DC: invalid category"));
    }
    let bits = reader.read_bits(category)?;
    Ok(extend(bits, category))
}

fn decode_ac(
    reader: &mut BitReader,
    table: &HuffTable,
    block: &mut [i32; 64],
) -> Result<(), ImageError> {
    let mut k = 1usize;
    while k < 64 {
        let rs = huff_decode(reader, table)?;
        let run = (rs >> 4) as usize;
        let category = rs & 0x0F;

        if category == 0 {
            if run == 0 {
                // EOB: fill rest with zeros
                while k < 64 {
                    block[k] = 0;
                    k += 1;
                }
                return Ok(());
            } else if run == 0x0F {
                // ZRL: 16 zeros
                for _ in 0..16 {
                    if k < 64 {
                        block[k] = 0;
                        k += 1;
                    }
                }
                continue;
            } else {
                return Err(decode_err("AC: invalid run/category combination"));
            }
        }

        // Skip `run` zeros
        for _ in 0..run {
            if k < 64 {
                block[k] = 0;
                k += 1;
            }
        }

        if k >= 64 {
            return Err(decode_err("AC: coefficient index out of range"));
        }

        let bits = reader.read_bits(category)?;
        block[k] = extend(bits, category);
        k += 1;
    }
    Ok(())
}

// ---------------------------------------------------------------------------
// Dequantization
// ---------------------------------------------------------------------------

fn dequantize(block: &mut [i32; 64], quant: &QuantTable) {
    for (b, &q) in block.iter_mut().zip(quant.values.iter()) {
        *b *= q as i32;
    }
}

// ---------------------------------------------------------------------------
// Inverse DCT
// ---------------------------------------------------------------------------

/// Reorder from zigzag to natural 8x8 row-major order.
fn unzigzag(zigzag: &[i32; 64]) -> [i32; 64] {
    let mut natural = [0i32; 64];
    for i in 0..64 {
        natural[ZIGZAG[i]] = zigzag[i];
    }
    natural
}

/// 1D IDCT on 8 values using the Loeffler/Ligtenberg/Moschytz algorithm.
/// Fixed-point with 12 bits of fractional precision.
///
/// Constants are scaled: C_k = cos(k*pi/16) * 2^12, rounded.
const FIX_0_298: i32 = 2446; // cos(7*pi/16) * 4096
const FIX_0_390: i32 = 3196; // sqrt(2) * (cos(6*pi/16) - cos(2*pi/16)) * 2048 -- see below
const FIX_0_541: i32 = 4433; // sqrt(2) * cos(6*pi/16) * 4096
const FIX_0_765: i32 = 6270; // sqrt(2) * cos(2*pi/16) - sqrt(2) * cos(6*pi/16)
const FIX_1_175: i32 = 9633; // sqrt(2) * cos(pi/8) -- used in stage 1 butterfly
const FIX_1_501: i32 = 12299; // sqrt(2) * (cos(pi/16) - cos(7*pi/16))
const FIX_1_847: i32 = 15137; // sqrt(2) * cos(3*pi/16)
const FIX_1_961: i32 = 16069; // sqrt(2) * (cos(3*pi/16) + cos(5*pi/16)) -- negative
const FIX_2_053: i32 = 16819; // sqrt(2) * (cos(pi/16) + cos(7*pi/16))
const FIX_2_562: i32 = 20995; // sqrt(2) * (cos(3*pi/16) - cos(5*pi/16)) -- negative
const FIX_3_072: i32 = 25172; // sqrt(2) * (cos(pi/16) + cos(3*pi/16))

const CONST_BITS: i32 = 13;
const PASS1_BITS: i32 = 2;

/// Perform the 2D IDCT and produce 64 pixel values (clamped to 0..255).
/// Input is in natural 8x8 row-major order, dequantized.
fn idct_block(coeffs: &[i32; 64]) -> [u8; 64] {
    let mut workspace = [0i32; 64];

    // Pass 1: process columns from input, store into workspace.
    for col in 0..8 {
        // If all AC terms are zero, short-circuit.
        if coeffs[col + 8] == 0
            && coeffs[col + 16] == 0
            && coeffs[col + 24] == 0
            && coeffs[col + 32] == 0
            && coeffs[col + 40] == 0
            && coeffs[col + 48] == 0
            && coeffs[col + 56] == 0
        {
            let dcval = coeffs[col] << PASS1_BITS;
            for row in 0..8 {
                workspace[row * 8 + col] = dcval;
            }
            continue;
        }

        // Even part: use the Loeffler method
        let z2 = coeffs[col + 16];
        let z3 = coeffs[col + 48];

        let z1 = (z2 + z3) * FIX_0_541;
        let tmp2 = z1 + z3 * (-FIX_1_847);
        let tmp3 = z1 + z2 * FIX_0_765;

        let z2 = coeffs[col];
        let z3 = coeffs[col + 32];

        let tmp0 = (z2 + z3) << CONST_BITS;
        let tmp1 = (z2 - z3) << CONST_BITS;

        let tmp10 = tmp0 + tmp3;
        let tmp13 = tmp0 - tmp3;
        let tmp11 = tmp1 + tmp2;
        let tmp12 = tmp1 - tmp2;

        // Odd part
        let tmp0 = coeffs[col + 56];
        let tmp1 = coeffs[col + 40];
        let tmp2 = coeffs[col + 24];
        let tmp3 = coeffs[col + 8];

        let z1 = tmp0 + tmp3;
        let z2 = tmp1 + tmp2;
        let z3 = tmp0 + tmp2;
        let z4 = tmp1 + tmp3;
        let z5 = (z3 + z4) * FIX_1_175;

        let tmp0 = tmp0 * FIX_0_298;
        let tmp1 = tmp1 * FIX_2_053;
        let tmp2 = tmp2 * FIX_3_072;
        let tmp3 = tmp3 * FIX_1_501;
        let z1 = z1 * (-FIX_0_390);
        let z2 = z2 * (-FIX_2_562);
        let z3 = z3 * (-FIX_1_961);
        let z4 = z4 * (-FIX_0_298);

        let z3 = z3 + z5;
        let z4 = z4 + z5;

        let tmp0 = tmp0 + z1 + z3;
        let tmp1 = tmp1 + z2 + z4;
        let tmp2 = tmp2 + z2 + z3;
        let tmp3 = tmp3 + z1 + z4;

        let shift = CONST_BITS - PASS1_BITS;
        workspace[col] = (tmp10 + tmp3 + (1 << (shift - 1))) >> shift;
        workspace[col + 56] = (tmp10 - tmp3 + (1 << (shift - 1))) >> shift;
        workspace[col + 8] = (tmp11 + tmp2 + (1 << (shift - 1))) >> shift;
        workspace[col + 48] = (tmp11 - tmp2 + (1 << (shift - 1))) >> shift;
        workspace[col + 16] = (tmp12 + tmp1 + (1 << (shift - 1))) >> shift;
        workspace[col + 40] = (tmp12 - tmp1 + (1 << (shift - 1))) >> shift;
        workspace[col + 24] = (tmp13 + tmp0 + (1 << (shift - 1))) >> shift;
        workspace[col + 32] = (tmp13 - tmp0 + (1 << (shift - 1))) >> shift;
    }

    // Pass 2: process rows from workspace, produce output.
    let mut output = [0u8; 64];
    for row in 0..8 {
        let base = row * 8;

        // Short-circuit for all-zero AC
        if workspace[base + 1] == 0
            && workspace[base + 2] == 0
            && workspace[base + 3] == 0
            && workspace[base + 4] == 0
            && workspace[base + 5] == 0
            && workspace[base + 6] == 0
            && workspace[base + 7] == 0
        {
            let dcval = clamp_to_u8(
                ((workspace[base] + (1 << (PASS1_BITS + 2))) >> (PASS1_BITS + 3)) + 128,
            );
            for col in 0..8 {
                output[base + col] = dcval;
            }
            continue;
        }

        let z2 = workspace[base + 2];
        let z3 = workspace[base + 6];

        let z1 = (z2 + z3) * FIX_0_541;
        let tmp2 = z1 + z3 * (-FIX_1_847);
        let tmp3 = z1 + z2 * FIX_0_765;

        let z2 = workspace[base];
        let z3 = workspace[base + 4];

        let tmp0 = (z2 + z3) << CONST_BITS;
        let tmp1 = (z2 - z3) << CONST_BITS;

        let tmp10 = tmp0 + tmp3;
        let tmp13 = tmp0 - tmp3;
        let tmp11 = tmp1 + tmp2;
        let tmp12 = tmp1 - tmp2;

        let tmp0 = workspace[base + 7];
        let tmp1 = workspace[base + 5];
        let tmp2 = workspace[base + 3];
        let tmp3 = workspace[base + 1];

        let z1 = tmp0 + tmp3;
        let z2 = tmp1 + tmp2;
        let z3 = tmp0 + tmp2;
        let z4 = tmp1 + tmp3;
        let z5 = (z3 + z4) * FIX_1_175;

        let tmp0 = tmp0 * FIX_0_298;
        let tmp1 = tmp1 * FIX_2_053;
        let tmp2 = tmp2 * FIX_3_072;
        let tmp3 = tmp3 * FIX_1_501;
        let z1 = z1 * (-FIX_0_390);
        let z2 = z2 * (-FIX_2_562);
        let z3 = z3 * (-FIX_1_961);
        let z4 = z4 * (-FIX_0_298);

        let z3 = z3 + z5;
        let z4 = z4 + z5;

        let tmp0 = tmp0 + z1 + z3;
        let tmp1 = tmp1 + z2 + z4;
        let tmp2 = tmp2 + z2 + z3;
        let tmp3 = tmp3 + z1 + z4;

        let shift = CONST_BITS + PASS1_BITS + 3;
        let round = 1 << (shift - 1);
        output[base] = clamp_to_u8(((tmp10 + tmp3 + round) >> shift) + 128);
        output[base + 7] = clamp_to_u8(((tmp10 - tmp3 + round) >> shift) + 128);
        output[base + 1] = clamp_to_u8(((tmp11 + tmp2 + round) >> shift) + 128);
        output[base + 6] = clamp_to_u8(((tmp11 - tmp2 + round) >> shift) + 128);
        output[base + 2] = clamp_to_u8(((tmp12 + tmp1 + round) >> shift) + 128);
        output[base + 5] = clamp_to_u8(((tmp12 - tmp1 + round) >> shift) + 128);
        output[base + 3] = clamp_to_u8(((tmp13 + tmp0 + round) >> shift) + 128);
        output[base + 4] = clamp_to_u8(((tmp13 - tmp0 + round) >> shift) + 128);
    }

    output
}

fn clamp_to_u8(val: i32) -> u8 {
    val.clamp(0, 255) as u8
}

// ---------------------------------------------------------------------------
// MCU-based scan decoding
// ---------------------------------------------------------------------------

/// Decode all MCUs from entropy-coded data.
/// Returns per-component sample buffers (at component resolution).
fn decode_scan(
    reader: &mut BitReader,
    frame: &FrameHeader,
    scan: &ScanHeader,
    quant_tables: &[Option<QuantTable>; 4],
    dc_tables: &[Option<HuffTable>; 4],
    ac_tables: &[Option<HuffTable>; 4],
    restart_interval: u16,
) -> Result<Vec<Vec<u8>>, ImageError> {
    let h_max = frame.h_max as usize;
    let v_max = frame.v_max as usize;

    // MCU dimensions in pixels
    let mcu_w = h_max * 8;
    let mcu_h = v_max * 8;

    // Number of MCUs in each direction
    let mcus_x = (frame.width as usize).div_ceil(mcu_w);
    let mcus_y = (frame.height as usize).div_ceil(mcu_h);

    // Allocate per-component sample buffers
    let mut comp_buffers: Vec<Vec<u8>> = Vec::with_capacity(frame.components.len());
    let mut comp_widths: Vec<usize> = Vec::with_capacity(frame.components.len());
    let mut comp_heights: Vec<usize> = Vec::with_capacity(frame.components.len());

    for comp in &frame.components {
        let cw = mcus_x * comp.h_sample as usize * 8;
        let ch = mcus_y * comp.v_sample as usize * 8;
        comp_buffers.push(vec![0u8; cw * ch]);
        comp_widths.push(cw);
        comp_heights.push(ch);
    }

    // DC predictors per component
    let mut dc_pred = vec![0i32; frame.components.len()];

    let mut mcu_count = 0u32;
    let restart_interval = restart_interval as u32;

    for mcu_y in 0..mcus_y {
        for mcu_x in 0..mcus_x {
            // Handle restart marker
            if restart_interval > 0 && mcu_count > 0 && mcu_count.is_multiple_of(restart_interval) {
                reader.align_to_byte();
                // Skip past the restart marker
                // The marker may already have been found by the bit reader
                if reader.marker_found.is_some() {
                    reader.marker_found = None;
                }
                // Reset DC predictors
                dc_pred.fill(0);
            }

            for scan_comp in &scan.components {
                let ci = scan_comp.component_index;
                let comp = &frame.components[ci];

                let dc_table = dc_tables[scan_comp.dc_table_id as usize]
                    .as_ref()
                    .ok_or_else(|| decode_err("missing DC Huffman table"))?;
                let ac_table = ac_tables[scan_comp.ac_table_id as usize]
                    .as_ref()
                    .ok_or_else(|| decode_err("missing AC Huffman table"))?;
                let quant = quant_tables[comp.quant_table_id as usize]
                    .as_ref()
                    .ok_or_else(|| decode_err("missing quantization table"))?;

                // Each component contributes h_sample * v_sample blocks per MCU
                for bv in 0..comp.v_sample as usize {
                    for bh in 0..comp.h_sample as usize {
                        // Decode one 8x8 block
                        let mut block = [0i32; 64];

                        // DC
                        let dc_diff = decode_dc(reader, dc_table)?;
                        dc_pred[ci] += dc_diff;
                        block[0] = dc_pred[ci];

                        // AC
                        decode_ac(reader, ac_table, &mut block)?;

                        // Dequantize (in zigzag order)
                        dequantize(&mut block, quant);

                        // Unzigzag to natural order
                        let natural = unzigzag(&block);

                        // IDCT
                        let pixels = idct_block(&natural);

                        // Write block to component buffer
                        let block_x = mcu_x * comp.h_sample as usize * 8 + bh * 8;
                        let block_y = mcu_y * comp.v_sample as usize * 8 + bv * 8;
                        let cw = comp_widths[ci];

                        for row in 0..8 {
                            let dst_y = block_y + row;
                            if dst_y < comp_heights[ci] {
                                let dst_offset = dst_y * cw + block_x;
                                for col in 0..8 {
                                    if block_x + col < cw {
                                        comp_buffers[ci][dst_offset + col] = pixels[row * 8 + col];
                                    }
                                }
                            }
                        }
                    }
                }
            }
            mcu_count += 1;
        }
    }

    Ok(comp_buffers)
}

// ---------------------------------------------------------------------------
// Chroma upsampling
// ---------------------------------------------------------------------------

#[allow(clippy::too_many_arguments)]
fn upsample(
    samples: &[u8],
    sample_width: usize,
    sample_height: usize,
    h_sample: u8,
    v_sample: u8,
    h_max: u8,
    v_max: u8,
    target_width: usize,
    target_height: usize,
) -> Vec<u8> {
    let h_ratio = (h_max / h_sample) as usize;
    let v_ratio = (v_max / v_sample) as usize;

    if h_ratio == 1 && v_ratio == 1 {
        // No upsampling needed — just crop if needed
        if sample_width == target_width && sample_height == target_height {
            return samples.to_vec();
        }
        let mut out = vec![0u8; target_width * target_height];
        for y in 0..target_height {
            for x in 0..target_width {
                out[y * target_width + x] = samples[y * sample_width + x];
            }
        }
        return out;
    }

    let mut out = vec![0u8; target_width * target_height];
    for y in 0..target_height {
        let sy = (y / v_ratio).min(sample_height - 1);
        for x in 0..target_width {
            let sx = (x / h_ratio).min(sample_width - 1);
            out[y * target_width + x] = samples[sy * sample_width + sx];
        }
    }
    out
}

// ---------------------------------------------------------------------------
// YCbCr to RGB conversion
// ---------------------------------------------------------------------------

fn ycbcr_to_rgba(
    y_samples: &[u8],
    cb_samples: &[u8],
    cr_samples: &[u8],
    width: usize,
    height: usize,
) -> Vec<u8> {
    let mut rgba = Vec::with_capacity(width * height * 4);
    for i in 0..width * height {
        let y = y_samples[i] as i32;
        let cb = cb_samples[i] as i32 - 128;
        let cr = cr_samples[i] as i32 - 128;

        // Fixed-point YCbCr->RGB (BT.601)
        let r = y + ((cr * 359 + 128) >> 8);
        let g = y - ((cb * 88 + cr * 183 + 128) >> 8);
        let b = y + ((cb * 454 + 128) >> 8);

        rgba.push(r.clamp(0, 255) as u8);
        rgba.push(g.clamp(0, 255) as u8);
        rgba.push(b.clamp(0, 255) as u8);
        rgba.push(255);
    }
    rgba
}

// ---------------------------------------------------------------------------
// Public API
// ---------------------------------------------------------------------------

/// Decode a JPEG image into an RGBA8 `Image`.
///
/// Supports baseline DCT (SOF0) with Huffman coding, 4:4:4/4:2:2/4:2:0
/// chroma subsampling, and restart markers.
pub fn decode_jpeg(data: &[u8]) -> Result<Image, ImageError> {
    let mut reader = JpegReader::new(data);

    // Verify SOI
    if reader.remaining() < 2 {
        return Err(decode_err("JPEG: too short"));
    }
    let soi1 = reader.read_byte()?;
    let soi2 = reader.read_byte()?;
    if soi1 != 0xFF || soi2 != MARKER_SOI {
        return Err(decode_err("JPEG: missing SOI marker"));
    }

    let mut quant_tables: [Option<QuantTable>; 4] = [None, None, None, None];
    let mut dc_tables: [Option<HuffTable>; 4] = [None, None, None, None];
    let mut ac_tables: [Option<HuffTable>; 4] = [None, None, None, None];
    let mut frame: Option<FrameHeader> = None;
    let mut restart_interval: u16 = 0;
    let mut comp_buffers: Option<Vec<Vec<u8>>> = None;

    loop {
        // Find next marker
        let mut b = reader.read_byte()?;
        if b != 0xFF {
            // Sometimes there is padding; scan for 0xFF
            while b != 0xFF {
                if reader.remaining() == 0 {
                    return Err(decode_err("JPEG: unexpected end searching for marker"));
                }
                b = reader.read_byte()?;
            }
        }
        // Skip fill bytes (multiple 0xFF)
        let mut marker = reader.read_byte()?;
        while marker == 0xFF {
            marker = reader.read_byte()?;
        }
        if marker == 0x00 {
            continue; // Stuffed byte outside scan, ignore
        }

        match marker {
            MARKER_EOI => break,

            MARKER_SOF0 => {
                frame = Some(parse_sof0(&mut reader)?);
            }

            // Reject progressive/lossless/etc
            0xC1..=0xC3 | 0xC5..=0xC7 | 0xC9..=0xCB | 0xCD..=0xCF => {
                return Err(decode_err("JPEG: only baseline DCT (SOF0) is supported"));
            }

            MARKER_DHT => {
                parse_dht(&mut reader, &mut dc_tables, &mut ac_tables)?;
            }

            MARKER_DQT => {
                parse_dqt(&mut reader, &mut quant_tables)?;
            }

            MARKER_DRI => {
                let _len = reader.read_u16_be()?;
                restart_interval = reader.read_u16_be()?;
            }

            MARKER_SOS => {
                let f = frame
                    .as_ref()
                    .ok_or_else(|| decode_err("JPEG: SOS before SOF"))?;
                let scan = parse_sos(&mut reader, f)?;

                let mut bit_reader = BitReader::new(reader.data, reader.pos);
                comp_buffers = Some(decode_scan(
                    &mut bit_reader,
                    f,
                    &scan,
                    &quant_tables,
                    &dc_tables,
                    &ac_tables,
                    restart_interval,
                )?);
                reader.pos = bit_reader.position();
                // If the bit reader found a marker, back up so the outer loop sees it
                if let Some(m) = bit_reader.marker_found {
                    if m == MARKER_EOI {
                        break;
                    }
                    // Back up to the 0xFF before the marker
                    reader.pos -= 1;
                    // We need to also account for the 0xFF
                    if reader.pos > 0 {
                        reader.pos -= 1;
                    }
                }
            }

            // APP0..APP15, COM: skip
            _ => {
                if reader.remaining() >= 2 {
                    let len = reader.read_u16_be()? as usize;
                    if len >= 2 {
                        reader.skip(len - 2)?;
                    }
                }
            }
        }
    }

    let f = frame.ok_or_else(|| decode_err("JPEG: no SOF0 frame found"))?;
    let buffers = comp_buffers.ok_or_else(|| decode_err("JPEG: no scan data"))?;

    let width = f.width as usize;
    let height = f.height as usize;

    if f.components.len() == 1 {
        // Grayscale
        let h_max = f.h_max;
        let v_max = f.v_max;
        let comp = &f.components[0];
        let mcus_x = width.div_ceil(h_max as usize * 8);
        let mcus_y = height.div_ceil(v_max as usize * 8);
        let cw = mcus_x * comp.h_sample as usize * 8;
        let ch = mcus_y * comp.v_sample as usize * 8;

        let gray = upsample(
            &buffers[0],
            cw,
            ch,
            comp.h_sample,
            comp.v_sample,
            h_max,
            v_max,
            width,
            height,
        );
        crate::pixel::from_grayscale(width as u32, height as u32, &gray)
    } else if f.components.len() == 3 {
        // YCbCr
        let h_max = f.h_max;
        let v_max = f.v_max;

        let mut upsampled = Vec::with_capacity(3);
        for (ci, comp) in f.components.iter().enumerate() {
            let mcus_x = width.div_ceil(h_max as usize * 8);
            let mcus_y = height.div_ceil(v_max as usize * 8);
            let cw = mcus_x * comp.h_sample as usize * 8;
            let ch = mcus_y * comp.v_sample as usize * 8;

            upsampled.push(upsample(
                &buffers[ci],
                cw,
                ch,
                comp.h_sample,
                comp.v_sample,
                h_max,
                v_max,
                width,
                height,
            ));
        }

        let rgba = ycbcr_to_rgba(&upsampled[0], &upsampled[1], &upsampled[2], width, height);
        Image::new(width as u32, height as u32, rgba)
    } else {
        Err(decode_err("JPEG: unsupported number of components"))
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    // -- Zigzag table --

    #[test]
    fn zigzag_covers_all_indices() {
        let mut seen = [false; 64];
        for &idx in &ZIGZAG {
            assert!(idx < 64);
            seen[idx] = true;
        }
        assert!(seen.iter().all(|&s| s), "ZIGZAG must cover 0..63");
    }

    #[test]
    fn zigzag_known_positions() {
        assert_eq!(ZIGZAG[0], 0);
        assert_eq!(ZIGZAG[1], 1);
        assert_eq!(ZIGZAG[2], 8);
        assert_eq!(ZIGZAG[3], 16);
        assert_eq!(ZIGZAG[63], 63);
    }

    // -- Unzigzag --

    #[test]
    fn unzigzag_dc_only() {
        let mut zigzag = [0i32; 64];
        zigzag[0] = 100;
        let natural = unzigzag(&zigzag);
        assert_eq!(natural[0], 100);
        for i in 1..64 {
            assert_eq!(natural[i], 0);
        }
    }

    // -- Extend (sign extension) --

    #[test]
    fn extend_values() {
        // Category 1: values 0,1 -> -1, 1
        assert_eq!(extend(0, 1), -1);
        assert_eq!(extend(1, 1), 1);
        // Category 2: values 0,1,2,3 -> -3,-2,2,3
        assert_eq!(extend(0, 2), -3);
        assert_eq!(extend(1, 2), -2);
        assert_eq!(extend(2, 2), 2);
        assert_eq!(extend(3, 2), 3);
        // Category 0
        assert_eq!(extend(0, 0), 0);
    }

    // -- IDCT: DC only block --

    #[test]
    fn idct_dc_only() {
        let mut coeffs = [0i32; 64];
        coeffs[0] = 128; // DC coefficient
        let pixels = idct_block(&coeffs);
        // All pixels should be DC/8 + 128 = 128/8 + 128 = 144
        // (IDCT of a DC-only block divides by 8 in each dimension)
        let expected = 128 + (128 / 8);
        for &p in &pixels {
            // Allow +-1 for rounding
            assert!(
                (p as i32 - expected).unsigned_abs() <= 1,
                "expected ~{expected}, got {p}"
            );
        }
    }

    #[test]
    fn idct_zero_block() {
        let coeffs = [0i32; 64];
        let pixels = idct_block(&coeffs);
        // All zeros + level shift of 128
        for &p in &pixels {
            assert_eq!(p, 128);
        }
    }

    // -- YCbCr to RGB --

    #[test]
    fn ycbcr_white() {
        let rgba = ycbcr_to_rgba(&[255], &[128], &[128], 1, 1);
        assert_eq!(rgba[0], 255); // R
        assert_eq!(rgba[1], 255); // G
        assert_eq!(rgba[2], 255); // B
        assert_eq!(rgba[3], 255); // A
    }

    #[test]
    fn ycbcr_black() {
        let rgba = ycbcr_to_rgba(&[0], &[128], &[128], 1, 1);
        assert_eq!(rgba[0], 0); // R
        assert_eq!(rgba[1], 0); // G
        assert_eq!(rgba[2], 0); // B
        assert_eq!(rgba[3], 255); // A
    }

    #[test]
    fn ycbcr_gray_128() {
        let rgba = ycbcr_to_rgba(&[128], &[128], &[128], 1, 1);
        assert_eq!(rgba[0], 128); // R
        assert_eq!(rgba[1], 128); // G
        assert_eq!(rgba[2], 128); // B
    }

    // -- Huffman table build and decode --

    #[test]
    fn huffman_build_and_decode() {
        // Simple Huffman table: 2 symbols
        // Symbol A has code 0 (1 bit), symbol B has code 1 (1 bit)
        let mut counts = [0u8; 16];
        counts[0] = 2; // 2 codes of length 1
        let symbols = vec![b'A', b'B'];
        let table = HuffTable::build(counts, symbols);

        // Encode "ABA" as bits: 0, 1, 0 = 0b010_00000 = 0x40
        let data = [0x40u8];
        let mut reader = BitReader::new(&data, 0);
        assert_eq!(huff_decode(&mut reader, &table).unwrap(), b'A');
        assert_eq!(huff_decode(&mut reader, &table).unwrap(), b'B');
        assert_eq!(huff_decode(&mut reader, &table).unwrap(), b'A');
    }

    #[test]
    fn huffman_multi_length() {
        // 3 symbols: A=0 (1 bit), B=10 (2 bits), C=11 (2 bits)
        let mut counts = [0u8; 16];
        counts[0] = 1; // 1 code of length 1
        counts[1] = 2; // 2 codes of length 2
        let symbols = vec![b'A', b'B', b'C'];
        let table = HuffTable::build(counts, symbols);

        // "ABCA" = 0, 10, 11, 0 = 0b_0_10_11_0_00 = 0x58 (MSB first)
        let data = [0x58u8];
        let mut reader = BitReader::new(&data, 0);
        assert_eq!(huff_decode(&mut reader, &table).unwrap(), b'A');
        assert_eq!(huff_decode(&mut reader, &table).unwrap(), b'B');
        assert_eq!(huff_decode(&mut reader, &table).unwrap(), b'C');
        assert_eq!(huff_decode(&mut reader, &table).unwrap(), b'A');
    }

    // -- Bit reader byte stuffing --

    #[test]
    fn bit_reader_stuffing() {
        // 0xFF 0x00 should produce a literal 0xFF byte
        let data = [0xFF, 0x00, 0x80];
        let mut reader = BitReader::new(&data, 0);
        // Read 8 bits -> should get 0xFF
        let val = reader.read_bits(8).unwrap();
        assert_eq!(val, 0xFF);
        // Read 8 more bits -> 0x80
        let val = reader.read_bits(8).unwrap();
        assert_eq!(val, 0x80);
    }

    // -- Dequantize --

    #[test]
    fn dequantize_basic() {
        let mut block = [0i32; 64];
        block[0] = 10;
        block[1] = -5;
        let quant = QuantTable {
            values: {
                let mut v = [1u16; 64];
                v[0] = 16;
                v[1] = 11;
                v
            },
        };
        dequantize(&mut block, &quant);
        assert_eq!(block[0], 160);
        assert_eq!(block[1], -55);
    }

    // -- Upsampling --

    #[test]
    fn upsample_identity() {
        let samples = vec![1, 2, 3, 4];
        let result = upsample(&samples, 2, 2, 2, 2, 2, 2, 2, 2);
        assert_eq!(result, samples);
    }

    #[test]
    fn upsample_2x_horizontal() {
        // 1x2 upsampled to 2x2
        let samples = vec![10, 20];
        let result = upsample(&samples, 2, 1, 1, 1, 2, 1, 4, 1);
        assert_eq!(result, vec![10, 10, 20, 20]);
    }

    #[test]
    fn upsample_2x_both() {
        // 1x1 upsampled to 2x2
        let samples = vec![42];
        let result = upsample(&samples, 1, 1, 1, 1, 2, 2, 2, 2);
        assert_eq!(result, vec![42, 42, 42, 42]);
    }

    // -- Minimal JPEG construction and decoding --

    /// Build a minimal 1-component (grayscale) 8x8 JPEG.
    fn build_minimal_grayscale_jpeg(dc_value: i32) -> Vec<u8> {
        let mut out = Vec::new();

        // SOI
        out.push(0xFF);
        out.push(MARKER_SOI);

        // DQT: one table, id 0, all values = 1
        out.push(0xFF);
        out.push(MARKER_DQT);
        let dqt_len: u16 = 2 + 1 + 64;
        out.push((dqt_len >> 8) as u8);
        out.push(dqt_len as u8);
        out.push(0x00); // precision=0 (8-bit), table_id=0
        for _ in 0..64 {
            out.push(1); // All quantization values = 1
        }

        // SOF0: 1 component, 8x8
        out.push(0xFF);
        out.push(MARKER_SOF0);
        let sof_len: u16 = 2 + 1 + 2 + 2 + 1 + 3;
        out.push((sof_len >> 8) as u8);
        out.push(sof_len as u8);
        out.push(8); // precision
        out.push(0);
        out.push(8); // height=8
        out.push(0);
        out.push(8); // width=8
        out.push(1); // 1 component
        out.push(1); // component id=1
        out.push(0x11); // H=1, V=1
        out.push(0); // quant table 0

        // DHT: DC table 0
        // Simple: category 0 has code "00" (2 bits), category `cat` has code "01..." etc.
        // For minimal test: we only need DC category for our value.
        // Use standard JPEG luminance DC table (simplified).
        out.push(0xFF);
        out.push(MARKER_DHT);

        // DC table: we define a minimal table that can encode category 0 and a few others.
        // Cat 0: code 00 (2 bits) -> value 0 (zero diff)
        // Cat 1: code 010 (3 bits) -> 1 additional bit for value +-1
        // Cat 2: code 011 (3 bits) -> 2 additional bits
        // Cat 3: code 100 (3 bits) -> 3 additional bits
        // Cat 4: code 101 (3 bits) -> 4 additional bits
        // Cat 5: code 110 (3 bits) -> 5 additional bits
        // Cat 6: code 1110 (4 bits) -> 6 additional bits
        // Cat 7: code 11110 (5 bits) -> 7 additional bits
        // Cat 8: code 111110 (6 bits) -> 8 additional bits
        let dc_counts: [u8; 16] = [0, 1, 5, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let dc_symbols: &[u8] = &[0, 1, 2, 3, 4, 5, 6, 7, 8];
        let dc_total: usize = dc_counts.iter().map(|&c| c as usize).sum();
        let dht_dc_len = 2 + 1 + 16 + dc_total;
        out.push((dht_dc_len >> 8) as u8);
        out.push(dht_dc_len as u8);
        out.push(0x00); // DC table, id 0
        for &c in &dc_counts {
            out.push(c);
        }
        for &s in dc_symbols {
            out.push(s);
        }

        // AC table 0: minimal - only EOB (0x00)
        // EOB: code 0 (shortest possible). Use: 1 code of length 1 = symbol 0x00
        // But we also need ZRL potentially. For DC-only block, EOB is all we need.
        out.push(0xFF);
        out.push(MARKER_DHT);
        let ac_counts: [u8; 16] = [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let ac_symbols: &[u8] = &[0x00]; // EOB
        let ac_total: usize = ac_counts.iter().map(|&c| c as usize).sum();
        let dht_ac_len = 2 + 1 + 16 + ac_total;
        out.push((dht_ac_len >> 8) as u8);
        out.push(dht_ac_len as u8);
        out.push(0x10); // AC table, id 0
        for &c in &ac_counts {
            out.push(c);
        }
        for &s in ac_symbols {
            out.push(s);
        }

        // SOS
        out.push(0xFF);
        out.push(MARKER_SOS);
        let sos_len: u16 = 2 + 1 + 2 + 3;
        out.push((sos_len >> 8) as u8);
        out.push(sos_len as u8);
        out.push(1); // 1 component in scan
        out.push(1); // component selector = 1
        out.push(0x00); // DC table 0, AC table 0
        out.push(0); // Ss
        out.push(63); // Se
        out.push(0); // Ah/Al

        // Entropy-coded data: encode DC category + value, then AC EOB
        // DC: for dc_value, we need the category and then the magnitude bits
        let (cat, magnitude) = if dc_value == 0 {
            (0u8, 0u16)
        } else {
            let abs_val = dc_value.unsigned_abs() as u16;
            let cat = 16 - abs_val.leading_zeros() as u8;
            let mag = if dc_value > 0 {
                dc_value as u16
            } else {
                ((1u16 << cat) - 1) - abs_val
            };
            (cat, mag)
        };

        // Now encode: DC Huffman code for category, then `cat` magnitude bits, then AC EOB code
        // DC codes per our table:
        // cat 0: 00 (2 bits)
        // cat 1: 010 (3 bits)
        // cat 2: 011 (3 bits)
        // cat 3: 100 (3 bits)
        // cat 4: 101 (3 bits)
        // cat 5: 110 (3 bits)
        // cat 6: 1110 (4 bits)
        // cat 7: 11110 (5 bits)
        // cat 8: 111110 (6 bits)
        // AC EOB: 0 (1 bit)

        // Build the DC Huffman code for our category
        let dc_huff_table = HuffTable::build(dc_counts, dc_symbols.to_vec());
        let mut dc_code = 0u32;
        let mut dc_code_len = 0u8;
        {
            // Find the code for our category symbol
            let mut code = 0u32;
            let mut si = 0usize;
            for i in 0..16 {
                for _ in 0..dc_counts[i] {
                    if dc_huff_table.symbols[si] == cat {
                        dc_code = code;
                        dc_code_len = (i + 1) as u8;
                    }
                    code += 1;
                    si += 1;
                }
                code <<= 1;
            }
        }

        // Assemble bits: dc_code (dc_code_len bits) + magnitude (cat bits) + EOB (1 bit = 0)
        let mut bits = 0u64;
        let mut bp = 0u32;

        // Write dc_code MSB first
        bits |= (dc_code as u64) << (64 - dc_code_len as u32 - bp);
        bp += dc_code_len as u32;

        // Write magnitude bits MSB first
        if cat > 0 {
            bits |= (magnitude as u64) << (64 - cat as u32 - bp);
            bp += cat as u32;
        }

        // Write AC EOB = 0 (1 bit)
        // Already 0, just advance
        bp += 1;

        // Convert to bytes
        let byte_count = (bp + 7) / 8;
        for i in 0..byte_count {
            let b = ((bits >> (64 - 8 - i * 8)) & 0xFF) as u8;
            out.push(b);
            // Byte-stuff if needed
            if b == 0xFF {
                out.push(0x00);
            }
        }

        // EOI
        out.push(0xFF);
        out.push(MARKER_EOI);

        out
    }

    #[test]
    fn decode_grayscale_8x8_dc_zero() {
        let jpeg = build_minimal_grayscale_jpeg(0);
        let img = decode_jpeg(&jpeg).unwrap();
        assert_eq!(img.width, 8);
        assert_eq!(img.height, 8);
        // DC=0, quant=1, so coefficient is 0, IDCT of zero block = 128 everywhere
        for i in 0..64 {
            let r = img.data[i * 4];
            let g = img.data[i * 4 + 1];
            let b = img.data[i * 4 + 2];
            let a = img.data[i * 4 + 3];
            assert_eq!(r, 128, "pixel {i}: R should be 128, got {r}");
            assert_eq!(g, 128);
            assert_eq!(b, 128);
            assert_eq!(a, 255);
        }
    }

    #[test]
    fn decode_grayscale_8x8_dc_positive() {
        let jpeg = build_minimal_grayscale_jpeg(64);
        let img = decode_jpeg(&jpeg).unwrap();
        assert_eq!(img.width, 8);
        assert_eq!(img.height, 8);
        // DC=64, IDCT of DC-only -> 64/8 + 128 = 136
        let expected = 128 + 64 / 8;
        for i in 0..64 {
            let r = img.data[i * 4];
            assert!(
                (r as i32 - expected).unsigned_abs() <= 1,
                "pixel {i}: expected ~{expected}, got {r}"
            );
        }
    }

    #[test]
    fn decode_grayscale_8x8_dc_negative() {
        let jpeg = build_minimal_grayscale_jpeg(-64);
        let img = decode_jpeg(&jpeg).unwrap();
        assert_eq!(img.width, 8);
        assert_eq!(img.height, 8);
        // DC=-64, IDCT of DC-only -> -64/8 + 128 = 120
        let expected = 128 - 64 / 8;
        for i in 0..64 {
            let r = img.data[i * 4];
            assert!(
                (r as i32 - expected).unsigned_abs() <= 1,
                "pixel {i}: expected ~{expected}, got {r}"
            );
        }
    }

    // -- Error cases --

    #[test]
    fn error_missing_soi() {
        let data = [0x00, 0x00];
        assert!(decode_jpeg(&data).is_err());
    }

    #[test]
    fn error_too_short() {
        let data = [0xFF];
        assert!(decode_jpeg(&data).is_err());
    }

    #[test]
    fn error_no_frame() {
        // SOI then EOI with no frame
        let data = [0xFF, MARKER_SOI, 0xFF, MARKER_EOI];
        let err = decode_jpeg(&data).unwrap_err();
        assert!(matches!(err, ImageError::Decode(_)));
    }

    #[test]
    fn error_progressive_rejected() {
        let mut data = vec![0xFF, MARKER_SOI, 0xFF, 0xC2]; // SOF2 = progressive
                                                           // Minimal SOF2 header
        data.extend_from_slice(&[0, 11, 8, 0, 8, 0, 8, 1, 1, 0x11, 0]);
        data.push(0xFF);
        data.push(MARKER_EOI);
        let err = decode_jpeg(&data).unwrap_err();
        match err {
            ImageError::Decode(msg) => assert!(msg.contains("baseline"), "got: {msg}"),
            _ => panic!("expected Decode error"),
        }
    }

    // -- Quantization table parsing --

    #[test]
    fn parse_dqt_8bit() {
        let mut data = vec![0, 67]; // length = 67
        data.push(0x00); // precision=0, table_id=0
        for i in 0..64u8 {
            data.push(i + 1);
        }
        let mut reader = JpegReader::new(&data);
        let mut tables: [Option<QuantTable>; 4] = [None, None, None, None];
        parse_dqt(&mut reader, &mut tables).unwrap();
        assert!(tables[0].is_some());
        assert_eq!(tables[0].as_ref().unwrap().values[0], 1);
        assert_eq!(tables[0].as_ref().unwrap().values[63], 64);
    }

    // -- Huffman table parsing --

    #[test]
    fn parse_dht_roundtrip() {
        let mut data = Vec::new();
        let counts: [u8; 16] = [0, 1, 5, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let symbols: &[u8] = &[0, 1, 2, 3, 4, 5];
        let total: usize = counts.iter().map(|&c| c as usize).sum();
        let len = 2 + 1 + 16 + total;
        data.push((len >> 8) as u8);
        data.push(len as u8);
        data.push(0x00); // DC, id 0
        data.extend_from_slice(&counts);
        data.extend_from_slice(symbols);

        let mut reader = JpegReader::new(&data);
        let mut dc_tables: [Option<HuffTable>; 4] = [None, None, None, None];
        let mut ac_tables: [Option<HuffTable>; 4] = [None, None, None, None];
        parse_dht(&mut reader, &mut dc_tables, &mut ac_tables).unwrap();
        assert!(dc_tables[0].is_some());
        assert_eq!(dc_tables[0].as_ref().unwrap().symbols, symbols);
    }

    // -- Frame header parsing --

    #[test]
    fn parse_sof0_basic() {
        let mut data = Vec::new();
        let len: u16 = 8 + 3;
        data.push((len >> 8) as u8);
        data.push(len as u8);
        data.push(8); // precision
        data.push(0);
        data.push(16); // height=16
        data.push(0);
        data.push(16); // width=16
        data.push(1); // 1 component
        data.push(1); // id=1
        data.push(0x11); // H=1, V=1
        data.push(0); // quant table 0

        let mut reader = JpegReader::new(&data);
        let frame = parse_sof0(&mut reader).unwrap();
        assert_eq!(frame.width, 16);
        assert_eq!(frame.height, 16);
        assert_eq!(frame.components.len(), 1);
        assert_eq!(frame.h_max, 1);
        assert_eq!(frame.v_max, 1);
    }

    #[test]
    fn parse_sof0_three_components() {
        let mut data = Vec::new();
        let len: u16 = 8 + 9; // 3 components * 3 bytes each
        data.push((len >> 8) as u8);
        data.push(len as u8);
        data.push(8);
        data.push(0);
        data.push(32); // height=32
        data.push(0);
        data.push(32); // width=32
        data.push(3); // 3 components
                      // Y: H=2, V=2
        data.push(1);
        data.push(0x22);
        data.push(0);
        // Cb: H=1, V=1
        data.push(2);
        data.push(0x11);
        data.push(1);
        // Cr: H=1, V=1
        data.push(3);
        data.push(0x11);
        data.push(1);

        let mut reader = JpegReader::new(&data);
        let frame = parse_sof0(&mut reader).unwrap();
        assert_eq!(frame.width, 32);
        assert_eq!(frame.height, 32);
        assert_eq!(frame.components.len(), 3);
        assert_eq!(frame.h_max, 2);
        assert_eq!(frame.v_max, 2);
        assert_eq!(frame.components[0].h_sample, 2);
        assert_eq!(frame.components[0].v_sample, 2);
        assert_eq!(frame.components[1].h_sample, 1);
    }

    // -- Clamp --

    #[test]
    fn clamp_values() {
        assert_eq!(clamp_to_u8(-10), 0);
        assert_eq!(clamp_to_u8(0), 0);
        assert_eq!(clamp_to_u8(128), 128);
        assert_eq!(clamp_to_u8(255), 255);
        assert_eq!(clamp_to_u8(300), 255);
    }
}