keeshond_datapack 0.13.0

use std::collections::HashMap;
use std::fs::File;
use std::hash::BuildHasherDefault;
use std::io::{BufReader, Write, BufWriter, Seek, SeekFrom, Read};
use std::iter::Iterator;
use std::path::{PathBuf, StripPrefixError};

use crate::ReadSeek;

use walkdir::{IntoIter, WalkDir};
use failure::Fail;
use rustc_hash::FxHasher;

use zip::ZipArchive;
use zip::result::ZipError;

/// A numeric ID used to refer to a [Source].
pub type SourceId = usize;

/// Trust settings for a given [Source]
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
pub enum TrustLevel
{
    /// The source is untrusted (it cannot load resources that require trust)
    UntrustedSource,
    /// The source is trusted (it can load resources that require trust)
    TrustedSource
}

/// Holds a list of [Source] objects and selects one to use when loading a package
pub struct SourceManager
{
    sources : Vec<Box<dyn Source>>
}

impl SourceManager
{
    /// Constructs a new, empty [SourceManager]
    pub fn new() -> SourceManager
    {
        SourceManager
        {
            sources : Vec::new()
        }
    }

    /// Adds the given [Source] to the end of the list.
    pub fn add_source(&mut self, source : Box<dyn Source>) -> SourceId
    {
        self.sources.push(source);

        self.sources.len()
    }

    /// Returns a reference to the [Source] of the given ID.
    pub fn source(&mut self, id : SourceId) -> Option<&mut Box<dyn Source>>
    {
        if id == 0 || id - 1 >= self.sources.len()
        {
            return None
        }

        Some(&mut self.sources[id - 1])
    }

    /// Returns the [SourceId] of the [Source] that has the given package, going in
    ///  reverse order of when they were added. If no suitable [Source] is available,
    ///  returns a "null" ID.
    pub fn package_source_id(&self, package_name : &str) -> SourceId
    {
        for i in (0..self.sources.len()).rev()
        {
            if self.sources[i].has_package(package_name)
            {
                return i + 1;
            }
        }

        0
    }

    /// Retrieves a mutable reference to a [Source] that has the given package, going in
    ///  reverse order of when they were added. If no suitable [Source] is available,
    ///  returns [None].
    pub fn package_source(&mut self, package_name : &str) -> Option<&mut Box<dyn Source>>
    {
        for source in self.sources.iter_mut().rev()
        {
            if source.has_package(package_name)
            {
                return Some(source);
            }
        }

        None
    }

    /// Removes all sources from the manager. Existing source IDs are invalidated.
    pub fn clear(&mut self)
    {
        self.sources.clear();
    }
}

/// An error returned when failing to access a package in a [Source]
#[derive(Debug, Fail)]
pub enum PackageError
{
    /// An error occurred while iterating through the packages, but no more information is available.
    ///  You should generally avoid using this unless you have to.
    #[fail(display = "Failed to access package source")]
    Generic,
    /// An error of type [std::io::Error] occurred.
    #[fail(display = "{}", _0)]
    IoError(#[cause] std::io::Error),
    /// A package item's path does not belong to the package's path.
    #[fail(display = "Package item does not belong to the given prefix: {}", _0)]
    PrefixMismatch(StripPrefixError),
    /// The package has no [crate::DataObject] by the given pathname
    #[fail(display = "The data object was not found")]
    DataNotFound,
    /// The operation is not supported by this implementation
    #[fail(display = "Operation not supported")]
    NotSupported,
    /// A logical error occurred while reading a source
    #[fail(display = "The package contained invalid data: {}", _0)]
    BadData(String)
}

impl From<std::io::Error> for PackageError
{
    fn from(error : std::io::Error) -> PackageError
    {
        PackageError::IoError(error)
    }
}

/// Represents a location that packages can be loaded from. For example, you could load packages from the
///  filesystem (via [FilesystemSource], or out of a special archive format.
pub trait Source
{
    /// The path that this [Source] originates from. Only used for debug purposes.
    fn get_uri(&self) -> &str;
    /// Returns true if the [Source] has a package of the given name, otherwise returns false
    fn has_package(&self, package_name : &str) -> bool;
    /// Returns a list of all packages available in this [Source]. Do not call this repeatedly!
    fn list_packages(&mut self) -> Vec<String>;
    /// Returns a [Read] + [Seek] for the file at the given pathname, if one exists.
    fn read_file<'a>(&'a mut self, package_name: &str, pathname: &str) -> Result<Box<dyn ReadSeek + 'a>, PackageError>;
    /// Returns a [Write] for the file at the given pathname.
    #[allow(unused_variables)]
    fn write_file<'a>(&'a mut self, package_name : &str, pathname : &str) -> Result<Box<dyn Write + 'a>, PackageError> { Err(PackageError::NotSupported) }
    /// Returns an iterator through the items in a given package, if the [Source] has said package
    fn iter_entries<'a>(&'a mut self, package_name : &str, type_folder : &str) -> Box<dyn Iterator<Item = Result<String, PackageError>> + 'a>;
    /// Returns the source's trust level for the given package. Trusted sources are able to load
    ///  resource types marked as requiring trust.
    fn trust_level(&self, package_name : &str) -> TrustLevel;
}


////////////////////////////////////////////////////////////////////////////////


struct EmptyEntryIter {}

impl EmptyEntryIter
{
    fn new() -> EmptyEntryIter { EmptyEntryIter {} }
}

impl Iterator for EmptyEntryIter
{
    type Item = Result<String, PackageError>;

    fn next(&mut self) -> Option<Result<String, PackageError>>
    {
        None
    }
}


struct FilesystemIter
{
    basepath : PathBuf,
    walkdir : IntoIter
}

impl FilesystemIter
{
    fn new<P : Into<PathBuf>>(basepath : P) -> FilesystemIter
    {
        let path = basepath.into();

        FilesystemIter
        {
            basepath : path.clone(),
            walkdir : WalkDir::new(path).into_iter(),
        }
    }
}

impl Iterator for FilesystemIter
{
    type Item = Result<String, PackageError>;

    fn next(&mut self) -> Option<Result<String, PackageError>>
    {
        while let Some(dir_next) = self.walkdir.next()
        {
            let dir_entry = match dir_next
            {
                Ok(dir_entry) => { dir_entry }
                Err(error) =>
                {
                    if let Some(io_error) = error.io_error()
                    {
                        if io_error.kind() == std::io::ErrorKind::NotFound
                        {
                            return None;
                        }
                    }

                    return Some(Err(PackageError::IoError(error.into())));
                }
            };

            if !dir_entry.file_type().is_file()
            {
                continue;
            }

            let path = dir_entry.path();
            let filepath = match path.strip_prefix(self.basepath.clone())
            {
                Ok(filepath) => { filepath }
                Err(error) =>
                {
                    return Some(Err(PackageError::PrefixMismatch(error)));
                }
            };

            let mut pathname = String::new();
            let mut first_part = true;

            for part in filepath.iter()
            {
                if !first_part
                {
                    pathname += "/";
                }
                pathname += &part.to_string_lossy();

                first_part = false;
            }

            return Some(Ok(pathname));
        }

        None
    }
}

/// A [Source] that loads packages from the filesystem. This is a good source to use for
///  development, or if you don't care about packaging your data files into an archive.
pub struct FilesystemSource
{
    basedir : String,
    package_list : Option<Vec<String>>,
    trust : TrustLevel
}

impl FilesystemSource
{
    /// Creates a new [FilesystemSource] using the given directory to look for packages in
    pub fn new(directory : &str, trust : TrustLevel) -> FilesystemSource
    {
        FilesystemSource
        {
            basedir : directory.to_string(),
            package_list : None,
            trust
        }
    }
}

impl Source for FilesystemSource
{
    fn get_uri(&self) -> &str
    {
        &self.basedir
    }

    fn has_package(&self, package_name : &str) -> bool
    {
        let mut path = PathBuf::from(&self.basedir);
        path.push(package_name);

        path.exists()
    }

    fn list_packages(&mut self) -> Vec<String>
    {
        if let Some(package_list) = &self.package_list
        {
            return package_list.clone();
        }

        let mut package_list : Vec<String> = Vec::new();
        let path = PathBuf::from(&self.basedir);

        let walkdir = WalkDir::new(path).min_depth(1).max_depth(1).into_iter();

        for dir_entry in walkdir
        {
            if let Ok(entry) = dir_entry
            {
                if !entry.file_type().is_dir()
                {
                    continue;
                }

                package_list.push(entry.file_name().to_string_lossy().to_string());
            }
        }

        package_list.sort();

        self.package_list = Some(package_list.clone());

        package_list
    }

    fn read_file<'a>(&'a mut self, package_name: &str, pathname: &str) -> Result<Box<dyn ReadSeek + 'a>, PackageError>
    {
        let mut path = PathBuf::from(&self.basedir);
        path.push(format!("{}/{}", package_name, pathname));

        let file = BufReader::new(File::open(path).map_err(
            |error| PackageError::IoError(error))?);

        Ok(Box::new(file))
    }

    fn write_file<'a>(&'a mut self, package_name : &str, pathname : &str) -> Result<Box<dyn Write + 'a>, PackageError>
    {
        let mut path = PathBuf::from(&self.basedir);
        path.push(format!("{}/{}", package_name, pathname));
        let mut folder_path = path.clone();

        if !pathname.ends_with("/")
        {
            folder_path.pop();
        }

        std::fs::create_dir_all(folder_path).map_err(
            |error| PackageError::IoError(error))?;

        let file = BufWriter::new(File::create(path).map_err(
            |error| PackageError::IoError(error))?);

        Ok(Box::new(file))
    }

    fn iter_entries<'a>(&'a mut self, package_name : &str, type_folder : &str) -> Box<dyn Iterator<Item = Result<String, PackageError>> + 'a>
    {
        let mut path = PathBuf::from(&self.basedir);
        path.push(package_name);
        path.push(type_folder);

        Box::new(FilesystemIter::new(path))
    }

    fn trust_level(&self, _package_name: &str) -> TrustLevel
    {
        self.trust
    }
}

fn package_error_for_zip_error(error : ZipError) -> PackageError
{
    return match error
    {
        ZipError::Io(io_error) => { PackageError::IoError(io_error) }
        ZipError::InvalidArchive(error_msg) => { PackageError::BadData(error_msg.to_string()) }
        ZipError::UnsupportedArchive(error_msg) => { PackageError::BadData(error_msg.to_string()) }
        ZipError::FileNotFound => { PackageError::DataNotFound }
    }
}

struct FakeSeekReader<R : Read>
{
    reader : R
}

impl<R : Read> FakeSeekReader<R>
{
    fn new(reader : R) -> FakeSeekReader<R>
    {
        FakeSeekReader
        {
            reader
        }
    }
}

impl<R : Read> Read for FakeSeekReader<R>
{
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize>
    {
        self.reader.read(buf)
    }
}

impl<R : Read> Seek for FakeSeekReader<R>
{
    fn seek(&mut self, _pos: SeekFrom) -> std::io::Result<u64>
    {
        Err(std::io::Error::new(std::io::ErrorKind::Other, "Seek not supported on ZIP archives"))
    }
}

struct ZipFolderIter<'a, T : Iterator<Item = &'a str>>
{
    iter : T,
    type_folder : String
}

impl<'a, T : Iterator<Item = &'a str>> ZipFolderIter<'a, T>
{
    fn new(iter : T, type_folder : &str) -> ZipFolderIter<'a, T>
    {
        ZipFolderIter
        {
            iter,
            type_folder : type_folder.to_string()
        }
    }
}

impl<'a, T : Iterator<Item = &'a str>> Iterator for ZipFolderIter<'a, T>
{
    type Item = Result<String, PackageError>;

    fn next(&mut self) -> Option<Result<String, PackageError>>
    {
        while let Some(item) = self.iter.next()
        {
            if let Some(filename) = item.strip_prefix(&format!("{}/", self.type_folder))
            {
                return Some(Ok(filename.to_string()));
            }
        }

        None
    }
}

/// A [Source] that loads packages from zip/pk3 files in the filesystem. This is a good source to
///  use for release builds. Each zip/pk3 file is a different package.
pub struct ZipFolderSource
{
    basedir : String,
    extension : String,
    package_list : Option<Vec<String>>,
    loaded_zips : HashMap<String, ZipArchive<BufReader<File>>, BuildHasherDefault<FxHasher>>,
    trust : TrustLevel
}

impl ZipFolderSource
{
    /// Creates a new [ZipFolderSource] using the given directory to look for packages in.
    pub fn new(directory : &str, trust : TrustLevel) -> ZipFolderSource
    {
        ZipFolderSource
        {
            basedir : directory.to_string(),
            extension : String::from(".zip"),
            package_list : None,
            loaded_zips : HashMap::with_hasher(BuildHasherDefault::<FxHasher>::default()),
            trust
        }
    }

    /// Creates a new [ZipFolderSource] using the given directory to look for packages in, and
    ///  the given file extension to look for zip files by package name. By default, new() will
    ///  use "zip" as the extension.
    pub fn with_extension(directory : &str, extension : &str, trust : TrustLevel) -> ZipFolderSource
    {
        let mut dot_extension = String::from(".");
        dot_extension.push_str(extension);

        ZipFolderSource
        {
            basedir : directory.to_string(),
            extension : dot_extension,
            package_list : None,
            loaded_zips : HashMap::with_hasher(BuildHasherDefault::<FxHasher>::default()),
            trust
        }
    }

    /// Returns true if there is an open file handle to the zip file for the given package name.
    pub fn zip_loaded(&self, package_name : &str) -> bool
    {
        self.loaded_zips.contains_key(package_name)
    }

    /// Opens an OS file handle to the zip file for the given package name.
    pub fn load_zip(&mut self, package_name : &str) -> Result<(), PackageError>
    {
        if self.zip_loaded(package_name)
        {
            return Ok(());
        }

        let mut path = PathBuf::from(&self.basedir);
        let mut zip_name = String::from(package_name);
        zip_name.push_str(&self.extension);
        path.push(zip_name);

        let file = BufReader::new(File::open(path).map_err(
            |error| PackageError::IoError(error))?);

        let archive = ZipArchive::new(file).map_err(
            |error| package_error_for_zip_error(error))?;

        self.loaded_zips.insert(package_name.to_string(), archive);

        Ok(())
    }

    /// Closes the OS file handle to the zip file for the given package name.
    pub fn unload_zip(&mut self, package_name : &str)
    {
        self.loaded_zips.remove(package_name);
    }
}

impl Source for ZipFolderSource
{
    fn get_uri(&self) -> &str
    {
        &self.basedir
    }

    fn has_package(&self, package_name: &str) -> bool
    {
        let mut path = PathBuf::from(&self.basedir);
        let mut zip_name = String::from(package_name);
        zip_name.push_str(&self.extension);
        path.push(zip_name);

        path.exists()
    }

    fn list_packages(&mut self) -> Vec<String>
    {
        if let Some(package_list) = &self.package_list
        {
            return package_list.clone();
        }

        let mut package_list: Vec<String> = Vec::new();
        let path = PathBuf::from(&self.basedir);

        let walkdir = WalkDir::new(path).min_depth(1).max_depth(1).into_iter();

        for dir_entry in walkdir
        {
            if let Ok(entry) = dir_entry
            {
                if !entry.file_type().is_file()
                {
                    continue;
                }

                if let Some(stripped) = entry.file_name().to_string_lossy()
                    .strip_suffix(&self.extension)
                {
                    package_list.push(stripped.to_string());
                }
            }
        }

        package_list.sort();

        self.package_list = Some(package_list.clone());

        package_list
    }

    fn read_file<'a>(&'a mut self, package_name: &str, pathname: &str) -> Result<Box<dyn ReadSeek + 'a>, PackageError>
    {
        self.load_zip(package_name)?;

        if let Some(archive) = self.loaded_zips.get_mut(package_name)
        {
            match archive.by_name(pathname)
            {
                Ok(reader) => { return Ok(Box::new(FakeSeekReader::new(reader))); }
                Err(error) => { return Err(package_error_for_zip_error(error)); }
            }
        }

        Err(PackageError::IoError(std::io::Error::new(std::io::ErrorKind::NotFound, "Archive not found")))
    }

    fn iter_entries<'a>(&'a mut self, package_name: &str, type_folder: &str) -> Box<dyn Iterator<Item=Result<String, PackageError>> + 'a>
    {
        if self.load_zip(package_name).is_err()
        {
            return Box::new(EmptyEntryIter::new());
        }

        let mut iter : Box<dyn Iterator<Item=Result<String, PackageError>>> = Box::new(EmptyEntryIter::new());

        if let Some(archive) = self.loaded_zips.get(package_name)
        {
            let archive_iter = archive.file_names();

            iter = Box::new(ZipFolderIter::new(archive_iter, type_folder));
        }

        iter
    }

    fn trust_level(&self, _package_name: &str) -> TrustLevel
    {
        self.trust
    }
}