dtmt/lib/sdk/src/bundle/file.rs

use std::io::{Cursor, Read, Seek, Write};
use std::path::Path;

use bitflags::bitflags;
use color_eyre::eyre::Context;
use color_eyre::{eyre, Result};
use futures::future::join_all;

use crate::binary::sync::*;
use crate::filetype::*;
use crate::murmur::{HashGroup, IdString64, Murmur64};

use super::filetype::BundleFileType;

#[derive(Debug)]
struct BundleFileHeader {
    variant: u32,
    unknown_1: u8,
    size: usize,
    len_data_file_name: usize,
}

#[derive(Clone, Debug)]
pub struct BundleFileVariant {
    property: u32,
    data: Vec<u8>,
    data_file_name: Option<String>,
    // Seems to be related to whether there is a data path.
    unknown_1: u8,
}

impl BundleFileVariant {
    // We will need a parameter for `property` eventually, so the `Default` impl would need to go
    // eventually anyways.
    #[allow(clippy::new_without_default)]
    pub fn new() -> Self {
        Self {
            // TODO: Hard coded for, as long as we don't support bundle properties
            property: 0,
            data: Vec::new(),
            data_file_name: None,
            unknown_1: 0,
        }
    }

    pub fn set_data(&mut self, data: Vec<u8>) {
        self.data = data;
    }

    pub fn size(&self) -> usize {
        self.data.len()
    }

    pub fn property(&self) -> u32 {
        self.property
    }

    pub fn data(&self) -> &[u8] {
        &self.data
    }

    pub fn data_file_name(&self) -> Option<&String> {
        self.data_file_name.as_ref()
    }

    #[tracing::instrument(skip_all)]
    fn read_header<R>(r: &mut R) -> Result<BundleFileHeader>
    where
        R: Read + Seek,
    {
        let variant = r.read_u32()?;
        let unknown_1 = r.read_u8()?;
        let size = r.read_u32()? as usize;
        r.skip_u8(1)?;
        let len_data_file_name = r.read_u32()? as usize;

        Ok(BundleFileHeader {
            size,
            unknown_1,
            variant,
            len_data_file_name,
        })
    }

    #[tracing::instrument(skip_all)]
    fn write_header<W>(&self, w: &mut W, props: Properties) -> Result<()>
    where
        W: Write + Seek,
    {
        w.write_u32(self.property)?;
        w.write_u8(self.unknown_1)?;

        let len_data_file_name = self.data_file_name.as_ref().map(|s| s.len()).unwrap_or(0);

        if props.contains(Properties::DATA) {
            w.write_u32(len_data_file_name as u32)?;
            w.write_u8(1)?;
            w.write_u32(0)?;
        } else {
            w.write_u32(self.data.len() as u32)?;
            w.write_u8(1)?;
            w.write_u32(len_data_file_name as u32)?;
        }

        Ok(())
    }
}

bitflags! {
    #[derive(Default, Clone, Copy, Debug)]
    pub struct Properties: u32 {
        const DATA = 0b100;
        // A custom flag used by DTMT to signify a file altered by mods.
        const MODDED = 1 << 31;
    }
}

#[derive(Clone, Debug)]
pub struct BundleFile {
    file_type: BundleFileType,
    name: IdString64,
    variants: Vec<BundleFileVariant>,
    props: Properties,
}

impl BundleFile {
    pub fn new(name: impl Into<IdString64>, file_type: BundleFileType) -> Self {
        Self {
            file_type,
            name: name.into(),
            variants: Vec::new(),
            props: Properties::empty(),
        }
    }

    pub fn add_variant(&mut self, variant: BundleFileVariant) {
        self.variants.push(variant)
    }

    pub fn set_variants(&mut self, variants: Vec<BundleFileVariant>) {
        self.variants = variants;
    }

    pub fn set_props(&mut self, props: Properties) {
        self.props = props;
    }

    pub fn set_modded(&mut self, is_modded: bool) {
        self.props.set(Properties::MODDED, is_modded);
    }

    #[tracing::instrument(name = "File::read", skip(ctx, r))]
    pub fn from_reader<R>(ctx: &crate::Context, r: &mut R, props: Properties) -> Result<Self>
    where
        R: Read + Seek,
    {
        let file_type = BundleFileType::from(r.read_u64()?);
        let hash = Murmur64::from(r.read_u64()?);
        let name = ctx.lookup_hash(hash, HashGroup::Filename);

        let header_count = r.read_u32()? as usize;
        tracing::trace!(header_count);
        let mut headers = Vec::with_capacity(header_count);
        r.skip_u32(0)?;

        for i in 0..header_count {
            let span = tracing::debug_span!("Read file header", i);
            let _enter = span.enter();

            let header = BundleFileVariant::read_header(r)
                .wrap_err_with(|| format!("Failed to read header {i}"))?;

            // TODO: Figure out how `header.unknown_1` correlates to `properties::DATA`
            // if props.contains(Properties::DATA) {
            //     tracing::debug!("props: {props:?} | unknown_1: {}", header.unknown_1)
            // }

            headers.push(header);
        }

        let mut variants = Vec::with_capacity(header_count);
        for (i, header) in headers.into_iter().enumerate() {
            let span = tracing::debug_span!(
                "Read file data {}",
                i,
                size = header.size,
                len_data_file_name = header.len_data_file_name
            );
            let _enter = span.enter();

            let (data, data_file_name) = if props.contains(Properties::DATA) {
                let data = vec![];
                let s = r
                    .read_string_len(header.size)
                    .wrap_err("Failed to read data file name")?;

                (data, Some(s))
            } else {
                let mut data = vec![0; header.size];
                r.read_exact(&mut data)
                    .wrap_err_with(|| format!("Failed to read file {i}"))?;

                let data_file_name = if header.len_data_file_name > 0 {
                    let s = r
                        .read_string_len(header.len_data_file_name)
                        .wrap_err("Failed to read data file name")?;
                    Some(s)
                } else {
                    None
                };

                (data, data_file_name)
            };

            let variant = BundleFileVariant {
                property: header.variant,
                data,
                data_file_name,
                unknown_1: header.unknown_1,
            };

            variants.push(variant);
        }

        Ok(Self {
            variants,
            file_type,
            name,
            props,
        })
    }

    #[tracing::instrument(name = "File::to_binary", skip_all)]
    pub fn to_binary(&self) -> Result<Vec<u8>> {
        let mut w = Cursor::new(Vec::new());

        w.write_u64(self.file_type.hash().into())?;
        w.write_u64(self.name.to_murmur64().into())?;
        w.write_u32(self.variants.len() as u32)?;

        // TODO: Figure out what this is
        w.write_u32(0x0)?;

        for variant in self.variants.iter() {
            w.write_u32(variant.property())?;
            w.write_u8(variant.unknown_1)?;

            let len_data_file_name = variant.data_file_name().map(|s| s.len()).unwrap_or(0);

            if self.props.contains(Properties::DATA) {
                w.write_u32(len_data_file_name as u32)?;
                w.write_u8(1)?;
                w.write_u32(0)?;
            } else {
                w.write_u32(variant.size() as u32)?;
                w.write_u8(1)?;
                w.write_u32(len_data_file_name as u32)?;
            }
        }

        for variant in self.variants.iter() {
            w.write_all(&variant.data)?;
            if let Some(s) = &variant.data_file_name {
                w.write_all(s.as_bytes())?;
            }
        }

        Ok(w.into_inner())
    }

    #[tracing::instrument("File::from_sjson", skip(sjson, name), fields(name = %name.display()))]
    pub async fn from_sjson(
        name: IdString64,
        file_type: BundleFileType,
        sjson: impl AsRef<str>,
        root: impl AsRef<Path> + std::fmt::Debug,
    ) -> Result<Self> {
        match file_type {
            BundleFileType::Lua => lua::compile(name, sjson).wrap_err("Failed to compile Lua file"),
            BundleFileType::Unknown(_) => {
                eyre::bail!("Unknown file type. Cannot compile from SJSON");
            }
            _ => {
                eyre::bail!(
                    "Compiling file type {} is not yet supported",
                    file_type.ext_name()
                )
            }
        }
    }

    pub fn props(&self) -> Properties {
        self.props
    }

    pub fn base_name(&self) -> &IdString64 {
        &self.name
    }

    pub fn name(&self, decompiled: bool, variant: Option<u32>) -> String {
        let mut s = self.name.display().to_string();
        s.push('.');

        if let Some(variant) = variant {
            s.push_str(&variant.to_string());
            s.push('.');
        }

        if decompiled {
            s.push_str(&self.file_type.decompiled_ext_name());
        } else {
            s.push_str(&self.file_type.ext_name());
        }

        s
    }

    pub fn matches_name(&self, name: &IdString64) -> bool {
        if self.name == *name {
            return true;
        }

        if let IdString64::String(name) = name {
            self.name(false, None) == *name || self.name(true, None) == *name
        } else {
            false
        }
    }

    pub fn file_type(&self) -> BundleFileType {
        self.file_type
    }

    pub fn variants(&self) -> &Vec<BundleFileVariant> {
        &self.variants
    }

    pub fn variants_mut(&mut self) -> impl Iterator<Item = &mut BundleFileVariant> {
        self.variants.iter_mut()
    }

    pub fn raw(&self) -> Result<Vec<UserFile>> {
        let files = self
            .variants
            .iter()
            .map(|variant| {
                let name = if self.variants.len() > 1 {
                    self.name(false, Some(variant.property()))
                } else {
                    self.name(false, None)
                };
                UserFile {
                    data: variant.data().to_vec(),
                    name: Some(name),
                }
            })
            .collect();

        Ok(files)
    }

    #[tracing::instrument(name = "File::decompiled", skip_all)]
    pub async fn decompiled(&self, ctx: &crate::Context) -> Result<Vec<UserFile>> {
        let file_type = self.file_type();

        if tracing::enabled!(tracing::Level::DEBUG) {
            tracing::debug!(
                name = self.name(true, None),
                variants = self.variants.len(),
                "Attempting to decompile"
            );
        }

        if file_type == BundleFileType::Strings {
            return strings::decompile(ctx, &self.variants);
        }

        let tasks = self.variants.iter().map(|variant| async move {
            let data = variant.data();
            let name = if self.variants.len() > 1 {
                self.name(true, Some(variant.property()))
            } else {
                self.name(true, None)
            };

            let res = match file_type {
                BundleFileType::Lua => lua::decompile(ctx, data).await,
                BundleFileType::Package => package::decompile(ctx, name.clone(), data),
                _ => {
                    tracing::debug!("Can't decompile, unknown file type");
                    Ok(vec![UserFile::with_name(data.to_vec(), name.clone())])
                }
            };

            let res = res.wrap_err_with(|| format!("Failed to decompile file {name}"));
            match res {
                Ok(files) => files,
                Err(err) => {
                    tracing::error!("{:?}", err);
                    vec![]
                }
            }
        });

        let results = join_all(tasks).await;

        Ok(results.into_iter().flatten().collect())
    }
}

impl PartialEq for BundleFile {
    fn eq(&self, other: &Self) -> bool {
        self.name == other.name && self.file_type == other.file_type
    }
}

pub struct UserFile {
    // TODO: Might be able to avoid some allocations with a Cow here
    data: Vec<u8>,
    name: Option<String>,
}

impl UserFile {
    pub fn new(data: Vec<u8>) -> Self {
        Self { data, name: None }
    }

    pub fn with_name(data: Vec<u8>, name: String) -> Self {
        Self {
            data,
            name: Some(name),
        }
    }

    pub fn data(&self) -> &[u8] {
        &self.data
    }

    pub fn name(&self) -> Option<&String> {
        self.name.as_ref()
    }
}