use std::{
collections::HashSet,
fs::{self, OpenOptions},
io::{self, Write as _},
sync::Mutex,
};
use camino::{Utf8Path, Utf8PathBuf};
use ltk_file::LeagueFileKind;
use crate::{ChunkDecoder, SubchunkToc, WadChunk, WadError};
use super::{
lock,
naming::{hashed_name, is_path_conflict, ltk_path, DirectoryPaths, NamingPolicy},
report::{ChunkOutcome, WriteOutcome},
Job, PathIssue,
};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
#[non_exhaustive]
pub enum ExtractLayout {
#[default]
Paths,
Flat,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
#[non_exhaustive]
pub enum ExistingFilePolicy {
#[default]
Overwrite,
Skip,
}
pub(super) struct ChunkWriter<'s> {
pub(super) layout: ExtractLayout,
pub(super) existing: ExistingFilePolicy,
pub(super) naming: NamingPolicy,
pub(super) type_filter: Option<&'s [LeagueFileKind]>,
pub(super) output_dir: &'s Utf8Path,
pub(super) directories: DirectoryPaths,
pub(super) output_occupied: bool,
pub(super) claimed: Mutex<HashSet<Utf8PathBuf>>,
pub(super) created: Mutex<HashSet<Utf8PathBuf>>,
pub(super) subchunk_toc: Option<SubchunkToc>,
}
impl ChunkWriter<'_> {
pub(super) fn write(
&self,
job: &Job,
decoder: &mut ChunkDecoder,
) -> Result<WriteOutcome, WadError> {
let data = decoder.decompress_chunk(&job.raw, &job.chunk, self.subchunk_toc.as_ref())?;
self.write_chunk(&job.chunk, &data, Utf8Path::new(&job.path), job.named)
}
pub(super) fn probe_skip(
&self,
chunk: &WadChunk,
path: &str,
named: bool,
) -> Option<WriteOutcome> {
if self.existing != ExistingFilePolicy::Skip
|| self.type_filter.is_some()
|| self.layout != ExtractLayout::Paths
{
return None;
}
let (relative, renamed) = self.final_path(chunk, Utf8Path::new(path), named, None)?;
let mut claimed = lock(&self.claimed);
if claimed.contains(&relative) {
return None;
}
if !self.output_dir.join(&relative).is_file() {
return None;
}
claimed.insert(relative.clone());
let issue = renamed.then(|| PathIssue::Renamed(relative.clone()));
Some(WriteOutcome {
outcome: ChunkOutcome::SkippedExisting,
path: Some(relative),
issue,
})
}
fn write_chunk(
&self,
chunk: &WadChunk,
chunk_data: &[u8],
chunk_path: &Utf8Path,
named: bool,
) -> Result<WriteOutcome, WadError> {
let chunk_kind = LeagueFileKind::identify_from_bytes(chunk_data);
if self
.type_filter
.is_some_and(|types| !types.contains(&chunk_kind))
{
return Ok(WriteOutcome::skipped(ChunkOutcome::SkippedByType, None));
}
let (mut relative_path, mut renamed) = match self.layout {
ExtractLayout::Paths => {
let (path, renamed) = self.resolve_final_path(chunk, chunk_path, named, chunk_kind);
match self.claim_path(path, renamed, chunk, chunk_kind) {
Some(claimed) => claimed,
None => {
return Ok(WriteOutcome::skipped(
ChunkOutcome::SkippedDuplicatePath,
Some(PathIssue::Duplicate),
))
}
}
}
ExtractLayout::Flat => self.resolve_flat_path(chunk, chunk_path, named, chunk_kind),
};
let full_path = self.output_dir.join(&relative_path);
let written = match self.place(&relative_path, chunk_data) {
Ok(written) => written,
Err(error) if is_path_conflict(&error, &full_path) => {
relative_path = hashed_name(chunk, chunk_kind, self.naming);
renamed = true;
lock(&self.claimed).insert(relative_path.clone());
self.place(&relative_path, chunk_data)?
}
Err(error) => return Err(WadError::IoError(error)),
};
let outcome = match written {
Written::Yes => ChunkOutcome::Written {
kind: chunk_kind,
bytes: chunk_data.len() as u64,
},
Written::Existed => ChunkOutcome::SkippedExisting,
};
let issue = renamed.then(|| PathIssue::Renamed(relative_path.clone()));
Ok(WriteOutcome {
outcome,
path: Some(relative_path),
issue,
})
}
fn place(&self, relative: &Utf8Path, data: &[u8]) -> io::Result<Written> {
let full_path = self.output_dir.join(relative);
if let Some(parent) = full_path.parent() {
self.create_dirs(parent)?;
}
write_file(&full_path, data, self.existing)
}
fn create_dirs(&self, parent: &Utf8Path) -> io::Result<()> {
if lock(&self.created).contains(parent) {
return Ok(());
}
fs::create_dir_all(parent)?;
let mut created = lock(&self.created);
for ancestor in parent.ancestors() {
if !ancestor.starts_with(self.output_dir) || !created.insert(ancestor.to_path_buf()) {
break;
}
}
Ok(())
}
fn resolve_final_path(
&self,
chunk: &WadChunk,
chunk_path: &Utf8Path,
named: bool,
chunk_kind: LeagueFileKind,
) -> (Utf8PathBuf, bool) {
self.final_path(chunk, chunk_path, named, Some(chunk_kind))
.expect("a path resolves whenever the chunk's kind is known")
}
fn final_path(
&self,
chunk: &WadChunk,
chunk_path: &Utf8Path,
named: bool,
chunk_kind: Option<LeagueFileKind>,
) -> Option<(Utf8PathBuf, bool)> {
let mut final_path = chunk_path.to_path_buf();
if !named && self.naming == NamingPolicy::Descriptive {
if let Some(ext) = chunk_kind?.extension() {
final_path.set_extension(ext);
}
}
if !self.directories.holds(final_path.as_str())
&& !(self.output_occupied && self.output_dir.join(&final_path).is_dir())
{
return Some((final_path, false));
}
let renamed = ltk_path(&final_path);
if self.directories.holds(renamed.as_str()) {
let kind = match self.naming {
NamingPolicy::Lossless => chunk_kind.unwrap_or(LeagueFileKind::Unknown),
_ => chunk_kind?,
};
return Some((hashed_name(chunk, kind, self.naming), true));
}
Some((renamed, true))
}
fn claim_path(
&self,
path: Utf8PathBuf,
renamed: bool,
chunk: &WadChunk,
chunk_kind: LeagueFileKind,
) -> Option<(Utf8PathBuf, bool)> {
let mut claimed = lock(&self.claimed);
if claimed.insert(path.clone()) {
return Some((path, renamed));
}
if self.naming != NamingPolicy::Lossless {
return None;
}
let suffixed = ltk_path(&path);
if claimed.insert(suffixed.clone()) {
return Some((suffixed, true));
}
let hashed = hashed_name(chunk, chunk_kind, self.naming);
claimed.insert(hashed.clone());
Some((hashed, true))
}
fn resolve_flat_path(
&self,
chunk: &WadChunk,
chunk_path: &Utf8Path,
named: bool,
chunk_kind: LeagueFileKind,
) -> (Utf8PathBuf, bool) {
let file_name = Utf8Path::new(chunk_path.file_name().unwrap_or_default());
let (resolved, renamed) = self.resolve_final_path(chunk, file_name, named, chunk_kind);
let mut claimed = lock(&self.claimed);
if claimed.insert(resolved.clone()) {
return (resolved, renamed);
}
let suffixed = Utf8PathBuf::from(match resolved.extension() {
Some(ext) => format!(
"{}.{:016x}.{ext}",
resolved.file_stem().unwrap_or_default(),
chunk.path_hash
),
None => format!("{}.{:016x}", resolved.as_str(), chunk.path_hash),
});
claimed.insert(suffixed.clone());
(suffixed, renamed)
}
}
enum Written {
Yes,
Existed,
}
fn write_file(path: &Utf8Path, data: &[u8], policy: ExistingFilePolicy) -> io::Result<Written> {
match policy {
ExistingFilePolicy::Overwrite => {
fs::write(path, data)?;
Ok(Written::Yes)
}
ExistingFilePolicy::Skip => {
match OpenOptions::new().write(true).create_new(true).open(path) {
Ok(mut file) => {
file.write_all(data)?;
Ok(Written::Yes)
}
Err(error) if error.kind() == io::ErrorKind::AlreadyExists && !path.is_dir() => {
Ok(Written::Existed)
}
Err(error) => Err(error),
}
}
}
}