hunpak 1.2.0

use adler32::adler32;
use anyhow::{anyhow, bail};
use log::debug;
use std::{
    collections::BTreeMap,
    fs::{File, create_dir_all},
    io::{BufReader, Cursor, Read, Write},
    path::{Component, Path, PathBuf},
};

use nom::{
    IResult,
    bytes::complete::{tag, take},
    error::{ErrorKind, make_error},
};

// TODO support alignment
/// Representation of a PAK file
///
/// The file structure is represented as a tree, which starts at the root.
/// The tree only contains the name, checksum, size and position of the file contents
/// in the data.
///
/// To create a new archive, use [PakFile::default] and then add files and directories
/// with [PakFile::add_directory] and [PakFile::add_file] or use
/// [PakFile::add_recursive] to add a folder recursively.
///
/// # Examples
///
/// ```
/// use hunpak::PakFile;
///
/// let mut pak = PakFile::default();
/// pak.add_directory("test").unwrap();
/// pak.add_directory("test/abc").unwrap();
/// pak.add_file("test/abc/test.txt", "Hallo".as_bytes()).unwrap();
/// pak.write("test.pak").unwrap();
/// ```
#[derive(Clone)]
pub struct PakFile {
    pub version: u8,
    pub root: PakEntry,
    pub data: Vec<u8>,
}

impl Default for PakFile {
    fn default() -> Self {
        Self {
            version: 0,
            root: PakEntry {
                name: String::new(),
                file: PakFileEntry::Directory(Vec::new()),
            },
            data: Vec::new(),
        }
    }
}

impl PakFile {
    /// Prints the file structure to the console
    pub fn print_dir(&self) {
        self.root.print(0);
    }

    fn extract_entry<P: AsRef<Path>>(&self, e: &PakEntry, output: P) -> anyhow::Result<()> {
        let path = output.as_ref().join(&e.name);
        debug!("Extracting {}", path.display());
        match &e.file {
            PakFileEntry::Directory(children) => {
                create_dir_all(&path)?;
                for c in children {
                    self.extract_entry(c, &path)?;
                }
            }
            PakFileEntry::File {
                position,
                size,
                checksum,
            } => {
                let mut file = File::create(&path)?;
                file.write_all(self.get_data(position, *size, *checksum)?)?;
            }
        }
        Ok(())
    }

    /// Extract the full archive to a folder on disk
    pub fn extract<P: AsRef<Path>>(&self, output: P) -> anyhow::Result<()> {
        self.extract_entry(&self.root, output)?;
        Ok(())
    }

    /// Read the archive from a given path
    pub fn read<P: AsRef<Path>>(input: P) -> anyhow::Result<Self> {
        let file = File::open(input)?;
        let size = file.metadata()?.len();
        let mut bytes = Vec::with_capacity(size as usize);
        let mut reader = BufReader::new(file);
        reader.read_to_end(&mut bytes)?;

        match pak_without_data(&bytes) {
            Ok((rest, (mut archive, data_size))) => {
                if rest.len() != data_size as usize {
                    bail!("Expected {data_size} data bytes, got {}", rest.len());
                }
                archive.data = rest.to_vec();
                Ok(archive)
            }
            Err(e) => Err(match e {
                nom::Err::Incomplete(_) => unreachable!(),
                nom::Err::Error(e) => anyhow!("{:?}", e.code),
                nom::Err::Failure(e) => anyhow!("{:?}", e.code),
            }),
        }
    }

    /// Write the archive to a file
    pub fn write<P: AsRef<Path>>(&self, output: P) -> anyhow::Result<()> {
        let mut file = File::create(output)?;

        let mut header = Vec::new();
        let mut header_writer = Cursor::new(&mut header);

        header_writer.write_all(b"PAK")?;
        header_writer.write_all(&[self.version])?;
        header_writer.write_all(&0u32.to_le_bytes())?;
        header_writer.write_all(&(self.data.len() as u32).to_le_bytes())?;
        self.root.write(&mut header_writer)?;
        header_writer.write_all(b"DATA")?;
        let header_size = header.len() as u32;
        header[4..8].copy_from_slice(&header_size.to_le_bytes());

        file.write_all(&header)?;
        file.write_all(&self.data)?;

        Ok(())
    }

    fn get_entry<P: AsRef<Path>>(&self, path: P) -> anyhow::Result<Option<&PakEntry>> {
        let path = path.as_ref();
        let mut res = &self.root;
        for component in path.components() {
            if component == Component::RootDir {
                continue;
            }
            let name = component
                .as_os_str()
                .to_str()
                .ok_or(anyhow!("Invalid character in path"))?;
            if let PakFileEntry::Directory(children) = &res.file {
                if let Some(new) = children.iter().find(|x| x.name == name) {
                    res = new;
                } else {
                    return Ok(None);
                }
            } else {
                bail!("Trying to search file inside file")
            }
        }
        Ok(Some(res))
    }

    fn get_entry_mut<P: AsRef<Path>>(&mut self, path: P) -> anyhow::Result<Option<&mut PakEntry>> {
        let path = path.as_ref();
        let mut res = &mut self.root;
        for component in path.components() {
            if component == Component::RootDir {
                continue;
            }
            let name = component
                .as_os_str()
                .to_str()
                .ok_or(anyhow!("Invalid character in path"))?;
            if let PakFileEntry::Directory(children) = &mut res.file {
                if let Some(new) = children.iter_mut().find(|x| x.name == name) {
                    res = new;
                } else {
                    return Ok(None);
                }
            } else {
                bail!("Trying to search file inside file")
            }
        }
        Ok(Some(res))
    }

    fn get_data(&self, position: &PakPosition, size: u32, checksum: u32) -> anyhow::Result<&[u8]> {
        let start = position.as_usize();
        let end = position.as_usize() + size as usize;
        let Some(data) = self.data.get(start..end) else {
            bail!("Trying to read outside of range");
        };
        let expected_checksum = adler32(data)?;
        if expected_checksum != checksum {
            bail!("Checksum mismatch");
        }
        Ok(data)
    }

    /// Get the data of a file at a given path
    pub fn get<P: AsRef<Path>>(&self, path: P) -> anyhow::Result<Option<&[u8]>> {
        if let Some(entry) = self.get_entry(&path)?
            && let PakFileEntry::File {
                position,
                size,
                checksum,
            } = &entry.file
        {
            Ok(Some(self.get_data(position, *size, *checksum)?))
        } else {
            Ok(None)
        }
    }

    /// Add a directory to the archive
    ///
    /// <div class="warning">This is not recursive!</div>
    pub fn add_directory<P: AsRef<Path>>(&mut self, path: P) -> anyhow::Result<()> {
        match self.get_entry(&path)? {
            Some(e) => {
                if !e.file.is_directory() {
                    bail!("Cannot add directory, file with same name existing")
                }
            }
            None => {
                if let Some(file_name) = path.as_ref().file_name().and_then(|x| x.to_str())
                    && let Some(parent) = path.as_ref().parent()
                    && let Some(parent_entry) = self.get_entry_mut(parent)?
                    && let PakFileEntry::Directory(children) = &mut parent_entry.file
                {
                    children.push(PakEntry {
                        name: file_name.to_string(),
                        file: PakFileEntry::Directory(Vec::new()),
                    });
                } else {
                    bail!("Invalid path")
                }
            }
        }
        Ok(())
    }

    /// Add a file with given contents at a given path
    ///
    /// The parent folder has to already exist.
    pub fn add_file<P: AsRef<Path>>(&mut self, path: P, data: &[u8]) -> anyhow::Result<()> {
        match self.get_entry(&path)? {
            Some(_) => {
                bail!("Cannot add file, entry with same name existing")
            }
            None => {
                let current_data_len = self.data.len();
                if let Some(file_name) = path.as_ref().file_name().and_then(|x| x.to_str())
                    && let Some(parent) = path.as_ref().parent()
                    && let Some(parent_entry) = self.get_entry_mut(parent)?
                    && let PakFileEntry::Directory(children) = &mut parent_entry.file
                {
                    let checksum = adler32(data)?;
                    let position = if current_data_len > i32::MAX as usize {
                        PakPosition::Float(current_data_len as f64)
                    } else {
                        PakPosition::Int(current_data_len as u32)
                    };
                    children.push(PakEntry {
                        name: file_name.to_string(),
                        file: PakFileEntry::File {
                            position,
                            size: data.len() as u32,
                            checksum,
                        },
                    });
                    self.data.extend_from_slice(data);
                } else {
                    bail!("Invalid path")
                }
            }
        }
        Ok(())
    }

    fn add_recursive_part<P: AsRef<Path>, Q: AsRef<Path>>(
        &mut self,
        root: &P,
        path: &Q,
    ) -> anyhow::Result<()> {
        debug!("Adding {} to archive", path.as_ref().display());
        for entry in std::fs::read_dir(path)? {
            let entry = entry?;
            let entry_path = entry.path();
            let relative = entry_path.strip_prefix(root)?;
            if entry_path.is_dir() {
                self.add_directory(relative)?;
                self.add_recursive_part(root, &entry_path)?;
            } else {
                let data = std::fs::read(&entry_path)?;
                self.add_file(relative, &data)?;
            }
        }
        Ok(())
    }

    /// Recursively add a folder which is given a certain path to the archive
    ///
    /// All paths in the archive will be relative to the given root.
    pub fn add_recursive<P: AsRef<Path>>(&mut self, root: P) -> anyhow::Result<()> {
        self.add_recursive_part(&root, &root)
    }

    fn entry_mappings<P: AsRef<Path>>(
        &self,
        entry: &PakEntry,
        path: P,
    ) -> anyhow::Result<BTreeMap<PathBuf, (PakPosition, u32, u32)>> {
        let mut res = BTreeMap::new();
        let entry_path = path.as_ref().join(&entry.name);
        match &entry.file {
            PakFileEntry::Directory(children) => {
                for c in children {
                    res.extend(self.entry_mappings(c, &entry_path)?);
                }
            }
            PakFileEntry::File {
                position,
                size,
                checksum,
            } => {
                res.insert(entry_path, (position.clone(), *size, *checksum));
            }
        }
        Ok(res)
    }

    /// Get metadata mappings of all files in the archives
    pub fn mappings(&self) -> anyhow::Result<BTreeMap<PathBuf, (PakPosition, u32, u32)>> {
        self.entry_mappings(&self.root, PathBuf::new())
    }
}

/// File or directory metadata and name
#[derive(Debug, Clone)]
pub struct PakEntry {
    pub name: String,
    pub file: PakFileEntry,
}

impl PakEntry {
    pub fn print(&self, depth: usize) {
        for i in 1..depth {
            if i == depth - 1 && depth > 1 {
                print!("└─ ");
            } else if depth > 0 {
                print!("│  ");
            }
        }
        println!("{}", self.name);
        if let PakFileEntry::Directory(children) = &self.file {
            for child in children {
                child.print(depth + 1);
            }
        }
    }

    fn write<W: std::io::Write>(&self, writer: &mut W) -> anyhow::Result<()> {
        let name_bytes = self.name.as_bytes();
        writer.write_all(&[name_bytes.len() as u8])?;
        writer.write_all(name_bytes)?;
        match &self.file {
            PakFileEntry::Directory(children) => {
                writer.write_all(&[1])?;
                writer.write_all(&(children.len() as u32).to_le_bytes())?;
                for c in children {
                    c.write(writer)?;
                }
            }
            PakFileEntry::File {
                position,
                size,
                checksum,
            } => {
                match position {
                    PakPosition::Float(x) => {
                        writer.write_all(&[2])?;
                        writer.write_all(&x.to_le_bytes())?
                    }
                    PakPosition::Int(x) => {
                        writer.write_all(&[0])?;
                        writer.write_all(&x.to_le_bytes())?
                    }
                }
                writer.write_all(&size.to_le_bytes())?;
                writer.write_all(&checksum.to_le_bytes())?;
            }
        }
        Ok(())
    }
}

/// File or directory metadata
#[derive(Debug, Clone)]
pub enum PakFileEntry {
    Directory(Vec<PakEntry>),
    File {
        position: PakPosition,
        size: u32,
        checksum: u32,
    },
}

impl PakFileEntry {
    pub fn is_directory(&self) -> bool {
        match &self {
            PakFileEntry::Directory(_) => true,
            PakFileEntry::File { .. } => false,
        }
    }
}

fn entry(input: &[u8]) -> IResult<&[u8], PakEntry> {
    let (input, name_len) = nom::number::complete::u8(input)?;
    let (input, name_bytes) = nom::bytes::complete::take(name_len)(input)?;
    let name = match str::from_utf8(name_bytes) {
        Ok(s) => s.to_string(),
        Err(_) => return Err(nom::Err::Error(make_error(input, ErrorKind::Alpha))),
    };
    let (input, flags) = nom::number::complete::u8(input)?;
    let is_directory = flags & 1 != 0;
    let (input, file) = if is_directory {
        let (mut input, num_children) = nom::number::complete::le_u32(input)?;
        let mut children = Vec::with_capacity(num_children as usize);
        for _ in 0..num_children {
            let (i, x) = entry(input)?;
            input = i;
            children.push(x);
        }
        (input, PakFileEntry::Directory(children))
    } else {
        let (input, position) = if flags & 2 != 0 {
            nom::number::complete::le_f64(input).map(|(a, b)| (a, PakPosition::Float(b)))?
        } else {
            nom::number::complete::le_u32(input).map(|(a, b)| (a, PakPosition::Int(b)))?
        };
        let (input, size) = nom::number::complete::le_u32(input)?;
        let (input, checksum) = nom::number::complete::le_u32(input)?;
        (
            input,
            PakFileEntry::File {
                position,
                size,
                checksum,
            },
        )
    };
    Ok((input, PakEntry { name, file }))
}

fn pak_without_data(input: &[u8]) -> IResult<&[u8], (PakFile, u32)> {
    let (input, _) = tag(b"PAK".as_slice())(input)?;
    let (input, version) = nom::number::complete::u8(input)?;
    let (input, header_size) = nom::number::complete::le_u32(input)?;
    let (input, data_size) = nom::number::complete::le_u32(input)?;
    let (input, header) = take(header_size - 16)(input)?;
    let (header, root) = entry(header)?;
    if !header.is_empty() {
        return Err(nom::Err::Failure(make_error(input, ErrorKind::LengthValue)));
    }
    let (input, _) = tag(b"DATA".as_slice())(input)?;
    Ok((
        input,
        (
            PakFile {
                version,
                root,
                data: Vec::new(),
            },
            data_size,
        ),
    ))
}

/// Position in data
///
/// This is needed, because the underlying data structure saves positions which
/// are bigger than [i32::MAX] as doubles.
#[derive(Debug, Clone)]
pub enum PakPosition {
    Float(f64),
    Int(u32),
}

impl PakPosition {
    pub fn as_usize(&self) -> usize {
        match self {
            PakPosition::Float(x) => *x as usize,
            PakPosition::Int(x) => *x as usize,
        }
    }
}