bevy_vrm 0.4.0

use std::{
    collections::{
        HashMap,
        HashSet,
    },
    sync::LazyLock,
};

use bevy::{
    camera::visibility::RenderLayers,
    mesh::{
        Indices,
        VertexAttributeValues,
        morph::MeshMorphWeights,
        skinning::SkinnedMesh,
    },
    prelude::*,
};
use bevy_shader_mtoon::{
    MtoonMaterial,
    OutlineSync,
};
use serde_vrm::vrm0::BoneName;
pub use serde_vrm::vrm0::FirstPersonFlag;

pub const FIRST_PERSON_LAYER: usize = 7;
pub const THIRD_PERSON_LAYER: usize = 8;

pub static DEFAULT_RENDER_LAYERS: LazyLock<HashMap<FirstPersonFlag, RenderLayers>> =
    LazyLock::new(|| {
        let mut map = HashMap::default();

        map.insert(
            FirstPersonFlag::Auto,
            RenderLayers::from_layers(&[0, FIRST_PERSON_LAYER, THIRD_PERSON_LAYER]),
        );
        map.insert(
            FirstPersonFlag::Both,
            RenderLayers::from_layers(&[0, FIRST_PERSON_LAYER, THIRD_PERSON_LAYER]),
        );
        map.insert(
            FirstPersonFlag::FirstPersonOnly,
            RenderLayers::layer(FIRST_PERSON_LAYER),
        );
        map.insert(
            FirstPersonFlag::ThirdPersonOnly,
            RenderLayers::layer(THIRD_PERSON_LAYER),
        );

        map
    });

/// Splits each VRM mesh into first- and third-person variants and assigns their
/// [`RenderLayers`]. Cameras must render
/// [`FIRST_PERSON_LAYER`]/[`THIRD_PERSON_LAYER`] to see the split meshes.
#[derive(EntityEvent)]
pub struct SetupFirstPerson {
    pub entity:        Entity,
    /// Render layer values to use for each first person flag.
    /// Defaults to [`DEFAULT_RENDER_LAYERS`] if not provided.
    pub render_layers: Option<HashMap<FirstPersonFlag, RenderLayers>>,
}

pub(crate) fn setup_first_person(
    event: On<SetupFirstPerson>,
    bones: Query<(Entity, &BoneName)>,
    mut flags: Query<(
        Entity,
        &mut FirstPersonFlag,
        &Mesh3d,
        Option<&Name>,
        Option<&MeshMaterial3d<StandardMaterial>>,
        Option<&MeshMaterial3d<MtoonMaterial>>,
        Option<&MeshMorphWeights>,
        Option<&OutlineSync>,
    )>,
    mut commands: Commands,
    mut meshes: ResMut<Assets<Mesh>>,
    parents: Query<&ChildOf>,
    skins: Query<&SkinnedMesh>,
) {
    if bones.is_empty() {
        return;
    }

    let render_layers = event
        .render_layers
        .as_ref()
        .unwrap_or(&DEFAULT_RENDER_LAYERS);

    let (head_ent, _) = bones
        .iter()
        .find(|(e, name)| **name == BoneName::Head && is_child(*e, event.entity, &parents))
        .expect("VRM must have Head bone");

    for (
        ent,
        mut flag,
        mesh_handle,
        name,
        standard_material,
        mtoon_material,
        morph_weights,
        outline_sync,
    ) in &mut flags
    {
        // Skip if not a child of this VRM entity.
        if !is_child(ent, event.entity, &parents) {
            continue;
        }
        // Auto splits into a head-removed first-person variant plus the full
        // third-person original, each on its own render layer.
        if *flag == FirstPersonFlag::Auto {
            let Some(mesh) = meshes.get(mesh_handle) else {
                warn!("Mesh not found");
                continue;
            };
            let Ok(skin) = skins.get(ent) else {
                continue;
            };
            let Some(first_person) = strip_head_vertices(mesh, skin, head_ent, &parents) else {
                continue;
            };

            let new_ent = commands
                .spawn((
                    Transform::default(),
                    Mesh3d(meshes.add(first_person)),
                    render_layers[&FirstPersonFlag::FirstPersonOnly].clone(),
                ))
                .id();

            insert_cloned(&mut commands, new_ent, mtoon_material);
            insert_cloned(&mut commands, new_ent, standard_material);
            insert_cloned(&mut commands, new_ent, name);
            insert_cloned(&mut commands, new_ent, morph_weights);
            insert_cloned(&mut commands, new_ent, outline_sync);

            let mut new_skin = skin.clone();
            for (i, e) in new_skin.joints.iter().enumerate() {
                if *e == ent {
                    new_skin.joints.insert(i, new_ent);
                    break;
                }
            }

            commands.entity(new_ent).insert(new_skin);
            commands.entity(ent).add_child(new_ent);

            *flag = FirstPersonFlag::ThirdPersonOnly;
        }

        commands
            .entity(ent)
            .insert(render_layers[flag.as_ref()].clone());
    }
}

fn insert_cloned<T: Component + Clone>(
    commands: &mut Commands,
    entity: Entity,
    component: Option<&T>,
) {
    if let Some(component) = component {
        commands.entity(entity).insert(component.clone());
    }
}

/// Clones a mesh with every triangle touching a head-weighted vertex removed.
fn strip_head_vertices(
    mesh: &Mesh,
    skin: &SkinnedMesh,
    head_ent: Entity,
    parents: &Query<&ChildOf>,
) -> Option<Mesh> {
    let Some(VertexAttributeValues::Uint16x4(joints)) = mesh.attribute(Mesh::ATTRIBUTE_JOINT_INDEX)
    else {
        return None;
    };
    let Some(VertexAttributeValues::Float32x4(weights)) =
        mesh.attribute(Mesh::ATTRIBUTE_JOINT_WEIGHT)
    else {
        return None;
    };

    let mut to_remove = HashSet::<usize>::default();
    for (i, item) in joints.iter().enumerate() {
        for (j, idx) in item.iter().enumerate() {
            let joint_ent = skin.joints[*idx as usize];
            if is_child(joint_ent, head_ent, parents) && weights[i][j] > 0.0 {
                to_remove.insert(i);
            }
        }
    }

    let mut to_remove = to_remove.into_iter().collect::<Vec<_>>();
    to_remove.sort_by(|a, b| b.cmp(a));

    let mut mesh = mesh.clone();
    if let Some(indices) = mesh.indices_mut() {
        match indices {
            Indices::U16(vec) => clean_indices(vec, &to_remove),
            Indices::U32(vec) => clean_indices(vec, &to_remove),
        }
    }
    Some(mesh)
}

trait ToUsize {
    fn to_usize(self) -> usize;
}

impl ToUsize for u16 {
    fn to_usize(self) -> usize {
        self as usize
    }
}

impl ToUsize for u32 {
    fn to_usize(self) -> usize {
        self as usize
    }
}

/// Remove the specified vertices from the indices.
fn clean_indices<T: Copy + PartialEq + ToUsize>(indices: &mut Vec<T>, vertices: &[usize]) {
    let mut to_remove = Vec::default();

    for (i, chunk) in indices.chunks(3).enumerate() {
        for n in chunk {
            if vertices.contains(&n.to_usize()) {
                to_remove.push(i);
                break;
            }
        }
    }

    for i in to_remove.into_iter().rev() {
        let start = i * 3;
        indices.remove(start);
        indices.remove(start);
        indices.remove(start);
    }
}

/// Walks up the parent tree, searching for a specific Entity.
fn is_child(target_child: Entity, target_parent: Entity, parents: &Query<&ChildOf>) -> bool {
    if target_child == target_parent {
        true
    } else if let Ok(child_of) = parents.get(target_child) {
        is_child(child_of.parent(), target_parent, parents)
    } else {
        false
    }
}