zarust 0.2.0

use std::{
    collections::HashMap,
    io::{Read, Write},
};

use rayon::prelude::*;
use sha2::{Digest, Sha256};

use crate::{
    Error, Result,
    format::{
        BLOCK_SIZE, ENTRIES_PER_OFFSET_RECORD, EntryData, Footer, MAGIC, OffsetRecord,
        ROOT_NAME_OFFSET, Section, VERSION_1, cmp_name, eq_name, path_components,
    },
};

#[derive(Debug)]
enum NodeData {
    File { offset: u64, size: u64 },
    Directory { children: Vec<usize> },
}

#[derive(Debug)]
struct Node {
    name_index: Option<usize>,
    data: NodeData,
}

/// The zstd level the format is written at. Fixed, not a tuning knob.
const COMPRESSION_LEVEL: i32 = 6;

/// A streaming ZArchive writer.
///
/// File contents are emitted as they are added. The directory tables and footer
/// are appended by [`finish`](Self::finish), so the destination never needs to
/// implement `Seek`. Full blocks are compressed in bounded parallel batches and
/// written in their original order.
pub struct ArchiveWriter<W> {
    output: W,
    hasher: Sha256,
    output_offset: u64,
    input_offset: u64,
    block_buffer: Vec<u8>,
    pending_blocks: Vec<Vec<u8>>,
    compression_batch_size: usize,
    offset_records: Vec<OffsetRecord>,
    block_count: u64,
    nodes: Vec<Node>,
    names: Vec<Vec<u8>>,
    name_lookup: HashMap<Vec<u8>, usize>,
    active_file: Option<usize>,
}

impl<W: Write> ArchiveWriter<W> {
    pub fn new(output: W) -> Self {
        Self {
            output,
            hasher: Sha256::new(),
            output_offset: 0,
            input_offset: 0,
            block_buffer: Vec::with_capacity(BLOCK_SIZE),
            pending_blocks: Vec::new(),
            compression_batch_size: rayon::current_num_threads().saturating_mul(2).max(1),
            offset_records: Vec::new(),
            block_count: 0,
            nodes: vec![Node {
                name_index: None,
                data: NodeData::Directory { children: vec![] },
            }],
            names: Vec::new(),
            name_lookup: HashMap::new(),
            active_file: None,
        }
    }

    /// Creates a directory. Missing parents are created when `recursive` is true.
    pub fn make_dir(&mut self, path: impl AsRef<[u8]>, recursive: bool) -> Result<()> {
        self.active_file = None;
        let components = path_components(path.as_ref()).collect::<Vec<_>>();
        if components.is_empty() {
            return Err(Error::InvalidPath("directory path is empty"));
        }

        if recursive {
            let mut parent = 0;
            for component in components {
                validate_component(component)?;
                match self.find_child(parent, component) {
                    Some(child) if matches!(self.nodes[child].data, NodeData::Directory { .. }) => {
                        parent = child;
                    }
                    Some(_) => return Err(Error::NotADirectory),
                    None => parent = self.insert_node(parent, component, false)?,
                }
            }
            Ok(())
        } else {
            let (name, parents) = components
                .split_last()
                .ok_or(Error::InvalidPath("directory path is empty"))?;
            validate_component(name)?;
            let parent = self.find_directory(parents)?;
            if self.find_child(parent, name).is_some() {
                return Err(Error::DuplicatePath);
            }
            self.insert_node(parent, name, false)?;
            Ok(())
        }
    }

    /// Creates a file and makes it the destination of subsequent `append_data` calls.
    pub fn start_file(&mut self, path: impl AsRef<[u8]>) -> Result<()> {
        self.active_file = None;
        let components = path_components(path.as_ref()).collect::<Vec<_>>();
        let (name, parents) = components
            .split_last()
            .ok_or(Error::InvalidPath("file path is empty"))?;
        validate_component(name)?;
        let parent = self.find_directory(parents)?;
        if self.find_child(parent, name).is_some() {
            return Err(Error::DuplicatePath);
        }
        let node = self.insert_node(parent, name, true)?;
        self.active_file = Some(node);
        Ok(())
    }

    /// Appends bytes to the active file.
    pub fn append_data(&mut self, mut bytes: &[u8]) -> Result<()> {
        let active = self.active_file.ok_or(Error::NotAFile)?;
        let data_len = bytes.len() as u64;
        while !bytes.is_empty() {
            let wanted = BLOCK_SIZE - self.block_buffer.len();
            let take = wanted.min(bytes.len());
            self.block_buffer.extend_from_slice(&bytes[..take]);
            bytes = &bytes[take..];
            if self.block_buffer.len() == BLOCK_SIZE {
                let block = std::mem::take(&mut self.block_buffer);
                self.queue_block(block)?;
                self.block_buffer = Vec::with_capacity(BLOCK_SIZE);
            }
        }
        let NodeData::File { size, .. } = &mut self.nodes[active].data else {
            return Err(Error::NotAFile);
        };
        *size = size.checked_add(data_len).ok_or(Error::ArchiveTooLarge)?;
        self.input_offset = self
            .input_offset
            .checked_add(data_len)
            .ok_or(Error::ArchiveTooLarge)?;
        Ok(())
    }

    /// Adds one complete file from a reader.
    pub fn add_file(&mut self, path: impl AsRef<[u8]>, mut source: impl Read) -> Result<u64> {
        self.start_file(path)?;
        let mut total = 0_u64;
        let mut buffer = [0_u8; BLOCK_SIZE];
        loop {
            let read = source.read(&mut buffer)?;
            if read == 0 {
                break;
            }
            self.append_data(&buffer[..read])?;
            total = total
                .checked_add(read as u64)
                .ok_or(Error::ArchiveTooLarge)?;
        }
        Ok(total)
    }

    /// Finalizes the archive and returns the underlying writer.
    pub fn finish(mut self) -> Result<W> {
        self.active_file = None;
        if !self.block_buffer.is_empty() {
            self.block_buffer.resize(BLOCK_SIZE, 0);
            let block = std::mem::take(&mut self.block_buffer);
            self.queue_block(block)?;
        }
        // Keep empty archives readable by the original implementation, which
        // requires at least one offset record.
        if self.block_count == 0 && self.pending_blocks.is_empty() {
            self.queue_block(vec![0; BLOCK_SIZE])?;
        }
        self.flush_blocks()?;

        let compressed = Section {
            offset: 0,
            size: self.output_offset,
        };
        while self.output_offset % 8 != 0 {
            self.write_body(&[0])?;
        }

        let offset_records = self.write_section(|writer| {
            for record in writer.offset_records.clone() {
                writer.write_body(&record.encode())?;
            }
            Ok(())
        })?;

        let (name_offsets, names) = self.write_name_table()?;
        let file_tree = self.write_file_tree(&name_offsets)?;
        let metadata_directory = Section {
            offset: self.output_offset,
            size: 0,
        };
        let metadata = metadata_directory;

        let total_size = self
            .output_offset
            .checked_add(crate::format::FOOTER_SIZE as u64)
            .ok_or(Error::ArchiveTooLarge)?;
        let mut footer = Footer {
            sections: [
                compressed,
                offset_records,
                names,
                file_tree,
                metadata_directory,
                metadata,
            ],
            integrity_hash: [0; 32],
            total_size,
            version: VERSION_1,
            magic: MAGIC,
        };
        let zeroed_footer = footer.encode(true);
        let mut footer_hasher = self.hasher.clone();
        footer_hasher.update(zeroed_footer);
        footer
            .integrity_hash
            .copy_from_slice(&footer_hasher.finalize());
        self.output.write_all(&footer.encode(false))?;
        self.output.flush()?;
        Ok(self.output)
    }

    fn find_directory(&self, components: &[&[u8]]) -> Result<usize> {
        let mut current = 0;
        for component in components {
            validate_component(component)?;
            current = self
                .find_child(current, component)
                .ok_or(Error::InvalidPath("parent directory does not exist"))?;
            if !matches!(self.nodes[current].data, NodeData::Directory { .. }) {
                return Err(Error::NotADirectory);
            }
        }
        Ok(current)
    }

    fn find_child(&self, parent: usize, name: &[u8]) -> Option<usize> {
        let NodeData::Directory { children } = &self.nodes[parent].data else {
            return None;
        };
        children.iter().copied().find(|child| {
            let name_index = self.nodes[*child].name_index.unwrap();
            eq_name(&self.names[name_index], name)
        })
    }

    fn insert_node(&mut self, parent: usize, name: &[u8], is_file: bool) -> Result<usize> {
        let name_index = self.name_index(name);
        let index = self.nodes.len();
        self.nodes.push(Node {
            name_index: Some(name_index),
            data: if is_file {
                NodeData::File {
                    offset: self.input_offset,
                    size: 0,
                }
            } else {
                NodeData::Directory { children: vec![] }
            },
        });
        let NodeData::Directory { children } = &mut self.nodes[parent].data else {
            return Err(Error::NotADirectory);
        };
        children.push(index);
        Ok(index)
    }

    fn name_index(&mut self, name: &[u8]) -> usize {
        if let Some(index) = self.name_lookup.get(name) {
            return *index;
        }
        let index = self.names.len();
        let owned = name.to_vec();
        self.names.push(owned.clone());
        self.name_lookup.insert(owned, index);
        index
    }

    fn queue_block(&mut self, block: Vec<u8>) -> Result<()> {
        debug_assert_eq!(block.len(), BLOCK_SIZE);
        self.pending_blocks.push(block);
        if self.pending_blocks.len() >= self.compression_batch_size {
            self.flush_blocks()?;
        }
        Ok(())
    }

    fn flush_blocks(&mut self) -> Result<()> {
        let blocks = std::mem::take(&mut self.pending_blocks);
        let compressed = blocks
            .into_par_iter()
            .map(compress_block)
            .collect::<std::io::Result<Vec<_>>>()
            .map_err(Error::Io)?;
        for block in compressed {
            self.write_stored_block(&block)?;
        }
        Ok(())
    }

    fn write_stored_block(&mut self, stored: &[u8]) -> Result<()> {
        let compressed_offset = self.output_offset;
        self.write_body(stored)?;

        let sub_index = self.block_count as usize % ENTRIES_PER_OFFSET_RECORD;
        if sub_index == 0 {
            self.offset_records.push(OffsetRecord {
                base_offset: compressed_offset,
                sizes: [0; ENTRIES_PER_OFFSET_RECORD],
            });
        }
        self.offset_records.last_mut().unwrap().sizes[sub_index] =
            u16::try_from(stored.len() - 1).map_err(|_| Error::ArchiveTooLarge)?;
        self.block_count += 1;
        Ok(())
    }

    fn write_name_table(&mut self) -> Result<(Vec<u32>, Section)> {
        let mut offsets = Vec::with_capacity(self.names.len());
        let names = self.names.clone();
        let section = self.write_section(|writer| {
            let mut relative = 0_u64;
            for name in names {
                offsets.push(u32::try_from(relative).map_err(|_| Error::ArchiveTooLarge)?);
                let length = name.len();
                if length < 0x80 {
                    writer.write_body(&[length as u8])?;
                    relative += 1;
                } else {
                    writer.write_body(&[(length as u8 & 0x7f) | 0x80, (length >> 7) as u8])?;
                    relative += 2;
                }
                writer.write_body(&name)?;
                relative = relative
                    .checked_add(length as u64)
                    .ok_or(Error::ArchiveTooLarge)?;
                if relative > ROOT_NAME_OFFSET as u64 {
                    return Err(Error::ArchiveTooLarge);
                }
            }
            Ok(())
        })?;
        Ok((offsets, section))
    }

    fn write_file_tree(&mut self, name_offsets: &[u32]) -> Result<Section> {
        let mut order = vec![0_usize];
        let mut starts = vec![0_u32; self.nodes.len()];
        let mut cursor = 0;
        while cursor < order.len() {
            let node_index = order[cursor];
            if let NodeData::Directory { children } = &self.nodes[node_index].data {
                let mut sorted = children.clone();
                sorted.sort_by(|left, right| {
                    let left = &self.names[self.nodes[*left].name_index.unwrap()];
                    let right = &self.names[self.nodes[*right].name_index.unwrap()];
                    cmp_name(left, right)
                });
                starts[node_index] =
                    u32::try_from(order.len()).map_err(|_| Error::ArchiveTooLarge)?;
                order.extend(sorted);
            }
            cursor += 1;
        }

        self.write_section(|writer| {
            for node_index in order {
                let node = &writer.nodes[node_index];
                let name_offset = node
                    .name_index
                    .map_or(ROOT_NAME_OFFSET, |index| name_offsets[index]);
                let data = match &node.data {
                    NodeData::File { offset, size } => EntryData::File {
                        offset: *offset,
                        size: *size,
                    },
                    NodeData::Directory { children } => EntryData::Directory {
                        start: starts[node_index],
                        count: u32::try_from(children.len()).map_err(|_| Error::ArchiveTooLarge)?,
                    },
                };
                let encoded = crate::format::FileEntry { name_offset, data }.encode()?;
                writer.write_body(&encoded)?;
            }
            Ok(())
        })
    }

    fn write_section(&mut self, write: impl FnOnce(&mut Self) -> Result<()>) -> Result<Section> {
        let offset = self.output_offset;
        write(self)?;
        Ok(Section {
            offset,
            size: self.output_offset - offset,
        })
    }

    fn write_body(&mut self, bytes: &[u8]) -> Result<()> {
        self.output.write_all(bytes)?;
        self.hasher.update(bytes);
        self.output_offset = self
            .output_offset
            .checked_add(bytes.len() as u64)
            .ok_or(Error::ArchiveTooLarge)?;
        Ok(())
    }
}

fn compress_block(block: Vec<u8>) -> std::io::Result<Vec<u8>> {
    debug_assert_eq!(block.len(), BLOCK_SIZE);
    let compressed = zstd::bulk::compress(&block, COMPRESSION_LEVEL)?;
    Ok(if compressed.len() < BLOCK_SIZE {
        compressed
    } else {
        block
    })
}

fn validate_component(component: &[u8]) -> Result<()> {
    if component.is_empty() {
        return Err(Error::InvalidPath("path contains an empty component"));
    }
    if component == b"." || component == b".." {
        return Err(Error::InvalidPath("dot components are not allowed"));
    }
    if component.len() > 0x7fff {
        return Err(Error::NameTooLong);
    }
    Ok(())
}