BlueArchiveToolkit/adapters/src/unity/unity_2021_3.rs

//! Unity 2021.3 Adapter
//!
//! 支持 Unity 2021.3.x 版本的 AssetBundle

use super::adapter::{
    ParsedAssetBundle, RawAssetBundle, UnityAdapter, UnityFsBlockInfo, UnityFsCompression,
    UnityFsDirectoryInfo, UnityFsHeader, VersionRange,
};
use async_trait::async_trait;
use std::io::Cursor;

const SERIALIZE_NOT_IMPLEMENTED: &str = "serialize() 将在 Phase 2 实现";
const UNITYFS_COMPRESSION_MASK: u32 = 0x3f;
const UNITYFS_BLOCK_INFO_AT_END_FLAG: u32 = 0x80;
const UNITYFS_ALIGNMENT: usize = 16;

/// Unity 2021.3 Adapter
pub struct Unity2021_3Adapter;

impl Unity2021_3Adapter {
    /// 创建新的适配器实例
    pub fn new() -> Self {
        Self
    }

    fn unity_version_bytes(data: &[u8]) -> Option<&[u8]> {
        if data.len() < 20 {
            return None;
        }

        if &data[0..7] != b"UnityFS" {
            return None;
        }

        let version_start = data.windows(7).position(|window| window == b"2021.3.")?;
        let version_bytes = &data[version_start..];
        let version_end = version_bytes
            .iter()
            .position(|&byte| byte == 0 || !byte.is_ascii())?;

        Some(&version_bytes[..version_end])
    }

    /// 检测 Unity 版本（从文件头）
    fn detect_unity_version(data: &[u8]) -> Option<String> {
        let version_bytes = Self::unity_version_bytes(data)?;
        std::str::from_utf8(version_bytes)
            .map(ToOwned::to_owned)
            .ok()
    }

    fn parse_unityfs(data: &[u8]) -> Result<ParsedAssetBundle, String> {
        let mut reader = UnityFsReader::new(data);
        let signature = reader.read_c_string("signature")?;
        if signature != "UnityFS" {
            return Err(format!("Unsupported AssetBundle signature: {}", signature));
        }

        let format_version = reader.read_u32("format_version")?;
        let target_version = reader.read_c_string("target_version")?;
        let unity_version = reader.read_c_string("unity_version")?;
        let total_size = reader.read_u64("total_size")?;
        let compressed_blocks_info_size = reader.read_u32("compressed_blocks_info_size")?;
        let uncompressed_blocks_info_size = reader.read_u32("uncompressed_blocks_info_size")?;
        let flags = reader.read_u32("flags")?;

        let header = UnityFsHeader {
            format_version,
            target_version,
            unity_version,
            total_size,
            compressed_blocks_info_size,
            uncompressed_blocks_info_size,
            flags,
        };

        if format_version >= 7 {
            reader.align(UNITYFS_ALIGNMENT)?;
        }

        let blocks_info_bytes = read_blocks_info_bytes(data, &mut reader, &header)?;
        let block_info = decompress_blocks_info(
            blocks_info_bytes,
            compressed_blocks_info_size,
            uncompressed_blocks_info_size,
            flags,
        )?;
        let (blocks, directories) = Self::parse_blocks_info(&block_info)?;

        Ok(ParsedAssetBundle {
            unity_version: header.unity_version.clone(),
            assets: directories
                .iter()
                .map(|directory| directory.path.clone())
                .collect(),
            raw_data: data.to_vec(),
            unityfs_header: Some(header),
            blocks,
            directories,
        })
    }

    fn parse_blocks_info(
        data: &[u8],
    ) -> Result<(Vec<UnityFsBlockInfo>, Vec<UnityFsDirectoryInfo>), String> {
        let mut reader = UnityFsReader::new(data);
        let _hash = reader.read_bytes(16, "blocks_info_hash")?;
        let block_count = reader.read_i32("block_count")?;
        if block_count < 0 {
            return Err(format!("Invalid UnityFS block count: {}", block_count));
        }

        let mut blocks = Vec::with_capacity(block_count as usize);
        for _ in 0..block_count {
            let uncompressed_size = reader.read_u32("block_uncompressed_size")?;
            let compressed_size = reader.read_u32("block_compressed_size")?;
            let flags = reader.read_u16("block_flags")?;
            blocks.push(UnityFsBlockInfo {
                uncompressed_size,
                compressed_size,
                flags,
                compression: compression_from_flags(flags),
            });
        }

        let directory_count = reader.read_i32("directory_count")?;
        if directory_count < 0 {
            return Err(format!(
                "Invalid UnityFS directory count: {}",
                directory_count
            ));
        }

        let mut directories = Vec::with_capacity(directory_count as usize);
        for _ in 0..directory_count {
            directories.push(UnityFsDirectoryInfo {
                offset: reader.read_u64("directory_offset")?,
                size: reader.read_u64("directory_size")?,
                flags: reader.read_u32("directory_flags")?,
                path: reader.read_c_string("directory_path")?,
            });
        }

        Ok((blocks, directories))
    }
}

fn read_blocks_info_bytes<'a>(
    data: &'a [u8],
    reader: &mut UnityFsReader<'a>,
    header: &UnityFsHeader,
) -> Result<&'a [u8], String> {
    let len = header.compressed_blocks_info_size as usize;
    if blocks_info_at_end(header.flags) {
        let start = data.len().checked_sub(len).ok_or_else(|| {
            format!(
                "UnityFS block info at end underflow: compressed size {}, file size {}",
                len,
                data.len()
            )
        })?;
        return Ok(&data[start..]);
    }

    reader.read_bytes(len, "blocks_info")
}

fn blocks_info_at_end(flags: u32) -> bool {
    flags & UNITYFS_BLOCK_INFO_AT_END_FLAG != 0
}

fn decompress_blocks_info(
    data: &[u8],
    compressed_size: u32,
    uncompressed_size: u32,
    flags: u32,
) -> Result<Vec<u8>, String> {
    if data.len() != compressed_size as usize {
        return Err(format!(
            "UnityFS block info size mismatch: header says {}, read {}",
            compressed_size,
            data.len()
        ));
    }

    let compression = compression_from_flags((flags & UNITYFS_COMPRESSION_MASK) as u16);
    match compression {
        UnityFsCompression::None => {
            if compressed_size != uncompressed_size {
                return Err(format!(
                    "Uncompressed UnityFS block info size mismatch: compressed {} != uncompressed {}",
                    compressed_size, uncompressed_size
                ));
            }
            Ok(data.to_vec())
        }
        UnityFsCompression::Lz4 | UnityFsCompression::Lz4Hc => {
            lz4::block::decompress(data, Some(uncompressed_size as i32))
                .map_err(|error| format!("Failed to decompress UnityFS LZ4 block info: {}", error))
        }
        UnityFsCompression::Lzma => {
            let mut output = Vec::with_capacity(uncompressed_size as usize);
            lzma_rs::lzma_decompress(&mut Cursor::new(data), &mut output).map_err(|error| {
                format!("Failed to decompress UnityFS LZMA block info: {}", error)
            })?;
            if output.len() != uncompressed_size as usize {
                return Err(format!(
                    "UnityFS LZMA block info size mismatch: expected {}, got {}",
                    uncompressed_size,
                    output.len()
                ));
            }
            Ok(output)
        }
        UnityFsCompression::Unknown(value) => Err(format!(
            "Unsupported UnityFS block info compression flag: {}",
            value
        )),
    }
}

fn compression_from_flags(flags: u16) -> UnityFsCompression {
    match flags & UNITYFS_COMPRESSION_MASK as u16 {
        0 => UnityFsCompression::None,
        1 => UnityFsCompression::Lzma,
        2 => UnityFsCompression::Lz4,
        3 | 4 => UnityFsCompression::Lz4Hc,
        value => UnityFsCompression::Unknown(value),
    }
}

struct UnityFsReader<'a> {
    data: &'a [u8],
    offset: usize,
}

impl<'a> UnityFsReader<'a> {
    fn new(data: &'a [u8]) -> Self {
        Self { data, offset: 0 }
    }

    fn read_bytes(&mut self, len: usize, field: &str) -> Result<&'a [u8], String> {
        let end = self
            .offset
            .checked_add(len)
            .ok_or_else(|| format!("{} length overflow at offset {}", field, self.offset))?;
        if end > self.data.len() {
            return Err(format!(
                "Unexpected end while reading {} at offset {}: need {}, have {}",
                field,
                self.offset,
                len,
                self.data.len().saturating_sub(self.offset)
            ));
        }

        let bytes = &self.data[self.offset..end];
        self.offset = end;
        Ok(bytes)
    }

    fn align(&mut self, alignment: usize) -> Result<(), String> {
        if alignment == 0 {
            return Ok(());
        }

        let remainder = self.offset % alignment;
        if remainder == 0 {
            return Ok(());
        }

        let padding = alignment - remainder;
        self.read_bytes(padding, "alignment padding").map(|_| ())
    }

    fn read_u16(&mut self, field: &str) -> Result<u16, String> {
        let bytes = self.read_bytes(2, field)?;
        Ok(u16::from_be_bytes([bytes[0], bytes[1]]))
    }

    fn read_u32(&mut self, field: &str) -> Result<u32, String> {
        let bytes = self.read_bytes(4, field)?;
        Ok(u32::from_be_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]))
    }

    fn read_i32(&mut self, field: &str) -> Result<i32, String> {
        let bytes = self.read_bytes(4, field)?;
        Ok(i32::from_be_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]))
    }

    fn read_u64(&mut self, field: &str) -> Result<u64, String> {
        let bytes = self.read_bytes(8, field)?;
        Ok(u64::from_be_bytes([
            bytes[0], bytes[1], bytes[2], bytes[3], bytes[4], bytes[5], bytes[6], bytes[7],
        ]))
    }

    fn read_c_string(&mut self, field: &str) -> Result<String, String> {
        let remaining = &self.data[self.offset..];
        let Some(length) = remaining.iter().position(|&byte| byte == 0) else {
            return Err(format!(
                "Missing null terminator while reading {} at offset {}",
                field, self.offset
            ));
        };
        let bytes = self.read_bytes(length, field)?;
        self.offset += 1;
        std::str::from_utf8(bytes)
            .map(ToOwned::to_owned)
            .map_err(|error| format!("Invalid UTF-8 in {}: {}", field, error))
    }
}

impl Default for Unity2021_3Adapter {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl UnityAdapter for Unity2021_3Adapter {
    fn name(&self) -> &str {
        "Unity-2021.3"
    }

    fn supported_versions(&self) -> VersionRange {
        VersionRange::new("2021.3.0", "2021.3.99")
    }

    fn can_handle(&self, bundle: &RawAssetBundle) -> bool {
        // 检测 Unity 版本
        if let Some(version) = Self::detect_unity_version(&bundle.data) {
            self.supported_versions().contains(&version)
        } else {
            false
        }
    }

    async fn parse(&self, bundle: &RawAssetBundle) -> Result<ParsedAssetBundle, String> {
        let parsed = Self::parse_unityfs(&bundle.data)?;
        if !self.supported_versions().contains(&parsed.unity_version) {
            return Err(format!(
                "Unsupported Unity version for {}: {}",
                self.name(),
                parsed.unity_version
            ));
        }
        Ok(parsed)
    }

    async fn serialize(&self, _parsed: &ParsedAssetBundle) -> Result<Vec<u8>, String> {
        // TODO: Phase 2 实现
        // 需要：
        // 1. 序列化 Asset 对象
        // 2. 重新构建 TypeTree
        // 3. 压缩数据块
        // 4. 写入 UnityFS 文件头
        Err(SERIALIZE_NOT_IMPLEMENTED.to_string())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn push_c_string(data: &mut Vec<u8>, value: &str) {
        data.extend_from_slice(value.as_bytes());
        data.push(0);
    }

    fn push_u16(data: &mut Vec<u8>, value: u16) {
        data.extend_from_slice(&value.to_be_bytes());
    }

    fn push_u32(data: &mut Vec<u8>, value: u32) {
        data.extend_from_slice(&value.to_be_bytes());
    }

    fn push_i32(data: &mut Vec<u8>, value: i32) {
        data.extend_from_slice(&value.to_be_bytes());
    }

    fn push_u64(data: &mut Vec<u8>, value: u64) {
        data.extend_from_slice(&value.to_be_bytes());
    }

    fn align(data: &mut Vec<u8>, alignment: usize) {
        let remainder = data.len() % alignment;
        if remainder != 0 {
            data.resize(data.len() + alignment - remainder, 0);
        }
    }

    fn synthetic_minimal_unityfs_bundle() -> Vec<u8> {
        let mut blocks_info = Vec::new();
        blocks_info.extend_from_slice(&[0; 16]);
        push_i32(&mut blocks_info, 1);
        push_u32(&mut blocks_info, 4);
        push_u32(&mut blocks_info, 4);
        push_u16(&mut blocks_info, 0);
        push_i32(&mut blocks_info, 1);
        push_u64(&mut blocks_info, 0);
        push_u64(&mut blocks_info, 4);
        push_u32(&mut blocks_info, 0);
        push_c_string(&mut blocks_info, "CAB-test");

        let mut data = Vec::new();
        push_c_string(&mut data, "UnityFS");
        push_u32(&mut data, 8);
        push_c_string(&mut data, "5.x.x");
        push_c_string(&mut data, "2021.3.56f2");
        push_u64(&mut data, 0);
        push_u32(&mut data, blocks_info.len() as u32);
        push_u32(&mut data, blocks_info.len() as u32);
        push_u32(&mut data, 0);
        align(&mut data, UNITYFS_ALIGNMENT);
        data.extend_from_slice(&blocks_info);
        data.extend_from_slice(b"data");

        let total_size = data.len() as u64;
        let total_size_offset = b"UnityFS\0".len() + 4 + b"5.x.x\0".len() + b"2021.3.56f2\0".len();
        data[total_size_offset..total_size_offset + 8].copy_from_slice(&total_size.to_be_bytes());
        data
    }

    #[test]
    fn test_adapter_name() {
        let adapter = Unity2021_3Adapter::new();
        assert_eq!(adapter.name(), "Unity-2021.3");
    }

    #[test]
    fn test_supported_versions() {
        let adapter = Unity2021_3Adapter::new();
        let range = adapter.supported_versions();

        assert!(range.contains("2021.3.0"));
        assert!(range.contains("2021.3.56"));
        assert!(range.contains("2021.3.9f1"));
        assert!(!range.contains("2021.2.0"));
    }

    #[test]
    fn test_can_handle_valid_bundle() {
        let adapter = Unity2021_3Adapter::new();

        let bundle = RawAssetBundle {
            data: synthetic_minimal_unityfs_bundle(),
            path: Some("synthetic-minimal.bundle".to_string()),
        };

        assert!(adapter.can_handle(&bundle));
    }

    #[test]
    fn test_can_handle_invalid_bundle() {
        let adapter = Unity2021_3Adapter::new();

        let bundle = RawAssetBundle {
            data: vec![0, 1, 2, 3],
            path: None,
        };

        assert!(!adapter.can_handle(&bundle));
    }

    #[tokio::test]
    async fn test_parse_unityfs_metadata() {
        let adapter = Unity2021_3Adapter::new();

        let bundle = RawAssetBundle {
            data: synthetic_minimal_unityfs_bundle(),
            path: Some("synthetic-minimal.bundle".to_string()),
        };

        let parsed = adapter.parse(&bundle).await.unwrap();
        let header = parsed.unityfs_header.unwrap();
        assert_eq!(parsed.unity_version, "2021.3.56f2");
        assert_eq!(header.format_version, 8);
        assert_eq!(parsed.blocks.len(), 1);
        assert_eq!(parsed.blocks[0].compression, UnityFsCompression::None);
        assert_eq!(parsed.directories.len(), 1);
        assert_eq!(parsed.directories[0].path, "CAB-test");
        assert_eq!(parsed.assets, vec!["CAB-test".to_string()]);
    }

    #[tokio::test]
    async fn test_parse_rejects_invalid_bundle() {
        let adapter = Unity2021_3Adapter::new();

        let bundle = RawAssetBundle {
            data: vec![],
            path: None,
        };

        let result = adapter.parse(&bundle).await;
        assert!(result.is_err());
        assert!(result.unwrap_err().contains("signature"));
    }
}