codec.rs

use std::{
    fs::File,
    io::{BufReader, BufWriter},
    path::Path,
};

use log::debug;

use crate::{
    archiver::Archiver,
    datamodel::{BlockPool, FileTable, IndexTable},
    unarchiver::UnArchiver,
};

pub const BLOCK_SIZE: usize = 128 * 1024;

#[derive(Default)]
pub struct ArCodec {
    pub(crate) index_table: IndexTable,
    pub(crate) file_table: FileTable,
    pub(crate) buffer: Vec<u8>,
}

impl ArCodec {
    pub fn new() -> Self {
        Self {
            buffer: vec![0; BLOCK_SIZE],
            ..Self::default()
        }
    }

    pub fn with_block_size(mut self, block_size: u32) -> Self {
        assert!(block_size <= u32::MAX / 2);
        self.buffer.resize(block_size as usize, 0);

        self
    }

    pub fn with_index_table(mut self, index_table: IndexTable) -> Self {
        self.index_table = index_table;
        self
    }

    pub fn with_file_table(mut self, file_table: FileTable) -> Self {
        self.file_table = file_table;
        self
    }

    pub fn with_index_table_from_path(
        mut self,
        it_path: impl AsRef<Path>,
    ) -> Result<Self, std::io::Error> {
        debug!(
            "Loading IndexTable from path '{}'",
            it_path.as_ref().display()
        );

        let it_file = File::open(it_path.as_ref())?;
        let it_file = BufReader::new(it_file);

        self.index_table = IndexTable::try_deserialize_from(it_file)?;
        self.index_table.file_path = Some(it_path.as_ref().to_path_buf());

        Ok(self)
    }

    pub fn with_index_table_from_path_or_create(
        mut self,
        it_path: impl AsRef<Path>,
    ) -> Result<Self, std::io::Error> {
        if it_path.as_ref().exists() {
            self.with_index_table_from_path(it_path)
        } else {
            debug!(
                "Creating new IndexTable for path '{}'",
                it_path.as_ref().display()
            );

            self.index_table = IndexTable::default();
            self.index_table.file_path = Some(it_path.as_ref().to_path_buf());
            Ok(self)
        }
    }

    pub fn with_file_table_from_path(
        mut self,
        ft_path: impl AsRef<Path>,
    ) -> Result<Self, std::io::Error> {
        debug!(
            "Loading FileTable from path '{}'",
            ft_path.as_ref().display()
        );

        let ft_file = File::open(ft_path.as_ref())?;
        let ft_file = BufReader::new(ft_file);

        let mut decoder = zstd::Decoder::new(ft_file)?;

        self.file_table = FileTable::try_deserialize_from(&mut decoder)?;
        self.file_table.file_path = Some(ft_path.as_ref().to_path_buf());

        Ok(self)
    }

    pub fn with_file_table_from_path_or_create(
        mut self,
        ft_path: impl AsRef<Path>,
    ) -> Result<Self, std::io::Error> {
        if ft_path.as_ref().exists() {
            self.with_file_table_from_path(ft_path)
        } else {
            debug!(
                "Creating new FileTable for path '{}'",
                ft_path.as_ref().display()
            );

            self.file_table = FileTable::default();
            self.file_table.file_path = Some(ft_path.as_ref().to_path_buf());
            Ok(self)
        }
    }

    pub fn archiver_from_bp_path(
        self,
        blocks_out: impl AsRef<Path>,
    ) -> std::io::Result<Archiver<File>> {
        let bpool = BlockPool::new_from_path_or_create(blocks_out)?;
        Ok(Archiver::new(self, bpool))
    }

    pub fn unarchiver_from_bp_path(
        self,
        blocks_in: impl AsRef<Path>,
    ) -> std::io::Result<UnArchiver<File>> {
        let bpool = BlockPool::new_from_path_readonly(blocks_in)?;
        Ok(UnArchiver::new(self, bpool))
    }

    pub fn write_file_table_to_path(
        &self,
        out_file: impl AsRef<Path>,
    ) -> Result<(), std::io::Error> {
        debug!(
            "Writing FileTable to file path '{}'",
            out_file.as_ref().display()
        );

        let out_file = File::create(out_file.as_ref())?;
        let out_file = BufWriter::new(out_file);

        let mut encoder = zstd::Encoder::new(out_file, 2)?;

        self.file_table.serialize_to(&mut encoder)?;

        encoder.finish()?;
        Ok(())
    }

    pub fn write_index_table_to_path(
        &self,
        out_file: impl AsRef<Path>,
    ) -> Result<(), std::io::Error> {
        debug!(
            "Writing IndexTable to file path '{}'",
            out_file.as_ref().display()
        );

        let out_file = File::create(out_file.as_ref())?;
        let out_file = BufWriter::new(out_file);
        self.index_table.serialize_to(out_file)
    }
}