1use std::path::{Path, PathBuf};
2
3use itertools::Itertools;
4use thiserror::Error;
5use crate::resource::partition_manager::PartitionManagerError::PartitionNotFound;
6
7use crate::resource::pdefs::{
8 GameDiscoveryError, GamePaths, PackageDefinitionError, PackageDefinitionSource, PartitionId,
9 PartitionInfo,
10};
11use crate::resource::resource_info::ResourceInfo;
12use crate::resource::runtime_resource_id::RuntimeResourceID;
13use crate::WoaVersion;
14
15use super::resource_partition::{PatchId, ResourcePartition, ResourcePartitionError};
16
17#[derive(Debug, Error)]
18pub enum PartitionManagerError {
19 #[error("Cannot use packagedefinition config: {0}")]
20 PackageDefinitionError(#[from] PackageDefinitionError),
21
22 #[error("partition '{0}' error: {1}")]
23 PartitionError(PartitionId, ResourcePartitionError),
24
25 #[error("partition {0} could not be found")]
26 PartitionNotFound(String),
27
28 #[error("resource {0} could not be found")]
29 ResourceNotFound(String),
30
31 #[error("Could not discover game paths: {0}")]
32 GameDiscoveryError(#[from] GameDiscoveryError),
33
34 #[error("Could not find a root partition")]
35 NoRootPartition(),
36}
37
38#[allow(dead_code)]
39#[derive(Clone, Debug, Copy)]
40pub struct PartitionState {
41 pub installing: bool,
42 pub mounted: bool,
43 pub install_progress: f32,
44}
45
46pub struct PartitionManager {
47 runtime_directory: PathBuf,
48 partition_infos: Vec<PartitionInfo>, pub partitions: Vec<ResourcePartition>, }
51
52impl PartitionManager {
53 pub fn new(
59 runtime_directory: PathBuf,
60 package_definition: &PackageDefinitionSource,
61 ) -> Result<Self, PartitionManagerError> {
62 let partition_infos = package_definition
63 .read()
64 .map_err(PartitionManagerError::PackageDefinitionError)?;
65
66 Ok(Self {
67 runtime_directory,
68 partition_infos,
69 partitions: vec![],
70 })
71 }
72
73 pub fn from_game(
80 retail_directory: PathBuf,
81 game_version: WoaVersion,
82 mount: bool,
83 ) -> Result<Self, PartitionManagerError> {
84 Self::from_game_with_callback(retail_directory, game_version, mount, |_, _| {})
85 }
86
87 pub fn from_game_with_callback<F>(
95 retail_directory: PathBuf,
96 game_version: WoaVersion,
97 mount: bool,
98 progress_callback: F,
99 ) -> Result<Self, PartitionManagerError>
100 where
101 F: FnMut(usize, &PartitionState),
102 {
103 let game_paths = GamePaths::from_retail_directory(retail_directory)?;
104 let package_definition =
105 PackageDefinitionSource::from_file(game_paths.package_definition_path, game_version)?;
106
107 let partition_infos = package_definition
109 .read()
110 .map_err(PartitionManagerError::PackageDefinitionError)?;
111
112 let mut package_manager = Self {
113 runtime_directory: game_paths.runtime_path,
114 partition_infos,
115 partitions: vec![],
116 };
117
118 if mount {
120 package_manager.mount_partitions(progress_callback)?;
121 }
122
123 Ok(package_manager)
124 }
125
126 fn try_read_partition<F>(
127 runtime_directory: &Path,
128 partition_info: PartitionInfo,
129 mut progress_callback: F,
130 ) -> Result<Option<ResourcePartition>, PartitionManagerError>
131 where
132 F: FnMut(&PartitionState),
133 {
134 let mut partition = ResourcePartition::new(partition_info.clone());
135 let mut state_result: PartitionState = PartitionState {
136 installing: false,
137 mounted: false,
138 install_progress: 0.0,
139 };
140
141 let callback = |state: &_| {
142 progress_callback(state);
143 state_result = *state;
144 };
145
146 partition
147 .mount_resource_packages_in_partition_with_callback(runtime_directory, callback)
148 .map_err(|e| PartitionManagerError::PartitionError(partition_info.id, e))?;
149
150 if state_result.mounted {
151 Ok(Some(partition))
152 } else {
153 Ok(None)
154 }
155 }
156
157 pub fn mount_partitions<F>(
162 &mut self,
163 mut progress_callback: F,
164 ) -> Result<(), PartitionManagerError>
165 where
166 F: FnMut(usize, &PartitionState),
167 {
168 let partitions = self
169 .partition_infos
170 .iter()
171 .enumerate()
172 .map(|(index, partition_info)| {
173 let callback = |state: &_| {
174 progress_callback(index + 1, state);
175 };
176
177 Self::try_read_partition(&self.runtime_directory, partition_info.clone(), callback)
178 })
179 .collect::<Result<Vec<Option<ResourcePartition>>, PartitionManagerError>>()?
180 .into_iter()
181 .flatten()
182 .collect::<Vec<ResourcePartition>>();
183
184 for partition in partitions {
185 self.partitions.push(partition);
186 }
187
188 Ok(())
189 }
190
191 pub fn mount_partition<F>(
197 &mut self,
198 partition_info: PartitionInfo,
199 progress_callback: F,
200 ) -> Result<(), PartitionManagerError>
201 where
202 F: FnMut(&PartitionState),
203 {
204 if let Some(partition) =
205 Self::try_read_partition(&self.runtime_directory, partition_info, progress_callback)?
206 {
207 self.partitions.push(partition)
208 }
209
210 Ok(())
211 }
212
213 pub fn read_resource_from(
214 &self,
215 partition_id: PartitionId,
216 rrid: RuntimeResourceID,
217 ) -> Result<Vec<u8>, PartitionManagerError> {
218 let partition = self
219 .partitions
220 .iter()
221 .find(|partition| partition.partition_info().id == partition_id);
222
223 if let Some(partition) = partition {
224 match partition.read_resource(&rrid) {
225 Ok(data) => Ok(data),
226 Err(e) => Err(PartitionManagerError::PartitionError(partition_id, e)),
227 }
228 } else {
229 Err(PartitionManagerError::PartitionNotFound(
230 partition_id.to_string(),
231 ))
232 }
233 }
234
235 pub fn find_partition(&self, partition_id: PartitionId) -> Option<&ResourcePartition> {
236 self.partitions
237 .iter()
238 .find(|partition| partition.partition_info().id == partition_id)
239 }
240
241 pub fn root_partition(
242 &self
243 ) -> Result<PartitionId, PartitionManagerError> {
244 if let Some(mut partition) = self.partition_infos.first(){
245 loop {
246 match &partition.parent{
247 Some(parent) => {
248 match self.find_partition(parent.clone()){
249 Some(part) => {partition = part.partition_info()}
250 None => {return Err(PartitionNotFound(parent.to_string()))}
251 };
252 },
253 None => return Ok(partition.id.clone()),
254 }
255 }
256 }
257 Err(PartitionManagerError::NoRootPartition())
258 }
259
260 pub fn partitions_with_resource(&self, rrid: &RuntimeResourceID) -> Vec<PartitionId> {
261 self.partitions
262 .iter()
263 .filter_map(|partition| {
264 if partition.contains(rrid) {
265 Some(partition.partition_info().id.clone())
266 } else {
267 None
268 }
269 })
270 .collect()
271 }
272
273 pub fn iter_all_runtime_resource_ids(&self) -> impl Iterator<Item = &RuntimeResourceID> + '_ {
278 self.partitions.iter().flat_map(|partition| {
279 partition.resources.keys()
280 })
281 }
282
283 pub fn resource_mounted(&self, rrid: &RuntimeResourceID) -> bool{
284 self.iter_all_runtime_resource_ids().contains(rrid)
285 }
286
287 pub fn resource_infos(&self, rrid: &RuntimeResourceID) -> Vec<(PartitionId, &ResourceInfo)> {
288 self.partitions_with_resource(rrid)
289 .into_iter()
290 .filter_map(|p_id| {
291 if let Ok(info) = self.resource_info_from(&p_id, rrid) {
292 Some((p_id, info))
293 } else {
294 None
295 }
296 })
297 .collect()
298 }
299
300 pub fn resource_info_from(
301 &self,
302 partition_id: &PartitionId,
303 rrid: &RuntimeResourceID,
304 ) -> Result<&ResourceInfo, PartitionManagerError> {
305 let partition = self
306 .partitions
307 .iter()
308 .find(|partition| partition.partition_info().id == *partition_id);
309
310 if let Some(partition) = partition {
311 match partition.get_resource_info(rrid) {
312 Ok(info) => Ok(info),
313 Err(e) => Err(PartitionManagerError::PartitionError(partition_id.clone(), e)),
314 }
315 } else {
316 Err(PartitionManagerError::PartitionNotFound(
317 partition_id.to_string(),
318 ))
319 }
320 }
321
322 pub fn resolve_resource_from(
328 &self,
329 partition_id: PartitionId,
330 resource_id: &RuntimeResourceID,
331 ) -> Result<(&ResourceInfo, PartitionId), PartitionManagerError> {
332 match self.find_partition(partition_id.clone()) {
333 Some(partition) => {
334 if partition.contains(resource_id) {
335 match partition.get_resource_info(resource_id) {
336 Ok(info) => Ok((info, partition_id.clone())),
337 Err(_) => Err(PartitionManagerError::ResourceNotFound(resource_id.to_string())),
338 }
339 } else {
340 match &partition.partition_info().parent {
341
342 Some(parent_id) => {
343 self.resolve_resource_from(parent_id.clone(), resource_id)
344 },
345 None => {
346 Err(PartitionManagerError::ResourceNotFound(resource_id.to_string()))
347 }
348 }
349 }
350 },
351 None => {
352 Err(PartitionNotFound(partition_id.to_string()))
353 }
354 }
355 }
356 #[deprecated(
357 since = "1.0.0",
358 note = "prefer direct access through the partitions field"
359 )]
360 pub fn partitions(&self) -> &Vec<ResourcePartition> {
361 &self.partitions
362 }
363
364 #[deprecated(
365 since = "1.1.0",
366 note = "please implement this yourself, it is out of scope for this struct"
367 )]
368 pub fn print_resource_changelog(&self, rrid: &RuntimeResourceID) {
369 println!("Resource: {rrid}");
370
371 for partition in &self.partitions {
372 let mut last_occurence: Option<&ResourceInfo> = None;
373
374 let size =
375 |info: &ResourceInfo| info.compressed_size().unwrap_or(info.header.data_size);
376
377 let changes = partition.resource_patch_indices(rrid);
378 let deletions = partition.resource_removal_indices(rrid);
379 let occurrences = changes
380 .clone()
381 .into_iter()
382 .chain(deletions.clone().into_iter())
383 .collect::<Vec<PatchId>>();
384
385 for occurence in occurrences.iter().sorted() {
386 println!(
387 "{}: {}",
388 match occurence {
389 PatchId::Base => {
390 "Base"
391 }
392 PatchId::Patch(_) => {
393 "Patch"
394 }
395 },
396 partition.partition_info().filename(*occurence)
397 );
398
399 if deletions.contains(occurence) {
400 println!("\t- Removal: resource deleted");
401 last_occurence = None;
402 }
403
404 if changes.contains(occurence) {
405 if let Ok(info) = partition.resource_info_from(rrid, *occurence) {
406 if let Some(last_info) = last_occurence {
407 println!(
408 "\t- Modification: Size changed from {} to {}",
409 size(last_info),
410 size(info)
411 );
412 } else {
413 println!("\t- Addition: New occurrence, Size {} bytes", size(info))
414 }
415 last_occurence = Some(info);
416 }
417 }
418 }
419 }
420 }
421}