1use std::marker::PhantomData;
2
3use bevy::{
4 ecs::system::{EntityCommand, EntityCommands},
5 prelude::{Commands, Component, Entity, World},
6};
7
8use crate::{stat_modification::ModificationType, StatData, StatIdentifier, Stats};
9
10pub struct ModifyStatEntityCommands<
12 'a,
13 StatCollection: AsMut<Stats> + Send + Sync + 'static + Component,
14> {
15 target_entity: Entity,
16 target_component: PhantomData<StatCollection>,
17 commands: Commands<'a, 'a>,
18}
19
20impl<'a, StatCollection: AsMut<Stats> + Send + Sync + 'static + Component>
21 ModifyStatEntityCommands<'a, StatCollection>
22{
23 pub fn id(&self) -> Entity {
25 self.target_entity
26 }
27
28 pub fn entity(&mut self, entity: Entity) {
30 self.target_entity = entity
31 }
32
33 pub fn commands(&mut self) -> &mut Commands<'a, 'a> {
35 &mut self.commands
36 }
37}
38
39impl<StatCollection: AsMut<Stats> + Send + Sync + 'static + Component>
40 ModifyStatEntityCommands<'_, StatCollection>
41{
42 pub fn entity_commands(&mut self) -> EntityCommands<'_> {
44 let id = self.id();
45 self.commands().entity(id)
46 }
47
48 pub fn add(
50 &mut self,
51 stat_id: impl StatIdentifier + 'static + Send + Sync,
52 stat_data: impl StatData,
53 ) -> &mut Self {
54 self.entity_commands()
55 .queue(modify_entity_stat::<StatCollection>(
56 stat_id,
57 ModificationType::add(stat_data),
58 ));
59 self
60 }
61
62 pub fn sub(
64 &mut self,
65 stat_id: impl StatIdentifier + 'static + Send + Sync,
66 stat_data: impl StatData,
67 ) -> &mut Self {
68 self.entity_commands()
69 .queue(modify_entity_stat::<StatCollection>(
70 stat_id,
71 ModificationType::sub(stat_data),
72 ));
73 self
74 }
75
76 pub fn set(
78 &mut self,
79 stat_id: impl StatIdentifier + 'static + Send + Sync,
80 stat_data: impl StatData,
81 ) -> &mut Self {
82 self.entity_commands()
83 .queue(modify_entity_stat::<StatCollection>(
84 stat_id,
85 ModificationType::set(stat_data),
86 ));
87 self
88 }
89
90 pub fn remove(&mut self, stat_id: impl StatIdentifier + 'static + Send + Sync) -> &mut Self {
92 self.entity_commands()
93 .queue(modify_entity_stat::<StatCollection>(
94 stat_id,
95 ModificationType::remove(),
96 ));
97 self
98 }
99
100 pub fn reset(&mut self, stat_id: impl StatIdentifier + 'static + Send + Sync) -> &mut Self {
102 self.entity_commands()
103 .queue(modify_entity_stat::<StatCollection>(
104 stat_id,
105 ModificationType::reset(),
106 ));
107 self
108 }
109}
110
111pub trait StatCommandsExt {
112 fn entity_stats<StatCollection: AsMut<Stats> + Send + Sync + 'static + Component>(
114 &mut self,
115 entity: Entity,
116 ) -> ModifyStatEntityCommands<'_, StatCollection>;
117
118 fn modify_stat<StatCollection: AsMut<Stats> + Send + Sync + 'static + Component>(
120 &mut self,
121 entity: Entity,
122 stat_id: impl StatIdentifier + 'static + Send + Sync,
123 modification_type: ModificationType,
124 );
125}
126
127impl<'a> StatCommandsExt for Commands<'a, 'a> {
128 fn entity_stats<StatCollection: AsMut<Stats> + Send + Sync + 'static + Component>(
130 &mut self,
131 entity: Entity,
132 ) -> ModifyStatEntityCommands<'_, StatCollection> {
133 ModifyStatEntityCommands {
134 target_entity: entity,
135 target_component: PhantomData,
136 commands: self.reborrow(),
137 }
138 }
139
140 fn modify_stat<StatCollection: AsMut<Stats> + Send + Sync + 'static + Component>(
142 &mut self,
143 entity: Entity,
144 stat_id: impl StatIdentifier + 'static + Send + Sync,
145 modification_type: ModificationType,
146 ) {
147 self.entity(entity)
148 .modify_stat::<StatCollection>(stat_id, modification_type);
149 }
150}
151
152pub trait StatEntityCommandsExt {
153 fn entity_stats<StatCollection: AsMut<Stats> + Send + Sync + 'static + Component>(
155 &mut self,
156 ) -> ModifyStatEntityCommands<'_, StatCollection>;
157
158 fn modify_stat<StatCollection: AsMut<Stats> + Send + Sync + 'static + Component>(
160 &mut self,
161 stat_id: impl StatIdentifier + 'static + Send + Sync,
162 modification_type: ModificationType,
163 );
164}
165
166impl<'a> StatEntityCommandsExt for EntityCommands<'a> {
167 fn modify_stat<StatCollection: AsMut<Stats> + Send + Sync + 'static + Component>(
168 &mut self,
169 stat_id: impl StatIdentifier + 'static + Send + Sync,
170 modification_type: ModificationType,
171 ) {
172 self.queue(modify_entity_stat::<StatCollection>(
173 stat_id,
174 modification_type,
175 ));
176 }
177
178 fn entity_stats<StatCollection: AsMut<Stats> + Send + Sync + 'static + Component>(
179 &mut self,
180 ) -> ModifyStatEntityCommands<'_, StatCollection> {
181 ModifyStatEntityCommands {
182 target_entity: self.id(),
183 target_component: PhantomData,
184 commands: self.commands(),
185 }
186 }
187}
188
189fn modify_entity_stat<StatCollection: AsMut<Stats> + Send + Sync + 'static + Component>(
190 stat_id: impl StatIdentifier + 'static + Send + Sync,
191 modification_type: ModificationType,
192) -> impl EntityCommand {
193 move |entity: Entity, world: &mut World| {
194 if let Ok(mut entity_mut) = world.get_entity_mut(entity) {
195 if let Some(mut stat_collection) = entity_mut.get_mut::<StatCollection>() {
196 let stats = stat_collection.as_mut().as_mut();
197
198 match modification_type {
199 ModificationType::Add(data) => {
200 stats.add_to_stat_manual(stat_id.identifier(), data)
201 }
202 ModificationType::Sub(data) => {
203 stats.sub_from_stat_manual(stat_id.identifier(), data)
204 }
205 ModificationType::Remove => stats.remove_stat_manual(stat_id.identifier()),
206 ModificationType::Set(data) => {
207 stats.set_stat_manual(stat_id.identifier(), data)
208 }
209 ModificationType::Reset => stats.reset_stat_manual(stat_id.identifier()),
210 }
211 }
212 }
213 }
214}
215
216#[cfg(test)]
217mod tests {
218 use super::*;
219
220 #[derive(Hash)]
221 pub struct EnemiesKilled;
222
223 impl StatIdentifier for EnemiesKilled {
224 fn identifier(&self) -> &'static str {
225 "Enemies Killed"
226 }
227 }
228
229 #[derive(Component)]
230 pub struct EntityStats {
231 stats: Stats,
232 }
233
234 impl AsMut<Stats> for EntityStats {
235 fn as_mut(&mut self) -> &mut Stats {
236 &mut self.stats
237 }
238 }
239
240 #[test]
241 fn entity_commands() {
242 let mut world = World::new();
243 let entity = world
244 .spawn(EntityStats {
245 stats: Stats::new(),
246 })
247 .id();
248
249 let mut commands = world.commands();
250
251 let mut stats = commands.entity_stats::<EntityStats>(entity);
252 stats.add(EnemiesKilled, 5u64);
253 stats.add(EnemiesKilled, 10u64);
254 world.flush();
255
256 assert_eq!(
257 *world
258 .entity(entity)
259 .get::<EntityStats>()
260 .unwrap()
261 .stats
262 .get_stat_downcast::<u64>(&EnemiesKilled)
263 .unwrap(),
264 15u64
265 );
266
267 let mut commands = world.commands();
268 let mut stats = commands.entity_stats::<EntityStats>(entity);
269 stats.reset(EnemiesKilled);
270 stats.add(EnemiesKilled, 15u64);
271 stats.sub(EnemiesKilled, 5u64);
272 stats.sub(EnemiesKilled, 7u64);
273 world.flush();
274
275 assert_eq!(
276 *world
277 .entity(entity)
278 .get::<EntityStats>()
279 .unwrap()
280 .stats
281 .get_stat_downcast::<u64>(&EnemiesKilled)
282 .unwrap(),
283 3u64
284 );
285
286 let mut commands = world.commands();
287 let mut stats = commands.entity_stats::<EntityStats>(entity);
288 stats.set(EnemiesKilled, 7u64);
289 world.flush();
290
291 assert_eq!(
292 *world
293 .entity(entity)
294 .get::<EntityStats>()
295 .unwrap()
296 .stats
297 .get_stat_downcast::<u64>(&EnemiesKilled)
298 .unwrap(),
299 7u64
300 );
301
302 let mut commands = world.commands();
303 let mut stats = commands.entity_stats::<EntityStats>(entity);
304 stats.remove(EnemiesKilled);
305 world.flush();
306
307 assert_eq!(
308 world
309 .entity(entity)
310 .get::<EntityStats>()
311 .unwrap()
312 .stats
313 .get_stat_downcast::<u64>(&EnemiesKilled),
314 None
315 );
316 }
317}