BlueArchiveToolkit/infrastructure/src/resources.rs

//! 内存资源仓储实现。

use async_trait::async_trait;
use bat_core::domain::{Resource, ResourceEntry, ResourceType};
use bat_core::repositories::resource_repository::{ResourceQuery, ResourceRepository};
use sqlx::sqlite::{SqliteConnectOptions, SqlitePoolOptions, SqliteQueryResult};
use sqlx::{QueryBuilder, Sqlite, SqlitePool};
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::str::FromStr;
use tokio::sync::RwLock;

/// 以内存 `HashMap` 保存资源索引的仓储实现。
#[derive(Debug, Default)]
pub struct InMemoryResourceRepository {
    resources: RwLock<HashMap<String, Resource>>,
}

impl InMemoryResourceRepository {
    /// 创建空资源仓储。
    pub fn new() -> Self {
        Self::default()
    }

    fn sorted_resources(resources: &HashMap<String, Resource>) -> Vec<Resource> {
        let mut resources = resources.values().cloned().collect::<Vec<_>>();
        resources.sort_by(|left, right| left.id.cmp(&right.id));
        resources
    }
}

#[async_trait]
impl ResourceRepository for InMemoryResourceRepository {
    async fn add(&self, resource: Resource) -> bat_core::Result<String> {
        let id = resource.id.clone();
        self.resources.write().await.insert(id.clone(), resource);
        Ok(id)
    }

    async fn find_by_id(&self, id: &str) -> bat_core::Result<Resource> {
        self.resources
            .read()
            .await
            .get(id)
            .cloned()
            .ok_or_else(|| bat_core::Error::NotFound(id.to_string()))
    }

    async fn find_by_hash(&self, hash: &str) -> bat_core::Result<Resource> {
        let resources = self.resources.read().await;
        Self::sorted_resources(&resources)
            .into_iter()
            .find(|resource| resource.entry.hash == hash)
            .ok_or_else(|| bat_core::Error::NotFound(hash.to_string()))
    }

    async fn list(&self, query: ResourceQuery) -> bat_core::Result<Vec<Resource>> {
        let resources = self.resources.read().await;
        let resources = Self::sorted_resources(&resources)
            .into_iter()
            .filter(|resource| query_matches(&query, resource))
            .collect();
        Ok(resources)
    }

    async fn update(&self, resource: Resource) -> bat_core::Result<()> {
        let mut resources = self.resources.write().await;
        if !resources.contains_key(&resource.id) {
            return Err(bat_core::Error::NotFound(resource.id));
        }

        resources.insert(resource.id.clone(), resource);
        Ok(())
    }

    async fn delete(&self, id: &str) -> bat_core::Result<()> {
        self.resources
            .write()
            .await
            .remove(id)
            .map(|_| ())
            .ok_or_else(|| bat_core::Error::NotFound(id.to_string()))
    }

    async fn count(&self, query: ResourceQuery) -> bat_core::Result<u64> {
        Ok(self.list(query).await?.len() as u64)
    }
}

/// SQLite 资源仓储实现。
#[derive(Debug, Clone)]
pub struct SqliteResourceRepository {
    pool: SqlitePool,
}

impl SqliteResourceRepository {
    /// 打开或创建 SQLite 资源仓储。
    pub async fn new(path: impl AsRef<Path>) -> bat_core::Result<Self> {
        if let Some(parent) = path.as_ref().parent() {
            tokio::fs::create_dir_all(parent).await?;
        }

        let options =
            SqliteConnectOptions::from_str(&format!("sqlite://{}", path.as_ref().display()))
                .map_err(|error| bat_core::Error::Other(error.into()))?
                .create_if_missing(true);

        let pool = SqlitePoolOptions::new()
            .max_connections(1)
            .connect_with(options)
            .await
            .map_err(|error| bat_core::Error::Other(error.into()))?;

        let repository = Self { pool };
        repository.init_schema().await?;
        Ok(repository)
    }

    async fn init_schema(&self) -> bat_core::Result<()> {
        Self::execute_query(
            &self.pool,
            sqlx::query(
                r#"
                CREATE TABLE IF NOT EXISTS resources (
                    id TEXT PRIMARY KEY NOT NULL,
                    path TEXT NOT NULL,
                    hash TEXT NOT NULL,
                    size INTEGER NOT NULL CHECK(size >= 0),
                    resource_type TEXT NOT NULL,
                    local_path TEXT NOT NULL,
                    address TEXT,
                    dependencies_json TEXT NOT NULL DEFAULT '[]'
                )
                "#,
            ),
        )
        .await?;

        Self::execute_query(
            &self.pool,
            sqlx::query(
                r#"
                CREATE INDEX IF NOT EXISTS idx_resources_hash
                ON resources(hash)
                "#,
            ),
        )
        .await?;

        Self::execute_query(
            &self.pool,
            sqlx::query(
                r#"
                CREATE INDEX IF NOT EXISTS idx_resources_type
                ON resources(resource_type)
                "#,
            ),
        )
        .await?;

        Self::execute_query(
            &self.pool,
            sqlx::query(
                r#"
                CREATE INDEX IF NOT EXISTS idx_resources_path
                ON resources(path)
                "#,
            ),
        )
        .await?;

        Ok(())
    }

    async fn execute_query<'q>(
        pool: &SqlitePool,
        query: sqlx::query::Query<'q, sqlx::Sqlite, sqlx::sqlite::SqliteArguments<'q>>,
    ) -> bat_core::Result<SqliteQueryResult> {
        query
            .execute(pool)
            .await
            .map_err(|error| bat_core::Error::Other(error.into()))
    }

    fn resource_type_to_str(resource_type: ResourceType) -> &'static str {
        match resource_type {
            ResourceType::AssetBundle => "AssetBundle",
            ResourceType::Manifest => "Manifest",
            ResourceType::TableBundle => "TableBundle",
            ResourceType::Other => "Other",
        }
    }

    fn resource_type_from_str(value: &str) -> bat_core::Result<ResourceType> {
        match value {
            "AssetBundle" => Ok(ResourceType::AssetBundle),
            "Manifest" => Ok(ResourceType::Manifest),
            "TableBundle" => Ok(ResourceType::TableBundle),
            "Other" => Ok(ResourceType::Other),
            other => Err(bat_core::Error::Serialization(format!(
                "Unknown resource type: {}",
                other
            ))),
        }
    }

    fn dependencies_to_json(dependencies: &[String]) -> bat_core::Result<String> {
        serde_json::to_string(dependencies)
            .map_err(|error| bat_core::Error::Serialization(error.to_string()))
    }

    fn dependencies_from_json(value: &str) -> bat_core::Result<Vec<String>> {
        serde_json::from_str(value)
            .map_err(|error| bat_core::Error::Serialization(error.to_string()))
    }

    fn resource_from_row(row: ResourceRow) -> bat_core::Result<Resource> {
        let (id, path, hash, size, resource_type, local_path, address, dependencies_json) = row;
        Ok(Resource {
            id,
            local_path: PathBuf::from(local_path),
            entry: ResourceEntry {
                path,
                hash,
                size: size as u64,
                resource_type: Self::resource_type_from_str(&resource_type)?,
                address,
                dependencies: Self::dependencies_from_json(&dependencies_json)?,
            },
        })
    }

    fn apply_filters<'a>(
        builder: &mut QueryBuilder<'a, Sqlite>,
        query: &'a ResourceQuery,
    ) -> bat_core::Result<()> {
        let mut has_where = false;

        if let Some(resource_type) = query.resource_type {
            push_condition_prefix(builder, &mut has_where);
            builder.push("resource_type = ");
            builder.push_bind(Self::resource_type_to_str(resource_type));
        }

        if let Some(hash) = &query.hash {
            push_condition_prefix(builder, &mut has_where);
            builder.push("hash = ");
            builder.push_bind(hash);
        }

        if let Some(pattern) = &query.path_pattern {
            push_condition_prefix(builder, &mut has_where);
            builder.push("path LIKE ");
            builder
                .push_bind(glob_to_like(pattern))
                .push(" ESCAPE '\\'");
        }

        Ok(())
    }

    async fn fetch_resources(
        &self,
        query: &ResourceQuery,
        limit: Option<usize>,
    ) -> bat_core::Result<Vec<Resource>> {
        let mut builder = QueryBuilder::<Sqlite>::new(
            "SELECT id, path, hash, size, resource_type, local_path, address, dependencies_json FROM resources",
        );
        Self::apply_filters(&mut builder, query)?;
        builder.push(" ORDER BY id");
        if let Some(limit) = limit {
            builder.push(" LIMIT ").push_bind(limit as i64);
        }

        let rows: Vec<ResourceRow> = builder
            .build_query_as()
            .fetch_all(&self.pool)
            .await
            .map_err(|error| bat_core::Error::Other(error.into()))?;

        rows.into_iter().map(Self::resource_from_row).collect()
    }

    async fn count_resources(&self, query: &ResourceQuery) -> bat_core::Result<u64> {
        let mut builder = QueryBuilder::<Sqlite>::new("SELECT COUNT(*) FROM resources");
        Self::apply_filters(&mut builder, query)?;

        let count: i64 = builder
            .build_query_scalar()
            .fetch_one(&self.pool)
            .await
            .map_err(|error| bat_core::Error::Other(error.into()))?;
        Ok(count as u64)
    }
}

#[async_trait]
impl ResourceRepository for SqliteResourceRepository {
    async fn add(&self, resource: Resource) -> bat_core::Result<String> {
        let dependencies = Self::dependencies_to_json(&resource.entry.dependencies)?;
        Self::execute_query(
            &self.pool,
            sqlx::query(
                r#"
                INSERT INTO resources (
                    id, path, hash, size, resource_type, local_path, address, dependencies_json
                )
                VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)
                ON CONFLICT(id) DO UPDATE SET
                    path = excluded.path,
                    hash = excluded.hash,
                    size = excluded.size,
                    resource_type = excluded.resource_type,
                    local_path = excluded.local_path,
                    address = excluded.address,
                    dependencies_json = excluded.dependencies_json
                "#,
            )
            .bind(resource.id.clone())
            .bind(resource.entry.path.clone())
            .bind(resource.entry.hash.clone())
            .bind(resource.entry.size as i64)
            .bind(Self::resource_type_to_str(resource.entry.resource_type))
            .bind(resource.local_path.to_string_lossy().to_string())
            .bind(resource.entry.address.clone())
            .bind(dependencies),
        )
        .await?;

        Ok(resource.id)
    }

    async fn find_by_id(&self, id: &str) -> bat_core::Result<Resource> {
        let row: Option<ResourceRow> = sqlx::query_as(
            r#"
            SELECT id, path, hash, size, resource_type, local_path, address, dependencies_json
            FROM resources
            WHERE id = ?1
            "#,
        )
        .bind(id)
        .fetch_optional(&self.pool)
        .await
        .map_err(|error| bat_core::Error::Other(error.into()))?;

        row.map(Self::resource_from_row)
            .transpose()?
            .ok_or_else(|| bat_core::Error::NotFound(id.to_string()))
    }

    async fn find_by_hash(&self, hash: &str) -> bat_core::Result<Resource> {
        let row: Option<ResourceRow> = sqlx::query_as(
            r#"
            SELECT id, path, hash, size, resource_type, local_path, address, dependencies_json
            FROM resources
            WHERE hash = ?1
            ORDER BY id
            LIMIT 1
            "#,
        )
        .bind(hash)
        .fetch_optional(&self.pool)
        .await
        .map_err(|error| bat_core::Error::Other(error.into()))?;

        row.map(Self::resource_from_row)
            .transpose()?
            .ok_or_else(|| bat_core::Error::NotFound(hash.to_string()))
    }

    async fn list(&self, query: ResourceQuery) -> bat_core::Result<Vec<Resource>> {
        self.fetch_resources(&query, None).await
    }

    async fn update(&self, resource: Resource) -> bat_core::Result<()> {
        let _existing = self.find_by_id(&resource.id).await?;
        self.add(resource).await?;
        Ok(())
    }

    async fn delete(&self, id: &str) -> bat_core::Result<()> {
        let result = Self::execute_query(
            &self.pool,
            sqlx::query("DELETE FROM resources WHERE id = ?1").bind(id),
        )
        .await?;
        if result.rows_affected() == 0 {
            return Err(bat_core::Error::NotFound(id.to_string()));
        }
        Ok(())
    }

    async fn count(&self, query: ResourceQuery) -> bat_core::Result<u64> {
        self.count_resources(&query).await
    }
}

fn push_condition_prefix<'a>(builder: &mut QueryBuilder<'a, Sqlite>, has_where: &mut bool) {
    if *has_where {
        builder.push(" AND ");
    } else {
        builder.push(" WHERE ");
        *has_where = true;
    }
}

type ResourceRow = (
    String,
    String,
    String,
    i64,
    String,
    String,
    Option<String>,
    String,
);

fn glob_to_like(pattern: &str) -> String {
    let mut escaped = String::new();
    let mut chars = pattern.chars().peekable();

    while let Some(ch) = chars.next() {
        match ch {
            '*' => {
                if matches!(chars.peek(), Some('*')) {
                    chars.next();
                }
                escaped.push('%');
            }
            '?' => escaped.push('_'),
            '%' | '_' | '\\' => {
                escaped.push('\\');
                escaped.push(ch);
            }
            other => escaped.push(other),
        }
    }

    escaped
}

fn query_matches(query: &ResourceQuery, resource: &Resource) -> bool {
    if let Some(resource_type) = query.resource_type {
        if resource.entry.resource_type != resource_type {
            return false;
        }
    }

    if let Some(hash) = &query.hash {
        if resource.entry.hash != *hash {
            return false;
        }
    }

    if let Some(pattern) = &query.path_pattern {
        if !wildcard_matches(pattern, &resource.entry.path) {
            return false;
        }
    }

    true
}

fn wildcard_matches(pattern: &str, value: &str) -> bool {
    wildcard_matches_bytes(pattern.as_bytes(), value.as_bytes())
}

fn wildcard_matches_bytes(pattern: &[u8], value: &[u8]) -> bool {
    match (pattern.first(), value.first()) {
        (None, None) => true,
        (None, Some(_)) => false,
        (Some(b'*'), _) => {
            wildcard_matches_bytes(&pattern[1..], value)
                || value
                    .first()
                    .is_some_and(|_| wildcard_matches_bytes(pattern, &value[1..]))
        }
        (Some(b'?'), Some(_)) => wildcard_matches_bytes(&pattern[1..], &value[1..]),
        (Some(pattern_byte), Some(value_byte)) if pattern_byte == value_byte => {
            wildcard_matches_bytes(&pattern[1..], &value[1..])
        }
        _ => false,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use bat_core::domain::{ResourceEntry, ResourceType};
    use std::path::PathBuf;

    fn resource(id: &str, path: &str, hash: &str, resource_type: ResourceType) -> Resource {
        Resource {
            id: id.to_string(),
            local_path: PathBuf::from(path),
            entry: ResourceEntry {
                path: path.to_string(),
                hash: hash.to_string(),
                size: 7,
                resource_type,
                address: None,
                dependencies: Vec::new(),
            },
        }
    }

    #[tokio::test]
    async fn add_find_update_and_delete_resource() {
        let repository = InMemoryResourceRepository::new();
        let id = repository
            .add(resource(
                "resource/synthetic-minimal.bundle",
                "synthetic-minimal.bundle",
                "hash-a",
                ResourceType::AssetBundle,
            ))
            .await
            .unwrap();

        assert_eq!(id, "resource/synthetic-minimal.bundle");
        assert_eq!(
            repository.find_by_id(&id).await.unwrap().entry.path,
            "synthetic-minimal.bundle"
        );

        let mut updated = repository.find_by_id(&id).await.unwrap();
        updated.entry.size = 12;
        repository.update(updated).await.unwrap();
        assert_eq!(repository.find_by_id(&id).await.unwrap().entry.size, 12);

        repository.delete(&id).await.unwrap();
        assert!(matches!(
            repository.find_by_id(&id).await,
            Err(bat_core::Error::NotFound(_))
        ));
    }

    #[tokio::test]
    async fn list_filters_by_type_hash_and_path() {
        let repository = InMemoryResourceRepository::new();
        repository
            .add(resource(
                "resource/a",
                "synthetic-minimal.bundle",
                "hash-a",
                ResourceType::AssetBundle,
            ))
            .await
            .unwrap();
        repository
            .add(resource(
                "resource/b",
                "catalog.json",
                "hash-b",
                ResourceType::Manifest,
            ))
            .await
            .unwrap();

        let query = ResourceQuery {
            resource_type: Some(ResourceType::AssetBundle),
            hash: Some("hash-a".to_string()),
            path_pattern: Some("synthetic-*.bundle".to_string()),
        };

        let results = repository.list(query).await.unwrap();
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].id, "resource/a");
        assert_eq!(
            repository.find_by_hash("hash-a").await.unwrap().id,
            "resource/a"
        );
        assert_eq!(repository.count(ResourceQuery::all()).await.unwrap(), 2);
    }

    async fn sqlite_repository() -> (tempfile::TempDir, SqliteResourceRepository) {
        let temp_dir = tempfile::tempdir().unwrap();
        let repository = SqliteResourceRepository::new(temp_dir.path().join("resources.sqlite"))
            .await
            .unwrap();
        (temp_dir, repository)
    }

    #[tokio::test]
    async fn sqlite_repository_persists_and_filters_resources() {
        let (_temp_dir, repository) = sqlite_repository().await;
        let mut resource = resource(
            "resource/sqlite-a",
            "assets/model.bundle",
            "hash-sqlite-a",
            ResourceType::AssetBundle,
        );
        resource.entry.address = Some("Character_001".to_string());
        resource
            .entry
            .dependencies
            .push("assets/shared.bundle".to_string());

        repository.add(resource.clone()).await.unwrap();

        let by_id = repository.find_by_id(&resource.id).await.unwrap();
        assert_eq!(by_id.entry.address.as_deref(), Some("Character_001"));
        assert_eq!(
            by_id.entry.dependencies,
            vec!["assets/shared.bundle".to_string()]
        );
        assert_eq!(
            repository.find_by_hash("hash-sqlite-a").await.unwrap().id,
            resource.id
        );

        let bundles = repository
            .list(ResourceQuery::by_type(ResourceType::AssetBundle))
            .await
            .unwrap();
        assert_eq!(bundles.len(), 1);

        let count = repository.count(ResourceQuery::all()).await.unwrap();
        assert_eq!(count, 1);

        repository.delete(&resource.id).await.unwrap();
        assert!(matches!(
            repository.find_by_id(&resource.id).await,
            Err(bat_core::Error::NotFound(_))
        ));
    }
}