BlueArchiveToolkit/infrastructure/src/import.rs

//! 资源导入服务。

use bat_adapters::manifest::GenericManifest;
use bat_adapters::unity::{RawAssetBundle, UnityAdapterRegistry};
use bat_core::domain::{Resource, ResourceType};
use bat_core::repositories::{CasRepository, ResourceRepository};
use std::collections::BTreeMap;
use std::path::{Path, PathBuf};

/// 待导入的 AssetBundle 数据。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct BundleSource {
    /// Bundle 路径或文件名。
    pub path: String,
    /// Bundle 原始字节。
    pub data: Vec<u8>,
}

impl BundleSource {
    /// 创建 Bundle 输入。
    pub fn new(path: impl Into<String>, data: impl Into<Vec<u8>>) -> Self {
        Self {
            path: path.into(),
            data: data.into(),
        }
    }
}

/// 单个导入结果。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ImportedResource {
    /// 写入 `ResourceRepository` 的资源 ID。
    pub id: String,
    /// Manifest 中的资源路径。
    pub source_path: String,
    /// Manifest 或路径推断出的资源类型。
    pub resource_type: ResourceType,
    /// 导入链路使用的稳定分类标签。
    pub category: ResourceImportCategory,
    /// CAS 返回的对象 ID。
    pub object_id: String,
    /// 写入 CAS 的字节数。
    pub bytes: u64,
    /// UnityFS 解析摘要，仅 AssetBundle 资源会填充。
    pub unityfs: Option<UnityFsImportSummary>,
}

/// Stable resource category produced by the import pipeline.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum ResourceImportCategory {
    /// UnityFS AssetBundle resource.
    AssetBundle,
    /// Unity TextAsset or text-like payload.
    TextAsset,
    /// Table bundle/database resource.
    Table,
    /// Media payload such as audio, video, texture, or media archive.
    Media,
    /// Manifest/catalog sidecar.
    Manifest,
    /// Resource that does not match a known category yet.
    Other,
}

impl ResourceImportCategory {
    /// Returns a stable lowercase label for reports and golden tests.
    pub fn as_str(self) -> &'static str {
        match self {
            Self::AssetBundle => "asset_bundle",
            Self::TextAsset => "text_asset",
            Self::Table => "table",
            Self::Media => "media",
            Self::Manifest => "manifest",
            Self::Other => "other",
        }
    }
}

/// 导入时解析到的 UnityFS 基础结构摘要。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UnityFsImportSummary {
    /// Bundle 声明的 Unity 版本。
    pub unity_version: String,
    /// UnityFS block 数量。
    pub block_count: usize,
    /// UnityFS directory 数量。
    pub directory_count: usize,
    /// UnityFS directory 路径。
    pub directories: Vec<String>,
}

/// Manifest 导入报告。
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct ResourceImportReport {
    /// 成功导入的资源。
    pub imported: Vec<ImportedResource>,
    /// 被跳过的资源路径。
    pub skipped: Vec<String>,
    /// 按稳定分类标签统计的成功导入数量。
    pub category_counts: BTreeMap<String, usize>,
}

/// 将解析后的资源清单写入 CAS 与资源仓储的应用服务。
pub struct ResourceImportService<'a> {
    cas: &'a dyn CasRepository,
    resources: &'a dyn ResourceRepository,
    unity_adapters: UnityAdapterRegistry,
}

impl<'a> ResourceImportService<'a> {
    /// 创建导入服务。
    pub fn new(cas: &'a dyn CasRepository, resources: &'a dyn ResourceRepository) -> Self {
        Self::with_unity_adapters(cas, resources, UnityAdapterRegistry::with_defaults())
    }

    /// 使用指定 Unity 适配器注册表创建导入服务。
    pub fn with_unity_adapters(
        cas: &'a dyn CasRepository,
        resources: &'a dyn ResourceRepository,
        unity_adapters: UnityAdapterRegistry,
    ) -> Self {
        Self {
            cas,
            resources,
            unity_adapters,
        }
    }

    /// 导入 manifest 中可用的资源数据。
    ///
    /// AssetBundle 条目必须能在 `bundles` 中按完整路径或文件名找到对应
    /// 数据并成功解析 UnityFS，否则返回错误。TextAsset/Table/Media 等
    /// 非 AssetBundle 条目在有数据时写入 CAS 和 `ResourceRepository`；缺少
    /// 数据时记录到 `skipped`，用于渐进式导入官方下载结果。
    pub async fn import_manifest_bundles(
        &self,
        manifest: &GenericManifest,
        bundles: &[BundleSource],
    ) -> bat_core::Result<ResourceImportReport> {
        let mut report = ResourceImportReport::default();

        for entry in &manifest.resources {
            let category = classify_import_category(entry.resource_type, &entry.path);
            let Some(bundle) = find_bundle(bundles, &entry.path) else {
                if entry.resource_type == ResourceType::AssetBundle {
                    return Err(bat_core::Error::InvalidArgument(format!(
                        "Missing bundle data for manifest resource: {}",
                        entry.path
                    )));
                }
                report.skipped.push(entry.path.clone());
                continue;
            };

            let unityfs = if entry.resource_type == ResourceType::AssetBundle {
                Some(self.parse_unityfs_summary(&entry.path, bundle).await?)
            } else {
                None
            };

            let object_id = self.cas.store(&bundle.data).await?;
            let mut stored_entry = entry.clone();
            stored_entry.hash = object_id.clone();
            stored_entry.size = bundle.data.len() as u64;

            let resource = Resource {
                id: resource_id_for_path(&entry.path),
                local_path: PathBuf::from(&entry.path),
                entry: stored_entry,
            };
            let id = self.resources.add(resource).await?;

            report.imported.push(ImportedResource {
                id,
                source_path: entry.path.clone(),
                resource_type: entry.resource_type,
                category,
                object_id,
                bytes: bundle.data.len() as u64,
                unityfs,
            });
            *report
                .category_counts
                .entry(category.as_str().to_string())
                .or_insert(0) += 1;
        }

        Ok(report)
    }

    async fn parse_unityfs_summary(
        &self,
        manifest_path: &str,
        bundle: &BundleSource,
    ) -> bat_core::Result<UnityFsImportSummary> {
        let raw_bundle = RawAssetBundle {
            data: bundle.data.clone(),
            path: Some(manifest_path.to_string()),
        };
        let adapter = self
            .unity_adapters
            .select_adapter(&raw_bundle)
            .map_err(|error| {
                bat_core::Error::InvalidArgument(format!(
                    "No Unity adapter for manifest resource {}: {}",
                    manifest_path, error
                ))
            })?;
        let parsed = adapter.parse(&raw_bundle).await.map_err(|error| {
            bat_core::Error::InvalidArgument(format!(
                "Invalid UnityFS bundle for manifest resource {}: {}",
                manifest_path, error
            ))
        })?;

        Ok(UnityFsImportSummary {
            unity_version: parsed.unity_version,
            block_count: parsed.blocks.len(),
            directory_count: parsed.directories.len(),
            directories: parsed
                .directories
                .into_iter()
                .map(|directory| directory.path)
                .collect(),
        })
    }
}

fn resource_id_for_path(path: &str) -> String {
    format!("resource/{}", path)
}

fn find_bundle<'a>(bundles: &'a [BundleSource], manifest_path: &str) -> Option<&'a BundleSource> {
    bundles
        .iter()
        .find(|bundle| bundle.path == manifest_path)
        .or_else(|| {
            file_name(manifest_path).and_then(|manifest_name| {
                bundles
                    .iter()
                    .find(|bundle| file_name(&bundle.path) == Some(manifest_name))
            })
        })
}

fn file_name(path: &str) -> Option<&str> {
    Path::new(path).file_name()?.to_str()
}

fn classify_import_category(resource_type: ResourceType, path: &str) -> ResourceImportCategory {
    match resource_type {
        ResourceType::AssetBundle => ResourceImportCategory::AssetBundle,
        ResourceType::TextAsset => ResourceImportCategory::TextAsset,
        ResourceType::TableBundle => ResourceImportCategory::Table,
        ResourceType::Media => ResourceImportCategory::Media,
        ResourceType::Manifest => ResourceImportCategory::Manifest,
        ResourceType::Other => classify_import_category_from_path(path),
    }
}

fn classify_import_category_from_path(path: &str) -> ResourceImportCategory {
    let normalized = path.replace('\\', "/").to_ascii_lowercase();
    if normalized.contains("tablebundles/") {
        ResourceImportCategory::Table
    } else if normalized.contains("mediaresources/")
        || normalized.contains("mediaresources-")
        || matches!(
            normalized.rsplit('.').next(),
            Some("mp3" | "mp4" | "ogg" | "wav" | "png" | "jpg" | "jpeg" | "webp" | "acb" | "awb")
        )
    {
        ResourceImportCategory::Media
    } else if normalized.contains("textassets/")
        || matches!(
            normalized.rsplit('.').next(),
            Some("txt" | "csv" | "xml" | "yaml" | "yml")
        )
    {
        ResourceImportCategory::TextAsset
    } else if normalized.contains("catalog") || normalized.ends_with(".hash") {
        ResourceImportCategory::Manifest
    } else {
        ResourceImportCategory::Other
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{FileSystemCasRepository, InMemoryResourceRepository};
    use bat_adapters::manifest::{GenericManifest, ManifestFormat, ManifestMetadata};
    use bat_core::domain::ResourceEntry;
    use bat_core::repositories::cas_repository::CasRepository;
    use bat_core::repositories::resource_repository::{ResourceQuery, ResourceRepository};
    use std::collections::HashMap;
    use tempfile::TempDir;

    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, "SYNTHETIC-CAB");

        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, 16);
        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
    }

    fn synthetic_manifest() -> GenericManifest {
        GenericManifest {
            format: ManifestFormat::AddressablesCatalog,
            resources: vec![
                ResourceEntry {
                    path: "synthetic/minimal.bundle".to_string(),
                    hash: "synthetic-manifest-hash".to_string(),
                    size: 99,
                    resource_type: ResourceType::AssetBundle,
                    address: None,
                    dependencies: Vec::new(),
                },
                ResourceEntry {
                    path: "synthetic/catalog.json".to_string(),
                    hash: "synthetic-catalog-hash".to_string(),
                    size: 1,
                    resource_type: ResourceType::Manifest,
                    address: None,
                    dependencies: Vec::new(),
                },
                ResourceEntry {
                    path: "TextAssets/dialogue.csv".to_string(),
                    hash: "synthetic-text-hash".to_string(),
                    size: 1,
                    resource_type: ResourceType::TextAsset,
                    address: Some("dialogue".to_string()),
                    dependencies: Vec::new(),
                },
                ResourceEntry {
                    path: "TableBundles/ExcelDB.db".to_string(),
                    hash: "synthetic-table-hash".to_string(),
                    size: 1,
                    resource_type: ResourceType::TableBundle,
                    address: None,
                    dependencies: Vec::new(),
                },
                ResourceEntry {
                    path: "MediaResources-Windows/voice/title.acb".to_string(),
                    hash: "synthetic-media-hash".to_string(),
                    size: 1,
                    resource_type: ResourceType::Media,
                    address: None,
                    dependencies: Vec::new(),
                },
            ],
            metadata: ManifestMetadata {
                locator_id: None,
                cdn_prefixes: Vec::new(),
                extra: HashMap::new(),
            },
        }
    }

    #[tokio::test]
    async fn imports_synthetic_manifest_bundles_into_cas_and_resource_repository() {
        let temp_dir = TempDir::new().unwrap();
        let cas = FileSystemCasRepository::new(temp_dir.path().join("cas"));
        let resources = InMemoryResourceRepository::new();
        let service = ResourceImportService::new(&cas, &resources);

        let report = service
            .import_manifest_bundles(
                &synthetic_manifest(),
                &[
                    BundleSource::new("minimal.bundle", synthetic_minimal_unityfs_bundle()),
                    BundleSource::new("dialogue.csv", b"id,text\n1,hello".to_vec()),
                    BundleSource::new("ExcelDB.db", b"table-fixture".to_vec()),
                    BundleSource::new("title.acb", b"media-fixture".to_vec()),
                ],
            )
            .await
            .unwrap();

        assert_eq!(report.imported.len(), 4);
        assert_eq!(report.skipped, vec!["synthetic/catalog.json".to_string()]);
        assert!(cas.exists(&report.imported[0].object_id).await);
        assert_eq!(report.imported[0].resource_type, ResourceType::AssetBundle);
        assert_eq!(
            report.imported[0].category,
            ResourceImportCategory::AssetBundle
        );
        let unityfs = report.imported[0].unityfs.as_ref().unwrap();
        assert_eq!(unityfs.unity_version, "2021.3.56f2");
        assert_eq!(unityfs.block_count, 1);
        assert_eq!(unityfs.directory_count, 1);
        assert_eq!(unityfs.directories, vec!["SYNTHETIC-CAB".to_string()]);
        assert_eq!(
            report.imported[1].category,
            ResourceImportCategory::TextAsset
        );
        assert_eq!(report.imported[2].category, ResourceImportCategory::Table);
        assert_eq!(report.imported[3].category, ResourceImportCategory::Media);
        assert_eq!(report.category_counts.get("asset_bundle"), Some(&1));
        assert_eq!(report.category_counts.get("text_asset"), Some(&1));
        assert_eq!(report.category_counts.get("table"), Some(&1));
        assert_eq!(report.category_counts.get("media"), Some(&1));

        let indexed = resources.list(ResourceQuery::all()).await.unwrap();
        assert_eq!(indexed.len(), 4);
        let indexed_bundle = indexed
            .iter()
            .find(|resource| resource.entry.path == "synthetic/minimal.bundle")
            .unwrap();
        assert_eq!(indexed_bundle.entry.hash, report.imported[0].object_id);
        assert_eq!(
            indexed_bundle.entry.size,
            synthetic_minimal_unityfs_bundle().len() as u64
        );
        assert_eq!(
            indexed
                .iter()
                .find(|resource| resource.entry.path.ends_with("title.acb"))
                .unwrap()
                .entry
                .resource_type,
            ResourceType::Media
        );
        assert_eq!(
            indexed
                .iter()
                .find(|resource| resource.entry.path.ends_with("ExcelDB.db"))
                .unwrap()
                .entry
                .resource_type,
            ResourceType::TableBundle
        );
        assert_eq!(
            indexed
                .iter()
                .find(|resource| resource.entry.path.ends_with("dialogue.csv"))
                .unwrap()
                .entry
                .resource_type,
            ResourceType::TextAsset
        );
    }

    #[tokio::test]
    async fn returns_error_when_bundle_data_is_missing() {
        let temp_dir = TempDir::new().unwrap();
        let cas = FileSystemCasRepository::new(temp_dir.path().join("cas"));
        let resources = InMemoryResourceRepository::new();
        let service = ResourceImportService::new(&cas, &resources);

        let error = service
            .import_manifest_bundles(&synthetic_manifest(), &[])
            .await
            .unwrap_err();

        assert!(matches!(error, bat_core::Error::InvalidArgument(_)));
    }

    #[tokio::test]
    async fn returns_error_when_bundle_is_not_unityfs() {
        let temp_dir = TempDir::new().unwrap();
        let cas = FileSystemCasRepository::new(temp_dir.path().join("cas"));
        let resources = InMemoryResourceRepository::new();
        let service = ResourceImportService::new(&cas, &resources);

        let error = service
            .import_manifest_bundles(
                &synthetic_manifest(),
                &[BundleSource::new(
                    "minimal.bundle",
                    b"not-a-unityfs-bundle".to_vec(),
                )],
            )
            .await
            .unwrap_err();

        assert!(matches!(error, bat_core::Error::InvalidArgument(_)));
        assert_eq!(resources.count(ResourceQuery::all()).await.unwrap(), 0);
    }

    #[tokio::test]
    async fn indexes_non_bundle_resources_without_unityfs_summary() {
        let temp_dir = TempDir::new().unwrap();
        let cas = FileSystemCasRepository::new(temp_dir.path().join("cas"));
        let resources = InMemoryResourceRepository::new();
        let service = ResourceImportService::new(&cas, &resources);
        let mut manifest = synthetic_manifest();
        manifest
            .resources
            .retain(|entry| entry.resource_type != ResourceType::AssetBundle);

        let report = service
            .import_manifest_bundles(
                &manifest,
                &[
                    BundleSource::new("dialogue.csv", b"id,text\n1,hello".to_vec()),
                    BundleSource::new("ExcelDB.db", b"table-fixture".to_vec()),
                    BundleSource::new("title.acb", b"media-fixture".to_vec()),
                ],
            )
            .await
            .unwrap();

        assert_eq!(report.imported.len(), 3);
        assert!(report.imported.iter().all(|item| item.unityfs.is_none()));
        assert_eq!(report.skipped, vec!["synthetic/catalog.json".to_string()]);
        assert_eq!(
            resources
                .count(ResourceQuery::by_type(ResourceType::TextAsset))
                .await
                .unwrap(),
            1
        );
        assert_eq!(
            resources
                .count(ResourceQuery::by_type(ResourceType::TableBundle))
                .await
                .unwrap(),
            1
        );
        assert_eq!(
            resources
                .count(ResourceQuery::by_type(ResourceType::Media))
                .await
                .unwrap(),
            1
        );
    }
}