HEDZReloaded/crates/hedz_reloaded/src/utils/trail.rs

use crate::GameState;
use bevy::prelude::*;
use serde::{Deserialize, Serialize};

#[derive(Copy, Clone, Component, Reflect, Deserialize, Serialize)]
#[reflect(Component)]
pub struct SpawnTrail {
    pub points: usize,
    pub col_start: LinearRgba,
    pub col_end: LinearRgba,
    pub width: f32,
    pub init_pos: bool,
}

impl SpawnTrail {
    pub fn new(points: usize, col_start: LinearRgba, col_end: LinearRgba, width: f32) -> Self {
        Self {
            points,
            col_start,
            col_end,
            width,
            init_pos: false,
        }
    }

    pub fn init_with_pos(mut self) -> Self {
        self.init_pos = true;
        self
    }
}

#[derive(Component, Reflect)]
#[reflect(Component)]
pub struct Trail {
    points: Vec<Vec3>,
    col_start: LinearRgba,
    col_end: LinearRgba,
}

impl Trail {
    pub fn new(trail: SpawnTrail) -> Self {
        Self {
            points: Vec::with_capacity(trail.points),
            col_start: trail.col_start,
            col_end: trail.col_end,
        }
    }

    pub fn with_pos(self, pos: Option<Vec3>) -> Self {
        let mut trail = self;
        if let Some(pos) = pos {
            trail.add(pos);
        }
        trail
    }

    pub fn add(&mut self, pos: Vec3) {
        if self.points.len() >= self.points.capacity() {
            self.points.pop();
        }
        self.points.insert(0, pos);
    }
}

pub fn plugin(app: &mut App) {
    app.add_systems(
        FixedUpdate,
        update_trail.run_if(in_state(GameState::Playing)),
    );
    #[cfg(feature = "client")]
    app.add_systems(Update, attach_trail.run_if(in_state(GameState::Playing)));
}

#[cfg(feature = "client")]
fn attach_trail(
    mut commands: Commands,
    query: Query<(Entity, &Transform, &SpawnTrail), Added<SpawnTrail>>,
    mut gizmo_assets: ResMut<Assets<GizmoAsset>>,
) {
    for (entity, transform, trail) in query.iter() {
        let width = trail.width;
        let init_pos = trail.init_pos.then_some(transform.translation);
        let id = commands
            .spawn((
                Trail::new(*trail).with_pos(init_pos),
                Gizmo {
                    handle: gizmo_assets.add(GizmoAsset::default()),
                    line_config: GizmoLineConfig { width, ..default() },
                    ..default()
                },
            ))
            .id();

        commands
            .entity(entity)
            .queue_silenced(move |mut world: EntityWorldMut| {
                world.add_child(id);
            });
    }
}

fn update_trail(
    mut query: Query<(Entity, &mut Trail, &Gizmo, &GlobalTransform)>,
    global_transform: Query<&GlobalTransform>,
    mut gizmo_assets: ResMut<Assets<GizmoAsset>>,
) -> Result {
    for (e, mut trail, gizmo, pos) in query.iter_mut() {
        trail.add(pos.translation());

        let parent_transform = global_transform.get(e)?;

        let Some(gizmo) = gizmo_assets.get_mut(gizmo.handle.id()) else {
            continue;
        };

        gizmo.clear();

        let lerp_denom = trail.points.len() as f32;

        gizmo.linestrip_gradient(trail.points.iter().enumerate().map(|(i, pos)| {
            (
                GlobalTransform::from_translation(*pos)
                    .reparented_to(parent_transform)
                    .translation,
                trail.col_start.mix(&trail.col_end, i as f32 / lerp_denom),
            )
        }));
    }

    Ok(())
}