yok 0.1.4

use std::{
    fs, io,
    path::{Path, PathBuf},
};

fn main() {
    use std::mem;

    #[derive(Debug, Clone, PartialEq, Eq)]
    pub struct DirEntry {
        path: String,
        is_dir: bool,
        is_file: bool,
        content: Vec<u8>,
    }

    #[derive(Debug, Clone, PartialEq, Eq)]
    pub struct Dir {
        data: Vec<DirEntry>,
    }

    // Implement as_bytes() for Dir
    impl Dir {
        pub fn as_bytes(&self) -> Vec<u8> {
            // Calculate the total size of the serialized data
            let mut size = mem::size_of::<u32>(); // Size of the length prefix

            for entry in &self.data {
                size += mem::size_of::<u32>(); // Size of the entry size prefix
                size += entry.path.as_bytes().len();
                size += mem::size_of::<bool>() * 2; // Sizes of is_dir and is_file fields
                size += mem::size_of::<u32>(); // Size of content size prefix
                size += entry.content.len();
            }

            let mut bytes = Vec::with_capacity(size);

            // Write the length prefix
            bytes.extend((self.data.len() as u32).to_le_bytes().iter());

            for entry in &self.data {
                // Write the size prefix for this entry
                bytes.extend((entry.path.as_bytes().len() as u32).to_le_bytes().iter());

                // Write the entry data
                bytes.extend(entry.path.as_bytes().iter());
                bytes.extend(&[entry.is_dir as u8, entry.is_file as u8]);
                bytes.extend((entry.content.len() as u32).to_le_bytes().iter());
                bytes.extend(entry.content.iter());
            }

            bytes
        }
    }

    // Implement from_bytes() for Dir
    impl Dir {
        #[allow(warnings)]
        pub fn from_bytes(bytes: &[u8]) -> Option<Self> {
            let mut cursor = 0;
            let len_prefix_size = mem::size_of::<u32>();

            if bytes.len() < len_prefix_size {
                return None;
            }

            let len = u32::from_le_bytes([
                bytes[cursor],
                bytes[cursor + 1],
                bytes[cursor + 2],
                bytes[cursor + 3],
            ]) as usize;

            cursor += len_prefix_size;

            let mut data = Vec::with_capacity(len);

            for _ in 0..len {
                if bytes.len() < cursor + len_prefix_size {
                    return None;
                }

                let entry_size = u32::from_le_bytes([
                    bytes[cursor],
                    bytes[cursor + 1],
                    bytes[cursor + 2],
                    bytes[cursor + 3],
                ]) as usize;

                cursor += len_prefix_size;

                let end_pos = cursor + entry_size;

                if bytes.len() < end_pos {
                    return None;
                }

                let path = String::from_utf8(bytes[cursor..end_pos].to_vec()).ok()?;
                cursor = end_pos;

                if bytes.len() < cursor + 2 {
                    return None;
                }

                let is_dir = bytes[cursor] != 0;
                let is_file = bytes[cursor + 1] != 0;
                cursor += 2;

                if bytes.len() < cursor + len_prefix_size {
                    return None;
                }

                let content_size = u32::from_le_bytes([
                    bytes[cursor],
                    bytes[cursor + 1],
                    bytes[cursor + 2],
                    bytes[cursor + 3],
                ]) as usize;

                cursor += len_prefix_size;

                if bytes.len() < cursor + content_size {
                    return None;
                }

                let content = bytes[cursor..cursor + content_size].to_vec();
                cursor += content_size;

                data.push(DirEntry {
                    path,
                    is_dir,
                    is_file,
                    content,
                });
            }

            Some(Dir { data })
        }
    }

    pub fn walk_dir(path: impl ToString) -> io::Result<Vec<PathBuf>> {
        let path = path.to_string();
        let dir = Path::new(&path);
        let mut files = Vec::new();
        for entry in fs::read_dir(dir)? {
            let entry = entry?;
            let newpath = entry.path();

            if newpath.is_dir() {
                files.extend(walk_dir(&newpath.display())?);
                files.push(newpath);
            } else {
                files.push(newpath);
            }
        }
        Ok(files)
    }

    use std::env;
    let yok_path = match env::var("YOK_PATH") {
        Ok(value) => value,
        Err(_) => ".".to_string(),
    };
    println!("cargo:rerun-if-env-changed={}", yok_path);
    println!("cargo:rerun-if-changed={}", yok_path);
    println!("cargo:rerun-if-changed=main.rs");
    println!("cargo:rerun-if-changed=lib.rs");
    let yok_path = std::path::Path::new(&yok_path);
    if yok_path.is_dir() {
        let mut res_dir: Vec<DirEntry> = vec![];
        let dir_list = walk_dir(yok_path.display()).expect("walk dir error");
        for dir in dir_list {
            if dir.is_file() {
                res_dir.push(DirEntry {
                    path: dir.display().to_string(),
                    is_dir: false,
                    is_file: true,
                    content: std::fs::read(dir).expect("read dir error"),
                });
            } else if dir.is_dir() {
                res_dir.push(DirEntry {
                    path: dir.display().to_string(),
                    is_dir: true,
                    is_file: false,
                    content: vec![],
                });
            }
        }
        let all_dir: Dir = Dir { data: res_dir };
        let all_dir_bytes = all_dir.as_bytes();
        std::fs::write(format!("../.yok"), all_dir_bytes).unwrap()
    } else {
        eprintln!("Environment variable 'YOK_PATH' not dir");
        return;
    }
}