chdlady-core 0.1.0

//! Hunk map table parsing and serialization for CHD v3/v4/v5.
use crate::crc16::crc16;
use crate::error::ChdError;
use crate::header::ChdHeader;
use chdlady_huffman::{BitReader, BitWriter, HuffmanDecoder, HuffmanEncoder};
use std::io::{Read, Seek, SeekFrom};

/// Type of hunk compression or reference in the container map.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum HunkType {
    /// Compressed using configured codec index (0 to 3).
    Compressed(u8),
    /// Stored uncompressed on disk.
    Uncompressed,
    /// Duplicate of an earlier hunk in this file.
    SelfRef(u64),
    /// Sourced from parent file.
    Parent(u64),
    /// Mini 8-byte fill pattern repeated across hunk (legacy v3/v4).
    Mini(u64),
}

/// Description of a single hunk in the container.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct MapEntry {
    /// Hunk storage or compression type.
    pub hunk_type: HunkType,
    /// Length in bytes of the block on disk.
    pub length: u32,
    /// File offset or referenced index.
    pub offset: u64,
    /// CRC-16 checksum of the hunk payload.
    pub crc16: u16,
}

const COMPRESSION_RLE_SMALL: u8 = 7;
const COMPRESSION_RLE_LARGE: u8 = 8;
const COMPRESSION_SELF_0: u8 = 9;
const COMPRESSION_SELF_1: u8 = 10;
const COMPRESSION_PARENT_SELF: u8 = 11;
const COMPRESSION_PARENT_0: u8 = 12;
const COMPRESSION_PARENT_1: u8 = 13;

/// Reads and decompresses the hunk map table for a CHD v5 container.
pub fn read_v5_map<R: Read + Seek>(
    reader: &mut R,
    header: &ChdHeader,
) -> Result<Vec<MapEntry>, ChdError> {
    let hunk_count = header.hunk_count();
    if hunk_count == 0 {
        return Ok(Vec::new());
    }
    if hunk_count > usize::MAX as u64 {
        return Err(ChdError::InvalidData(
            "hunk count exceeds platform address space".into(),
        ));
    }

    let file_len = reader.seek(SeekFrom::End(0))?;

    if !header.is_compressed() {
        let map_bytes = hunk_count.saturating_mul(4);
        if header.map_offset + map_bytes > file_len {
            return Err(ChdError::InvalidData(format!(
                "uncompressed map bounds ({}) exceed file length ({})",
                header.map_offset + map_bytes,
                file_len
            )));
        }
        reader.seek(SeekFrom::Start(header.map_offset))?;
        let mut raw_map = vec![0u8; (hunk_count as usize) * 4];
        reader.read_exact(&mut raw_map)?;

        let mut entries = Vec::with_capacity(hunk_count as usize);
        for (hunknum, chunk) in raw_map.as_chunks::<4>().0.iter().enumerate() {
            let blocknum = u32::from_be_bytes(*chunk);
            if blocknum != 0 {
                let offset = (blocknum as u64) * (header.hunk_bytes as u64);
                entries.push(MapEntry {
                    hunk_type: HunkType::Uncompressed,
                    length: header.hunk_bytes,
                    offset,
                    crc16: 0,
                });
            } else {
                entries.push(MapEntry {
                    hunk_type: HunkType::Parent(hunknum as u64),
                    length: 0,
                    offset: 0,
                    crc16: 0,
                });
            }
        }
        return Ok(entries);
    }

    if header.map_offset + 16 > file_len {
        return Err(ChdError::InvalidData(
            "map offset exceeds file length".into(),
        ));
    }
    reader.seek(SeekFrom::Start(header.map_offset))?;
    let mut raw_hdr = [0u8; 16];
    reader.read_exact(&mut raw_hdr)?;

    let complen = u32::from_be_bytes([raw_hdr[0], raw_hdr[1], raw_hdr[2], raw_hdr[3]]) as usize;
    if header.map_offset + 16 + (complen as u64) > file_len {
        return Err(ChdError::InvalidData(format!(
            "compressed map payload ({} bytes) exceeds file length ({})",
            complen, file_len
        )));
    }

    // A compressed map with complen bytes cannot expand to arbitrarily many hunks.
    // In MAME format, maximum RLE expansion is bounded (<2200 hunks per byte).
    let max_hunks = (complen as u64).saturating_mul(2200).saturating_add(1024);
    if hunk_count > max_hunks {
        return Err(ChdError::InvalidData(format!(
            "hunk count ({}) exceeds maximum possible for compressed map payload ({} bytes)",
            hunk_count, complen
        )));
    }

    let firstoffs = u64::from_be_bytes([
        0, 0, raw_hdr[4], raw_hdr[5], raw_hdr[6], raw_hdr[7], raw_hdr[8], raw_hdr[9],
    ]);
    let mapcrc = u16::from_be_bytes([raw_hdr[10], raw_hdr[11]]);
    let lengthbits = raw_hdr[12] as usize;
    let selfbits = raw_hdr[13] as usize;
    let parentbits = raw_hdr[14] as usize;

    let mut compressed = vec![0u8; complen];
    reader.read_exact(&mut compressed)?;

    let mut bitbuf = BitReader::new(&compressed);
    let mut decoder = HuffmanDecoder::<16, 8>::new();
    decoder.import_tree_rle(&mut bitbuf)?;

    let mut raw_types = Vec::with_capacity(hunk_count as usize);
    let mut lastcomp = 0u8;
    let mut repcount = 0usize;

    for _ in 0..hunk_count {
        if repcount > 0 {
            raw_types.push(lastcomp);
            repcount -= 1;
        } else {
            let val = decoder.decode_one(&mut bitbuf) as u8;
            if val == COMPRESSION_RLE_SMALL {
                repcount = 2 + (decoder.decode_one(&mut bitbuf) as usize);
                raw_types.push(lastcomp);
            } else if val == COMPRESSION_RLE_LARGE {
                repcount = 2 + 16 + ((decoder.decode_one(&mut bitbuf) as usize) << 4);
                repcount += decoder.decode_one(&mut bitbuf) as usize;
                raw_types.push(lastcomp);
            } else {
                lastcomp = val;
                raw_types.push(val);
            }
        }
    }

    let mut entries = Vec::with_capacity(hunk_count as usize);
    let mut curoffset = firstoffs;
    let mut last_self = 0u64;
    let mut last_parent = 0u64;
    let mut raw_map_bytes = Vec::with_capacity((hunk_count as usize) * 12);

    for (hunknum, &raw_type) in raw_types.iter().enumerate() {
        let final_type;
        let length;
        let offset;
        let crc;

        match raw_type {
            0..=3 => {
                final_type = raw_type;
                length = bitbuf.read(lengthbits);
                if length > header.hunk_bytes {
                    return Err(ChdError::InvalidData(format!(
                        "compressed hunk length ({}) exceeds hunk_bytes ({})",
                        length, header.hunk_bytes
                    )));
                }
                crc = bitbuf.read(16) as u16;
                offset = curoffset;
                curoffset = curoffset
                    .checked_add(length as u64)
                    .ok_or_else(|| ChdError::InvalidData("map offset overflow".into()))?;
            }
            4 => {
                final_type = raw_type;
                length = header.hunk_bytes;
                crc = bitbuf.read(16) as u16;
                offset = curoffset;
                curoffset = curoffset
                    .checked_add(length as u64)
                    .ok_or_else(|| ChdError::InvalidData("map offset overflow".into()))?;
            }
            5 => {
                final_type = raw_type;
                length = 0;
                crc = 0;
                offset = bitbuf.read(selfbits) as u64;
                last_self = offset;
            }
            6 => {
                final_type = raw_type;
                length = 0;
                crc = 0;
                offset = bitbuf.read(parentbits) as u64;
                last_parent = offset;
            }
            COMPRESSION_SELF_0 => {
                final_type = 5;
                length = 0;
                crc = 0;
                offset = last_self;
            }
            COMPRESSION_SELF_1 => {
                final_type = 5;
                length = 0;
                crc = 0;
                last_self = last_self
                    .checked_add(1)
                    .ok_or_else(|| ChdError::InvalidData("self ref offset overflow".into()))?;
                offset = last_self;
            }
            COMPRESSION_PARENT_SELF => {
                final_type = 6;
                length = 0;
                crc = 0;
                offset =
                    ((hunknum as u64) * (header.hunk_bytes as u64)) / (header.unit_bytes as u64);
                last_parent = offset;
            }
            COMPRESSION_PARENT_0 => {
                final_type = 6;
                length = 0;
                crc = 0;
                offset = last_parent;
            }
            COMPRESSION_PARENT_1 => {
                final_type = 6;
                length = 0;
                crc = 0;
                let step = (header.hunk_bytes as u64) / (header.unit_bytes as u64);
                last_parent = last_parent
                    .checked_add(step)
                    .ok_or_else(|| ChdError::InvalidData("parent offset overflow".into()))?;
                offset = last_parent;
            }
            _ => {
                return Err(ChdError::InvalidData(format!(
                    "unknown compression type in map: {}",
                    raw_type
                )));
            }
        }

        raw_map_bytes.push(final_type);
        raw_map_bytes.push((length >> 16) as u8);
        raw_map_bytes.push((length >> 8) as u8);
        raw_map_bytes.push(length as u8);
        raw_map_bytes.push((offset >> 40) as u8);
        raw_map_bytes.push((offset >> 32) as u8);
        raw_map_bytes.push((offset >> 24) as u8);
        raw_map_bytes.push((offset >> 16) as u8);
        raw_map_bytes.push((offset >> 8) as u8);
        raw_map_bytes.push(offset as u8);
        raw_map_bytes.push((crc >> 8) as u8);
        raw_map_bytes.push(crc as u8);

        let hunk_type = match final_type {
            0..=3 => HunkType::Compressed(final_type),
            4 => HunkType::Uncompressed,
            5 => HunkType::SelfRef(offset),
            6 => HunkType::Parent(offset),
            _ => unreachable!(),
        };

        entries.push(MapEntry {
            hunk_type,
            length,
            offset,
            crc16: crc,
        });
    }

    if bitbuf.overflow() {
        return Err(ChdError::InvalidData("map bitstream overflow".into()));
    }

    let computed_map_crc = crc16(&raw_map_bytes, 0xffff);
    if computed_map_crc != mapcrc {
        return Err(ChdError::CrcMismatch {
            expected: mapcrc,
            found: computed_map_crc,
        });
    }

    Ok(entries)
}

/// Reads the hunk map table for a CHD v3 or v4 container.
pub fn read_v34_map<R: Read + Seek>(
    reader: &mut R,
    header: &ChdHeader,
) -> Result<Vec<MapEntry>, ChdError> {
    let hunk_count = header.hunk_count();
    if hunk_count == 0 {
        return Ok(Vec::new());
    }
    if hunk_count > usize::MAX as u64 {
        return Err(ChdError::InvalidData(
            "hunk count exceeds platform address space".into(),
        ));
    }

    let file_len = reader.seek(SeekFrom::End(0))?;
    let map_bytes = hunk_count.saturating_mul(16);
    if header.map_offset + map_bytes > file_len {
        return Err(ChdError::InvalidData(format!(
            "v3/v4 map bounds ({}) exceed file length ({})",
            header.map_offset + map_bytes,
            file_len
        )));
    }

    reader.seek(SeekFrom::Start(header.map_offset))?;
    let mut raw_map = vec![0u8; (hunk_count as usize) * 16];
    reader.read_exact(&mut raw_map)?;

    let mut entries = Vec::with_capacity(hunk_count as usize);
    for chunk in raw_map.as_chunks::<16>().0 {
        let [b0, b1, b2, b3, b4, b5, b6, b7, c0, c1, c2, c3, l0, l1, l2, flags] = *chunk;
        let blockoffs = u64::from_be_bytes([b0, b1, b2, b3, b4, b5, b6, b7]);
        let _blockcrc = u32::from_be_bytes([c0, c1, c2, c3]);
        let blocklen = u32::from(u16::from_be_bytes([l0, l1])) | (u32::from(l2) << 16);
        let entry_type = flags & 0x0f;

        let (hunk_type, length, offset) = match entry_type {
            1 => {
                if blocklen > header.hunk_bytes {
                    return Err(ChdError::InvalidData(format!(
                        "v3/v4 compressed block length ({}) exceeds hunk_bytes ({})",
                        blocklen, header.hunk_bytes
                    )));
                }
                (HunkType::Compressed(0), blocklen, blockoffs)
            }
            2 => (HunkType::Uncompressed, header.hunk_bytes, blockoffs),
            3 => (HunkType::Mini(blockoffs), 0, blockoffs),
            4 => (HunkType::SelfRef(blockoffs), 0, blockoffs),
            5 => (HunkType::Parent(blockoffs), 0, blockoffs),
            6 => {
                if blocklen > header.hunk_bytes {
                    return Err(ChdError::InvalidData(format!(
                        "v3/v4 compressed block length ({}) exceeds hunk_bytes ({})",
                        blocklen, header.hunk_bytes
                    )));
                }
                (HunkType::Compressed(1), blocklen, blockoffs)
            }
            _ => {
                return Err(ChdError::InvalidData(format!(
                    "unknown v3/v4 map entry type: {}",
                    entry_type
                )))
            }
        };

        entries.push(MapEntry {
            hunk_type,
            length,
            offset,
            crc16: 0,
        });
    }

    Ok(entries)
}

/// Serializes the hunk map into CHD v5 uncompressed format (4 bytes per hunk).
pub fn write_v5_uncompressed_map(entries: &[MapEntry], hunk_bytes: u32) -> Vec<u8> {
    let mut raw_map = Vec::with_capacity(entries.len() * 4);
    for entry in entries {
        match entry.hunk_type {
            HunkType::Uncompressed => {
                let blocknum = (entry.offset / hunk_bytes as u64) as u32;
                raw_map.extend_from_slice(&blocknum.to_be_bytes());
            }
            _ => {
                raw_map.extend_from_slice(&0u32.to_be_bytes());
            }
        }
    }
    raw_map
}

/// Serializes the hunk map into CHD v5 compressed format (RLE + Huffman + bitstream).
pub fn write_v5_map(
    entries: &[MapEntry],
    hunk_bytes: u32,
    unit_bytes: u32,
) -> Result<(Vec<u8>, u64), ChdError> {
    let hunk_count = entries.len();
    if hunk_count == 0 {
        return Ok((vec![0u8; 16], 0));
    }

    let mut raw_map_bytes = Vec::with_capacity(hunk_count * 12);
    for entry in entries {
        let final_type = match entry.hunk_type {
            HunkType::Compressed(c) => c,
            HunkType::Uncompressed | HunkType::Mini(_) => 4,
            HunkType::SelfRef(_) => 5,
            HunkType::Parent(_) => 6,
        };
        raw_map_bytes.push(final_type);
        raw_map_bytes.push((entry.length >> 16) as u8);
        raw_map_bytes.push((entry.length >> 8) as u8);
        raw_map_bytes.push(entry.length as u8);
        raw_map_bytes.push((entry.offset >> 40) as u8);
        raw_map_bytes.push((entry.offset >> 32) as u8);
        raw_map_bytes.push((entry.offset >> 24) as u8);
        raw_map_bytes.push((entry.offset >> 16) as u8);
        raw_map_bytes.push((entry.offset >> 8) as u8);
        raw_map_bytes.push(entry.offset as u8);
        raw_map_bytes.push((entry.crc16 >> 8) as u8);
        raw_map_bytes.push(entry.crc16 as u8);
    }
    let map_crc = crc16(&raw_map_bytes, 0xffff);

    let mut compression_rle = Vec::with_capacity(hunk_count);
    let mut encoder = HuffmanEncoder::<16, 8>::new();

    let mut max_self = 0u64;
    let mut last_self = 0u64;
    let mut max_parent = 0u64;
    let mut last_parent = 0u64;
    let mut max_complen = 0u32;
    let mut lastcomp = 0u8;
    let mut count = 0usize;
    let mut firstoffs = 0u64;

    for (hunknum, entry) in entries.iter().enumerate() {
        let hunknum = hunknum as u64;
        let curcomp = match entry.hunk_type {
            HunkType::Compressed(c) => {
                max_complen = max_complen.max(entry.length);
                if firstoffs == 0 {
                    firstoffs = entry.offset;
                }
                c
            }
            HunkType::Uncompressed | HunkType::Mini(_) => {
                max_complen = max_complen.max(entry.length);
                if firstoffs == 0 {
                    firstoffs = entry.offset;
                }
                4
            }
            HunkType::SelfRef(refhunk) => {
                let code = if refhunk == last_self {
                    COMPRESSION_SELF_0
                } else if refhunk == last_self + 1 {
                    COMPRESSION_SELF_1
                } else {
                    max_self = max_self.max(refhunk);
                    5
                };
                last_self = refhunk;
                code
            }
            HunkType::Parent(refunit) => {
                let expected_unit = (hunknum * hunk_bytes as u64) / (unit_bytes as u64);
                let code = if refunit == expected_unit {
                    COMPRESSION_PARENT_SELF
                } else if refunit == last_parent {
                    COMPRESSION_PARENT_0
                } else if refunit == last_parent + (hunk_bytes as u64) / (unit_bytes as u64) {
                    COMPRESSION_PARENT_1
                } else {
                    max_parent = max_parent.max(refunit);
                    6
                };
                last_parent = refunit;
                code
            }
        };

        if curcomp == lastcomp {
            count += 1;
        }

        if curcomp != lastcomp || hunknum == (hunk_count - 1) as u64 {
            while count != 0 {
                if count < 3 {
                    encoder.histo_one(lastcomp as usize);
                    compression_rle.push(lastcomp);
                    count -= 1;
                } else if count <= 3 + 15 {
                    encoder.histo_one(COMPRESSION_RLE_SMALL as usize);
                    compression_rle.push(COMPRESSION_RLE_SMALL);
                    encoder.histo_one(count - 3);
                    compression_rle.push((count - 3) as u8);
                    count = 0;
                } else {
                    let this_count = count.min(3 + 16 + 255);
                    encoder.histo_one(COMPRESSION_RLE_LARGE as usize);
                    compression_rle.push(COMPRESSION_RLE_LARGE);
                    encoder.histo_one((this_count - 3 - 16) >> 4);
                    compression_rle.push(((this_count - 3 - 16) >> 4) as u8);
                    encoder.histo_one((this_count - 3 - 16) & 15);
                    compression_rle.push(((this_count - 3 - 16) & 15) as u8);
                    count -= this_count;
                }
            }
            if curcomp != lastcomp {
                lastcomp = curcomp;
                encoder.histo_one(lastcomp as usize);
                compression_rle.push(lastcomp);
            }
        }
    }

    let lengthbits = if max_complen == 0 {
        0
    } else {
        32 - max_complen.leading_zeros() as u8
    };
    let selfbits = if max_self == 0 {
        0
    } else {
        64 - max_self.leading_zeros() as u8
    };
    let parentbits = if max_parent == 0 {
        0
    } else {
        64 - max_parent.leading_zeros() as u8
    };

    encoder
        .compute_tree_from_histo()
        .map_err(|e| ChdError::Codec(format!("map huffman error: {:?}", e)))?;
    let mut bitbuf = BitWriter::new();
    encoder
        .export_tree_rle(&mut bitbuf)
        .map_err(|e| ChdError::Codec(format!("map huffman export error: {:?}", e)))?;

    for &symbol in &compression_rle {
        encoder.encode_one(&mut bitbuf, symbol as usize);
    }

    let mut src_idx = 0;
    let mut count = 0usize;
    let mut lastcomp = 0u8;

    for entry in &entries[..hunk_count] {
        if count == 0 {
            let val = compression_rle[src_idx];
            src_idx += 1;
            if val == COMPRESSION_RLE_SMALL {
                count = 2 + (compression_rle[src_idx] as usize);
                src_idx += 1;
            } else if val == COMPRESSION_RLE_LARGE {
                count = 2 + 16 + ((compression_rle[src_idx] as usize) << 4);
                src_idx += 1;
                count += compression_rle[src_idx] as usize;
                src_idx += 1;
            } else {
                lastcomp = val;
            }
        } else {
            count -= 1;
        }

        match lastcomp {
            0..=3 => {
                bitbuf.write(entry.length, lengthbits as usize);
                bitbuf.write(entry.crc16 as u32, 16);
            }
            4 => {
                bitbuf.write(entry.crc16 as u32, 16);
            }
            5 => {
                bitbuf.write(entry.offset as u32, selfbits as usize);
            }
            6 => {
                bitbuf.write(entry.offset as u32, parentbits as usize);
            }
            COMPRESSION_SELF_0
            | COMPRESSION_SELF_1
            | COMPRESSION_PARENT_SELF
            | COMPRESSION_PARENT_0
            | COMPRESSION_PARENT_1 => {}
            _ => unreachable!(),
        }
    }

    let compressed_data = bitbuf.into_bytes();
    let complen = compressed_data.len() as u32;

    let mut result = Vec::with_capacity(16 + compressed_data.len());
    result.extend_from_slice(&complen.to_be_bytes());
    result.extend_from_slice(&firstoffs.to_be_bytes()[2..8]);
    result.extend_from_slice(&map_crc.to_be_bytes());
    result.push(lengthbits);
    result.push(selfbits);
    result.push(parentbits);
    result.push(0);
    result.extend_from_slice(&compressed_data);

    Ok((result, firstoffs))
}

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

    #[test]
    fn test_map_round_trip() {
        let entries = vec![
            MapEntry {
                hunk_type: HunkType::Compressed(0),
                length: 1024,
                offset: 124,
                crc16: 0x1234,
            },
            MapEntry {
                hunk_type: HunkType::Compressed(0),
                length: 2048,
                offset: 1148,
                crc16: 0x5678,
            },
            MapEntry {
                hunk_type: HunkType::SelfRef(0),
                length: 0,
                offset: 0,
                crc16: 0,
            },
            MapEntry {
                hunk_type: HunkType::SelfRef(1),
                length: 0,
                offset: 1,
                crc16: 0,
            },
            MapEntry {
                hunk_type: HunkType::Uncompressed,
                length: 4096,
                offset: 3196,
                crc16: 0xabcd,
            },
        ];

        let (map_bytes, _) = write_v5_map(&entries, 4096, 512).expect("write map");

        let header = ChdHeader {
            tag: *b"MComprHD",
            length: 124,
            version: 5,
            compressors: [u32::from_be_bytes(*b"zlib"), 0, 0, 0],
            logical_bytes: 5 * 4096,
            map_offset: 0,
            meta_offset: 0,
            hunk_bytes: 4096,
            unit_bytes: 512,
            raw_sha1: [0u8; 20],
            sha1: [0u8; 20],
            parent_sha1: [0u8; 20],
        };

        let mut cursor = Cursor::new(map_bytes);
        let read_entries = read_v5_map(&mut cursor, &header).expect("read map");
        assert_eq!(entries.len(), read_entries.len());
        for (i, (expected, actual)) in entries.iter().zip(&read_entries).enumerate() {
            assert_eq!(expected, actual, "mismatch at hunk {}", i);
        }
    }
}