shinespark/src/blend_asset.rs

//! Blender `.blend` source processing.
//!
//! Blender is used as an offline importer. The processed result is a normal
//! GLB, so the runtime continues to use Bevy's built-in glTF loader.

use bevy::{
    asset::{
        AssetLoader, AssetPath, AsyncWriteExt, LoadContext,
        io::{Reader, Writer},
        processor::LoadTransformAndSave,
        saver::{AssetSaver, SavedAsset},
        transformer::{AssetTransformer, TransformedAsset},
    },
    gltf::GltfLoader,
    prelude::*,
    reflect::TypePath,
};
use serde::Deserialize;
use std::{
    collections::HashMap,
    fs, io,
    path::{Path, PathBuf},
    process::Command,
    time::{SystemTime, UNIX_EPOCH},
};

use crate::fighter::animation::AnimKind;

const EXPORT_SCRIPT: &str = include_str!("blend_asset/export.py");

#[derive(Asset, TypePath)]
struct BlendSource {
    glb: Vec<u8>,
}

#[derive(Default, TypePath)]
struct BlendSourceLoader;

#[derive(Debug, Deserialize)]
struct ExportManifest {
    dependencies: Vec<String>,
}

#[derive(Debug, Deserialize)]
struct AnimationManifest {
    animations: HashMap<AnimKind, String>,
}

impl AssetLoader for BlendSourceLoader {
    type Asset = BlendSource;
    type Settings = ();
    type Error = io::Error;

    async fn load(
        &self,
        reader: &mut dyn Reader,
        _: &Self::Settings,
        load_context: &mut LoadContext<'_>,
    ) -> Result<Self::Asset, Self::Error> {
        let mut source_bytes = Vec::new();
        reader.read_to_end(&mut source_bytes).await?;

        let temp = TempDir::new()?;
        let source_path = locate_source(load_context.path().path());
        let blend_path = if let Some(path) = source_path.as_deref() {
            path.to_path_buf()
        } else {
            let path = temp.path().join("source.blend");
            fs::write(&path, source_bytes)?;
            path
        };
        let output = temp.path().join("export.gltf");
        let script = temp.path().join("export.py");
        let actions = referenced_actions(load_context, &asset_root()).await?;
        let actions_path = temp.path().join("actions.json");
        fs::write(
            &actions_path,
            serde_json::to_vec(&actions).map_err(io::Error::other)?,
        )?;
        fs::write(&script, EXPORT_SCRIPT)?;

        let result = Command::new(std::env::var_os("BLENDER").unwrap_or_else(|| "blender".into()))
            .arg("--background")
            .arg(&blend_path)
            .arg("--python")
            .arg(&script)
            .arg("--")
            .arg(&output)
            .arg(&actions_path)
            .arg(&blend_path)
            .output()
            .map_err(|error| io::Error::other(format!("failed to start Blender: {error}")))?;
        if !result.status.success() {
            return Err(io::Error::other(format!(
                "Blender export failed with {}:\n{}\n{}",
                result.status,
                String::from_utf8_lossy(&result.stdout),
                String::from_utf8_lossy(&result.stderr)
            )));
        }

        let manifest: ExportManifest =
            serde_json::from_slice(&fs::read(temp.path().join("export-manifest.json"))?)
                .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?;

        let asset_root = asset_root();
        let mut dependencies = manifest.dependencies;
        let document: serde_json::Value = serde_json::from_slice(&fs::read(&output)?)
            .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?;
        validate_export(&document, &output)?;
        if let Some(images) = document.get("images").and_then(serde_json::Value::as_array) {
            for image in images {
                if let Some(uri) = image.get("uri").and_then(serde_json::Value::as_str) {
                    let path = temp.path().join(uri);
                    dependencies.push(path.to_string_lossy().into_owned());
                }
            }
        }

        let mut seen = std::collections::HashSet::new();
        for dependency in dependencies {
            let physical = PathBuf::from(&dependency);
            let Some(asset_path) = physical_to_asset_path(&physical, &asset_root) else {
                continue;
            };
            if seen.insert(asset_path.clone()) {
                // This read is intentionally discarded: it registers the file and
                // its processed hash as a processor dependency.
                load_context
                    .read_asset_bytes(asset_path)
                    .await
                    .map_err(|error| {
                        io::Error::other(format!("failed to read blend dependency: {error}"))
                    })?;
            }
        }

        let glb = make_glb(&output, &temp, load_context.path().path(), &asset_root)?;
        Ok(BlendSource { glb })
    }

    fn extensions(&self) -> &[&str] {
        &["blend"]
    }
}

fn validate_export(document: &serde_json::Value, output: &Path) -> io::Result<()> {
    for key in ["scenes", "nodes", "meshes"] {
        let count = document
            .get(key)
            .and_then(serde_json::Value::as_array)
            .map_or(0, Vec::len);
        if count == 0 {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                format!("Blender export {} contains no {key}", output.display()),
            ));
        }
    }
    Ok(())
}

async fn referenced_actions(
    load_context: &mut LoadContext<'_>,
    root: &Path,
) -> io::Result<Vec<String>> {
    let model_path = load_context.path().path();
    let Some(stem) = model_path.file_stem().and_then(|stem| stem.to_str()) else {
        return Ok(Vec::new());
    };
    let manifest_path = model_path
        .parent()
        .unwrap_or_else(|| Path::new(""))
        .join(format!("{stem}.fighter.ron"));
    if !root.join(&manifest_path).is_file() {
        return Ok(Vec::new());
    }
    let bytes = load_context
        .read_asset_bytes(manifest_path)
        .await
        .map_err(|error| io::Error::other(error.to_string()))?;
    let manifest: AnimationManifest = ron::de::from_bytes(&bytes)
        .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?;
    Ok(manifest.animations.into_values().collect())
}

#[derive(Asset, TypePath)]
struct ProcessedBlend {
    bytes: Vec<u8>,
}

#[derive(Default, TypePath)]
struct BlendTransformer;

impl AssetTransformer for BlendTransformer {
    type AssetInput = BlendSource;
    type AssetOutput = ProcessedBlend;
    type Settings = ();
    type Error = io::Error;

    async fn transform<'a>(
        &'a self,
        source: TransformedAsset<Self::AssetInput>,
        _: &'a Self::Settings,
    ) -> Result<TransformedAsset<Self::AssetOutput>, Self::Error> {
        let bytes = source.glb.clone();
        Ok(source.replace_asset(ProcessedBlend { bytes }))
    }
}

#[derive(Default, TypePath)]
struct BlendSaver;

impl AssetSaver for BlendSaver {
    type Asset = ProcessedBlend;
    type Settings = ();
    type OutputLoader = GltfLoader;
    type Error = io::Error;

    async fn save(
        &self,
        writer: &mut Writer,
        asset: SavedAsset<'_, '_, Self::Asset>,
        _: &Self::Settings,
        _: AssetPath<'_>,
    ) -> Result<<Self::OutputLoader as AssetLoader>::Settings, Self::Error> {
        writer.write_all(&asset.bytes).await?;
        Ok(default())
    }
}

type BlendProcessor = LoadTransformAndSave<BlendSourceLoader, BlendTransformer, BlendSaver>;

pub struct BlendAssetProcessingPlugin;

impl Plugin for BlendAssetProcessingPlugin {
    fn build(&self, app: &mut App) {
        app.init_asset::<BlendSource>()
            .init_asset::<ProcessedBlend>()
            .register_asset_loader(BlendSourceLoader)
            .register_asset_processor(BlendProcessor::new(BlendTransformer, BlendSaver))
            .set_default_asset_processor::<BlendProcessor>("blend");
    }
}

struct TempDir(PathBuf);

impl TempDir {
    fn new() -> io::Result<Self> {
        let stamp = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_nanos();
        let path =
            std::env::temp_dir().join(format!("shinespark-blend-{}-{stamp}", std::process::id()));
        fs::create_dir_all(&path)?;
        Ok(Self(path))
    }
    fn path(&self) -> &Path {
        &self.0
    }
}

impl Drop for TempDir {
    fn drop(&mut self) {
        let _ = fs::remove_dir_all(&self.0);
    }
}

fn asset_root() -> PathBuf {
    std::env::var_os("BEVY_ASSET_ROOT")
        .map(PathBuf::from)
        .unwrap_or_else(|| std::env::current_dir().unwrap_or_default().join("assets"))
}

fn locate_source(path: &Path) -> Option<PathBuf> {
    let root = asset_root();
    [root.join(path), PathBuf::from(path)]
        .into_iter()
        .find(|candidate| candidate.is_file())
}

fn physical_to_asset_path(path: &Path, root: &Path) -> Option<String> {
    let canonical = fs::canonicalize(path).ok()?;
    let root = fs::canonicalize(root).ok()?;
    let relative = canonical.strip_prefix(root).ok()?;
    Some(relative.to_string_lossy().replace('\\', "/"))
}

fn make_glb(
    gltf_path: &Path,
    temp: &TempDir,
    model_asset_path: &Path,
    asset_root: &Path,
) -> io::Result<Vec<u8>> {
    let mut json: serde_json::Value = serde_json::from_slice(&fs::read(gltf_path)?)
        .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?;
    let bin_name = json["buffers"][0]["uri"].as_str().ok_or_else(|| {
        io::Error::new(io::ErrorKind::InvalidData, "glTF has no binary buffer URI")
    })?;
    let mut bin = fs::read(temp.path().join(bin_name))?;
    let model_parent = model_asset_path.parent().unwrap_or_else(|| Path::new(""));
    let image_count = json["images"].as_array().map_or(0, Vec::len);
    if image_count > 0 && !json["bufferViews"].is_array() {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            "glTF bufferViews are not an array",
        ));
    }
    for index in 0..image_count {
        let uri = json["images"][index]
            .get("uri")
            .and_then(serde_json::Value::as_str)
            .map(str::to_owned);
        let Some(uri) = uri else { continue };
        if uri.starts_with("data:") {
            continue;
        }
        let uri_path = temp.path().join(uri);
        if let Some(asset_path) = physical_to_asset_path(&uri_path, asset_root) {
            json["images"][index]["uri"] =
                serde_json::Value::String(relative_uri(model_parent, Path::new(&asset_path)));
        } else {
            let bytes = fs::read(&uri_path)?;
            while bin.len() % 4 != 0 {
                bin.push(0);
            }
            let offset = bin.len();
            bin.extend_from_slice(&bytes);
            let mime = mime_type(&uri_path);
            let view = json["bufferViews"].as_array().unwrap().len();
            json["bufferViews"]
                .as_array_mut()
                .unwrap()
                .push(serde_json::json!({
                    "buffer": 0,
                    "byteOffset": offset,
                    "byteLength": bytes.len(),
                }));
            json["images"][index].as_object_mut().unwrap().remove("uri");
            json["images"][index]["bufferView"] = serde_json::Value::from(view);
            json["images"][index]["mimeType"] = serde_json::Value::String(mime.to_owned());
        }
    }
    json["buffers"][0]["byteLength"] = serde_json::Value::from(bin.len());
    json["buffers"][0].as_object_mut().unwrap().remove("uri");
    let mut json_bytes = serde_json::to_vec(&json).map_err(io::Error::other)?;
    while json_bytes.len() % 4 != 0 {
        json_bytes.push(b' ');
    }
    while bin.len() % 4 != 0 {
        bin.push(0);
    }
    let length = 12 + 8 + json_bytes.len() + 8 + bin.len();
    let mut glb = Vec::with_capacity(length);
    glb.extend_from_slice(b"glTF");
    glb.extend_from_slice(&2u32.to_le_bytes());
    glb.extend_from_slice(&(length as u32).to_le_bytes());
    glb.extend_from_slice(&(json_bytes.len() as u32).to_le_bytes());
    glb.extend_from_slice(b"JSON");
    glb.extend_from_slice(&json_bytes);
    glb.extend_from_slice(&(bin.len() as u32).to_le_bytes());
    glb.extend_from_slice(b"BIN\0");
    glb.extend_from_slice(&bin);
    Ok(glb)
}

fn relative_uri(from: &Path, to: &Path) -> String {
    let from: Vec<_> = from.components().collect();
    let to: Vec<_> = to.components().collect();
    let common = from.iter().zip(&to).take_while(|(a, b)| a == b).count();
    let mut result = PathBuf::new();
    for _ in common..from.len() {
        result.push("..");
    }
    for component in &to[common..] {
        result.push(component.as_os_str());
    }
    result.to_string_lossy().replace('\\', "/")
}

fn mime_type(path: &Path) -> &'static str {
    match path
        .extension()
        .and_then(|e| e.to_str())
        .unwrap_or_default()
        .to_ascii_lowercase()
        .as_str()
    {
        "jpg" | "jpeg" => "image/jpeg",
        "webp" => "image/webp",
        _ => "image/png",
    }
}