1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145

use byteorder::{LittleEndian, ReadBytesExt};
use rayon::prelude::*;

use std::collections::BTreeMap;
use std::io::{Cursor, Read};
use std::path::Path;

use anyhow::Result;

pub mod fat;
pub mod fnt;

use self::fat::FileAllocTable;
use self::fnt::{Directory, FileEntry, ROOT_ID};

/// Represents a NitroROM file system.
#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub struct FileSystem {
    pub dirs: BTreeMap<u16, Directory>,
    overlays: Vec<FileEntry>,
}

impl FileSystem {
    pub fn new(fnt: &[u8], fat: &[u8]) -> Result<Self> {
        let mut cursor = Cursor::new(fnt);
        let mut dirs = BTreeMap::new();

        cursor.set_position(6);
        let count = cursor.read_u16::<LittleEndian>()?;

        cursor.set_position(0);

        for index in 0..count {
            let id = ROOT_ID + index;
            dirs.insert(id, Directory::new(&mut cursor, id)?);
        }

        let fat = FileAllocTable::new(fat)?;

        let mut fnt = Self { 
            dirs,
            overlays: Vec::new(),
        };

        fnt.populate(&mut cursor, &fat)?;

        Ok(fnt)
    }

    /// How many directories there are
    pub fn count(&self) -> usize {
        self.dirs.len()
    }
    
    /// Get a Vec of all files
    pub fn files(&self) -> Vec<&FileEntry> {
        self.dirs
            .par_iter()
            .flat_map(|(_, ref dir)| {
                &dir.files
            })
            .collect::<_>()
    }

    /// The lowest ID in the File System. Any ID lower than this in 
    /// the FAT is an overlay file.
    pub fn start_id(&self) -> u16 {
        self.dirs[&ROOT_ID].start_id()
    }

    /// Get a Vec of all overlays
    pub fn overlays(&self) -> &[FileEntry] {
        &self.overlays
    }

    fn populate(&mut self, cursor: &mut Cursor<&[u8]>, fat: &FileAllocTable) -> Result<()> {
        self._populate(cursor, "", ROOT_ID, fat)?;

        self.overlays = (0..self.start_id())
            .into_par_iter()
            .map(|id| {
                let alloc_info = fat.get(id).unwrap();
                FileEntry::new(id, &format!("overlay_{:04}", id), alloc_info)
            })
            .collect::<_>();

        Ok(())
    }

    fn _populate<P: AsRef<Path>>(&mut self, mut cursor: &mut Cursor<&[u8]>, path: P, id: u16, fat: &FileAllocTable) -> Result<()> {
        let mut file_id = {
            let dir = self.dirs.get_mut(&id).unwrap();
            dir.set_path(&path);
            cursor.set_position(dir.offset() as u64);
            dir.start_id()
        };

        let mut files = Vec::new();

        let mut len = cursor.read_u8()?;

        while len != 0 {
            let name = self.read_name(&mut cursor, len)?;

            if len > 0x80 {
                //  Read the directory ID that this name goes to
                let dir_id = cursor.read_u16::<LittleEndian>()?;

                let pos = cursor.position();
                let new_path = path.as_ref().join(name);
                
                self._populate(&mut cursor, new_path, dir_id, fat)?;

                cursor.set_position(pos);
            } else {
                let file_path = path.as_ref().join(name);
                let alloc_info = fat.get(file_id).unwrap();

                files.push(FileEntry::new(file_id, &file_path, alloc_info));
                file_id += 1;
            }

            len = cursor.read_u8()?;
        }

        let dir = self.dirs.get_mut(&id).unwrap();

        dir.append_files(&files);

        Ok(())
    }

    fn read_name<R: Read>(&self, cursor: &mut R, mut len: u8) -> Result<String> {
        let mut name = String::new();

        if len > 0x80 {
            len -= 0x80;
        }

        cursor.take(u64::from(len))
            .read_to_string(&mut name)?;

        Ok(name)
    }
}