zarust 0.2.1

use std::{
    collections::{HashMap, HashSet, VecDeque},
    fs::File,
    io::{Read, Seek, SeekFrom},
    path::Path,
};

use sha2::{Digest, Sha256};

use crate::{
    Error, Result,
    format::{
        BLOCK_SIZE, ENTRIES_PER_OFFSET_RECORD, EntryData, FILE_ENTRY_SIZE, FOOTER_SIZE, FileEntry,
        Footer, MAGIC, OFFSET_RECORD_SIZE, OffsetRecord, ROOT_NAME_OFFSET, VERSION_1, eq_name,
        path_components,
    },
};

/// A stable index into an archive's file tree.
pub type NodeHandle = u32;

/// The root directory handle.
pub const ROOT_NODE: NodeHandle = 0;

/// The kind of an archive entry.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum EntryKind {
    File,
    Directory,
}

/// A borrowed directory entry.
#[derive(Clone, Copy, Debug)]
pub struct DirEntry<'a> {
    /// The component name as stored in the archive.
    pub name: &'a [u8],
    pub kind: EntryKind,
    /// File size in bytes, or zero for a directory.
    pub size: u64,
    pub handle: NodeHandle,
}

impl DirEntry<'_> {
    #[must_use]
    pub fn is_file(&self) -> bool {
        self.kind == EntryKind::File
    }

    #[must_use]
    pub fn is_directory(&self) -> bool {
        self.kind == EntryKind::Directory
    }
}

/// A random-access ZArchive reader.
pub struct ArchiveReader<R> {
    source: R,
    file_size: u64,
    expected_hash: [u8; 32],
    compressed_data_offset: u64,
    compressed_data_size: u64,
    offset_records: Vec<OffsetRecord>,
    names: Vec<u8>,
    entries: Vec<FileEntry>,
    cache: BlockCache,
}

impl ArchiveReader<File> {
    /// Opens an archive from the filesystem.
    pub fn open(path: impl AsRef<Path>) -> Result<Self> {
        Self::new(File::open(path)?)
    }
}

impl<R: Read + Seek> ArchiveReader<R> {
    /// Opens an archive from any readable, seekable source.
    pub fn new(mut source: R) -> Result<Self> {
        let file_size = source.seek(SeekFrom::End(0))?;
        if file_size < FOOTER_SIZE as u64 {
            return Err(Error::InvalidArchive(
                "file is too small to contain a footer",
            ));
        }

        let footer_offset = file_size - FOOTER_SIZE as u64;
        let mut footer_bytes = [0; FOOTER_SIZE];
        read_exact_at(&mut source, footer_offset, &mut footer_bytes)?;
        let footer = Footer::decode(&footer_bytes);
        validate_footer(&footer, file_size, footer_offset)?;

        let offset_section = footer.sections[1];
        let offset_len = usize::try_from(offset_section.size)
            .map_err(|_| Error::InvalidArchive("offset table is too large"))?;
        if offset_len % OFFSET_RECORD_SIZE != 0 {
            return Err(Error::InvalidArchive("offset table has a partial record"));
        }
        let mut bytes = vec![0; offset_len];
        read_exact_at(&mut source, offset_section.offset, &mut bytes)?;
        let offset_records = bytes
            .chunks_exact(OFFSET_RECORD_SIZE)
            .map(OffsetRecord::decode)
            .collect::<Vec<_>>();

        let names_section = footer.sections[2];
        let names_len = usize::try_from(names_section.size)
            .map_err(|_| Error::InvalidArchive("name table is too large"))?;
        let mut names = vec![0; names_len];
        read_exact_at(&mut source, names_section.offset, &mut names)?;

        let tree_section = footer.sections[3];
        let tree_len = usize::try_from(tree_section.size)
            .map_err(|_| Error::InvalidArchive("file tree is too large"))?;
        if tree_len == 0 || tree_len % FILE_ENTRY_SIZE != 0 {
            return Err(Error::InvalidArchive("file tree has an invalid size"));
        }
        let mut bytes = vec![0; tree_len];
        read_exact_at(&mut source, tree_section.offset, &mut bytes)?;
        let entries = bytes
            .chunks_exact(FILE_ENTRY_SIZE)
            .map(FileEntry::decode)
            .collect::<Vec<_>>();

        validate_tree(&entries, &names, &offset_records)?;

        Ok(Self {
            source,
            file_size,
            expected_hash: footer.integrity_hash,
            compressed_data_offset: footer.sections[0].offset,
            compressed_data_size: footer.sections[0].size,
            offset_records,
            names,
            entries,
            cache: BlockCache::new(64),
        })
    }

    /// Returns the node at `path`, using ASCII case-insensitive matching.
    pub fn lookup(&self, path: impl AsRef<[u8]>) -> Option<NodeHandle> {
        self.lookup_kind(path, None)
    }

    /// Returns the node at `path` only if it has the requested kind.
    pub fn lookup_kind(
        &self,
        path: impl AsRef<[u8]>,
        expected: Option<EntryKind>,
    ) -> Option<NodeHandle> {
        let mut current = ROOT_NODE;
        for component in path_components(path.as_ref()) {
            let entry = self.entries.get(current as usize)?;
            let EntryData::Directory { start, count } = entry.data else {
                return None;
            };
            let end = start.checked_add(count)?;
            current = (start..end).find(|index| {
                self.entries
                    .get(*index as usize)
                    .and_then(|entry| self.name(entry.name_offset).ok())
                    .is_some_and(|name| eq_name(name, component))
            })?;
        }

        expected
            .is_none_or(|kind| self.entry_kind(current) == Some(kind))
            .then_some(current)
    }

    /// The total size of the archive file, including its tables and footer.
    #[must_use]
    pub fn archive_size(&self) -> u64 {
        self.file_size
    }

    /// The size of the compressed data section, excluding the tables and footer.
    #[must_use]
    pub fn compressed_data_size(&self) -> u64 {
        self.compressed_data_size
    }

    #[must_use]
    pub fn entry_kind(&self, handle: NodeHandle) -> Option<EntryKind> {
        self.entries
            .get(handle as usize)
            .map(|entry| match entry.data {
                EntryData::File { .. } => EntryKind::File,
                EntryData::Directory { .. } => EntryKind::Directory,
            })
    }

    pub fn directory_len(&self, handle: NodeHandle) -> Result<u32> {
        match self.entry(handle)?.data {
            EntryData::Directory { count, .. } => Ok(count),
            EntryData::File { .. } => Err(Error::NotADirectory),
        }
    }

    pub fn directory_entry(&self, handle: NodeHandle, index: u32) -> Result<DirEntry<'_>> {
        let EntryData::Directory { start, count } = self.entry(handle)?.data else {
            return Err(Error::NotADirectory);
        };
        if index >= count {
            return Err(Error::InvalidPath("directory entry index is out of range"));
        }
        let child_handle = start
            .checked_add(index)
            .ok_or(Error::InvalidArchive("directory range overflows"))?;
        let child = self.entry(child_handle)?;
        let (kind, size) = match child.data {
            EntryData::File { size, .. } => (EntryKind::File, size),
            EntryData::Directory { .. } => (EntryKind::Directory, 0),
        };
        Ok(DirEntry {
            name: self.name(child.name_offset)?,
            kind,
            size,
            handle: child_handle,
        })
    }

    pub fn directory_entries(&self, handle: NodeHandle) -> Result<Vec<DirEntry<'_>>> {
        (0..self.directory_len(handle)?)
            .map(|index| self.directory_entry(handle, index))
            .collect()
    }

    pub fn file_size(&self, handle: NodeHandle) -> Result<u64> {
        match self.entry(handle)?.data {
            EntryData::File { size, .. } => Ok(size),
            EntryData::Directory { .. } => Err(Error::NotAFile),
        }
    }

    /// Reads at most `buffer.len()` bytes beginning at an offset within a file.
    pub fn read_file(
        &mut self,
        handle: NodeHandle,
        offset: u64,
        buffer: &mut [u8],
    ) -> Result<usize> {
        let EntryData::File {
            offset: file_offset,
            size,
        } = self.entry(handle)?.data
        else {
            return Err(Error::NotAFile);
        };
        if offset >= size || buffer.is_empty() {
            return Ok(0);
        }

        let available = usize::try_from((size - offset).min(buffer.len() as u64)).unwrap();
        let mut raw_offset = file_offset
            .checked_add(offset)
            .ok_or(Error::InvalidArchive("file offset overflows"))?;
        let mut written = 0;
        while written < available {
            let block_index = raw_offset / BLOCK_SIZE as u64;
            let block_offset = raw_offset as usize % BLOCK_SIZE;
            let step = (available - written).min(BLOCK_SIZE - block_offset);
            let block = self.load_block(block_index)?;
            buffer[written..written + step]
                .copy_from_slice(&block[block_offset..block_offset + step]);
            written += step;
            raw_offset += step as u64;
        }
        Ok(written)
    }

    /// Reads an entire archived file into memory.
    pub fn read_file_to_end(&mut self, handle: NodeHandle) -> Result<Vec<u8>> {
        let size = usize::try_from(self.file_size(handle)?).map_err(|_| Error::ArchiveTooLarge)?;
        let mut bytes = vec![0; size];
        let read = self.read_file(handle, 0, &mut bytes)?;
        if read != size {
            return Err(Error::InvalidArchive("file data ended unexpectedly"));
        }
        Ok(bytes)
    }

    /// Verifies the SHA-256 digest stored in the archive footer.
    pub fn verify_integrity(&mut self) -> Result<bool> {
        self.verify_integrity_with(|_| {})
    }

    /// Verifies the footer digest, reporting bytes hashed as it goes.
    ///
    /// `progress` is called with the size of each chunk as it is consumed, so
    /// callers can drive a progress indicator over what is otherwise a single
    /// long pass across the whole file.
    pub fn verify_integrity_with(&mut self, mut progress: impl FnMut(u64)) -> Result<bool> {
        const HASH_OFFSET_IN_FOOTER: u64 = 6 * 16;
        let hash_start = self.file_size - FOOTER_SIZE as u64 + HASH_OFFSET_IN_FOOTER;
        let hash_end = hash_start + 32;
        self.source.seek(SeekFrom::Start(0))?;
        let mut hasher = Sha256::new();
        let mut absolute = 0_u64;
        let mut buffer = [0_u8; 64 * 1024];
        while absolute < self.file_size {
            let wanted =
                usize::try_from((self.file_size - absolute).min(buffer.len() as u64)).unwrap();
            self.source.read_exact(&mut buffer[..wanted])?;
            let overlap_start = absolute.max(hash_start);
            let overlap_end = (absolute + wanted as u64).min(hash_end);
            if overlap_start < overlap_end {
                buffer[(overlap_start - absolute) as usize..(overlap_end - absolute) as usize]
                    .fill(0);
            }
            hasher.update(&buffer[..wanted]);
            absolute += wanted as u64;
            progress(wanted as u64);
        }
        Ok(hasher.finalize().as_slice() == self.expected_hash)
    }

    pub fn into_inner(self) -> R {
        self.source
    }

    fn entry(&self, handle: NodeHandle) -> Result<&FileEntry> {
        self.entries
            .get(handle as usize)
            .ok_or(Error::InvalidPath("node handle is out of range"))
    }

    fn name(&self, offset: u32) -> Result<&[u8]> {
        if offset == ROOT_NAME_OFFSET {
            return Ok(&[]);
        }
        let offset = offset as usize;
        let first = *self
            .names
            .get(offset)
            .ok_or(Error::InvalidArchive("name offset is out of range"))?;
        let (length, header_len) = if first & 0x80 == 0 {
            (usize::from(first), 1)
        } else {
            let second = *self
                .names
                .get(offset + 1)
                .ok_or(Error::InvalidArchive("extended name header is truncated"))?;
            (usize::from(first & 0x7f) | (usize::from(second) << 7), 2)
        };
        let start = offset + header_len;
        let end = start
            .checked_add(length)
            .ok_or(Error::InvalidArchive("name length overflows"))?;
        self.names
            .get(start..end)
            .ok_or(Error::InvalidArchive("name is truncated"))
    }

    fn load_block(&mut self, block_index: u64) -> Result<&[u8]> {
        if self.cache.contains(block_index) {
            return Ok(self.cache.get(block_index).unwrap());
        }
        let record_index = usize::try_from(block_index / ENTRIES_PER_OFFSET_RECORD as u64)
            .map_err(|_| Error::InvalidArchive("block index is too large"))?;
        let sub_index = block_index as usize % ENTRIES_PER_OFFSET_RECORD;
        let record = self
            .offset_records
            .get(record_index)
            .ok_or(Error::InvalidArchive("file references a missing block"))?;
        let relative_offset = record.sizes[..sub_index]
            .iter()
            .try_fold(record.base_offset, |offset, size| {
                offset.checked_add(u64::from(*size) + 1)
            })
            .ok_or(Error::InvalidArchive("compressed block offset overflows"))?;
        let compressed_size = usize::from(record.sizes[sub_index]) + 1;
        let relative_end = relative_offset
            .checked_add(compressed_size as u64)
            .ok_or(Error::InvalidArchive("compressed block range overflows"))?;
        if relative_end > self.compressed_data_size {
            return Err(Error::InvalidArchive(
                "compressed block is outside its section",
            ));
        }

        let mut compressed = vec![0; compressed_size];
        read_exact_at(
            &mut self.source,
            self.compressed_data_offset + relative_offset,
            &mut compressed,
        )?;
        let block = if compressed_size == BLOCK_SIZE {
            compressed
        } else {
            zstd::bulk::decompress(&compressed, BLOCK_SIZE).map_err(Error::Io)?
        };
        if block.len() != BLOCK_SIZE {
            return Err(Error::InvalidArchive(
                "decompressed block has the wrong size",
            ));
        }
        self.cache.insert(block_index, block);
        Ok(self.cache.get(block_index).unwrap())
    }
}

fn read_exact_at(source: &mut (impl Read + Seek), offset: u64, bytes: &mut [u8]) -> Result<()> {
    if bytes.is_empty() {
        return Ok(());
    }
    source.seek(SeekFrom::Start(offset))?;
    source.read_exact(bytes)?;
    Ok(())
}

fn validate_footer(footer: &Footer, file_size: u64, footer_offset: u64) -> Result<()> {
    if footer.magic != MAGIC {
        return Err(Error::InvalidArchive("footer magic does not match"));
    }
    if footer.version != VERSION_1 {
        return Err(Error::InvalidArchive("archive version is not supported"));
    }
    if footer.total_size != file_size {
        return Err(Error::InvalidArchive("footer size does not match the file"));
    }
    if footer
        .sections
        .iter()
        .any(|section| !section.is_within(footer_offset))
    {
        return Err(Error::InvalidArchive(
            "a section is outside the archive body",
        ));
    }
    Ok(())
}

fn validate_tree(entries: &[FileEntry], names: &[u8], offsets: &[OffsetRecord]) -> Result<()> {
    let root = entries
        .first()
        .ok_or(Error::InvalidArchive("file tree is empty"))?;
    if !matches!(root.data, EntryData::Directory { .. }) || root.name_offset != ROOT_NAME_OFFSET {
        return Err(Error::InvalidArchive(
            "first file-tree entry is not the root",
        ));
    }
    let max_uncompressed = (offsets.len() as u64)
        .checked_mul(ENTRIES_PER_OFFSET_RECORD as u64)
        .and_then(|blocks| blocks.checked_mul(BLOCK_SIZE as u64))
        .ok_or(Error::InvalidArchive("block table size overflows"))?;
    for (index, entry) in entries.iter().enumerate() {
        if index != 0 {
            validate_name(names, entry.name_offset)?;
        }
        match entry.data {
            EntryData::Directory { start, count } => {
                let end = start
                    .checked_add(count)
                    .ok_or(Error::InvalidArchive("directory range overflows"))?;
                if end as usize > entries.len() {
                    return Err(Error::InvalidArchive(
                        "directory range is outside the file tree",
                    ));
                }
            }
            EntryData::File { offset, size } => {
                if offset
                    .checked_add(size)
                    .is_none_or(|end| end > max_uncompressed)
                {
                    return Err(Error::InvalidArchive(
                        "file range is outside the block table",
                    ));
                }
            }
        }
    }

    let mut visited = vec![false; entries.len()];
    let mut pending = vec![0_usize];
    visited[0] = true;
    while let Some(index) = pending.pop() {
        let EntryData::Directory { start, count } = entries[index].data else {
            continue;
        };
        let mut child_names = HashSet::with_capacity(count as usize);
        for child in start..start + count {
            let child = child as usize;
            if visited[child] {
                return Err(Error::InvalidArchive(
                    "file tree contains a cycle or shared child",
                ));
            }
            visited[child] = true;
            let name = decoded_name(names, entries[child].name_offset)?;
            let folded = name.iter().map(u8::to_ascii_lowercase).collect::<Vec<_>>();
            if !child_names.insert(folded) {
                return Err(Error::InvalidArchive(
                    "directory contains duplicate case-insensitive names",
                ));
            }
            pending.push(child);
        }
    }
    if visited.contains(&false) {
        return Err(Error::InvalidArchive(
            "file tree contains an unreachable entry",
        ));
    }
    Ok(())
}

fn validate_name(names: &[u8], offset: u32) -> Result<()> {
    decoded_name(names, offset).map(|_| ())
}

fn decoded_name(names: &[u8], offset: u32) -> Result<&[u8]> {
    if offset == ROOT_NAME_OFFSET {
        return Err(Error::InvalidArchive(
            "non-root entry uses the root name marker",
        ));
    }
    let offset = offset as usize;
    let first = *names
        .get(offset)
        .ok_or(Error::InvalidArchive("name offset is out of range"))?;
    let (length, header) = if first & 0x80 == 0 {
        (usize::from(first), 1)
    } else {
        let second = *names
            .get(offset + 1)
            .ok_or(Error::InvalidArchive("extended name header is truncated"))?;
        (usize::from(first & 0x7f) | (usize::from(second) << 7), 2)
    };
    if length == 0 || offset + header + length > names.len() {
        return Err(Error::InvalidArchive("entry name is empty or truncated"));
    }
    Ok(&names[offset + header..offset + header + length])
}

struct BlockCache {
    capacity: usize,
    blocks: HashMap<u64, Vec<u8>>,
    order: VecDeque<u64>,
}

impl BlockCache {
    fn new(capacity: usize) -> Self {
        Self {
            capacity,
            blocks: HashMap::with_capacity(capacity),
            order: VecDeque::with_capacity(capacity),
        }
    }

    fn contains(&self, index: u64) -> bool {
        self.blocks.contains_key(&index)
    }

    fn get(&mut self, index: u64) -> Option<&[u8]> {
        if self.blocks.contains_key(&index) {
            if let Some(position) = self.order.iter().position(|cached| *cached == index) {
                self.order.remove(position);
            }
            self.order.push_back(index);
        }
        self.blocks.get(&index).map(Vec::as_slice)
    }

    fn insert(&mut self, index: u64, block: Vec<u8>) {
        if self.blocks.len() == self.capacity
            && let Some(oldest) = self.order.pop_front()
        {
            self.blocks.remove(&oldest);
        }
        self.blocks.insert(index, block);
        self.order.push_back(index);
    }
}