1use serde::Serialize;
4use wowlab_types::{
5 data::{
6 TraitNode, TraitNodeEntry, TraitNodeType, TraitSelection, TraitTreeIndex,
7 TraitTreeWithSelections,
8 },
9 sim::{FastBuildHasher, FastMap, FastSet},
10};
11
12#[derive(Clone, Debug, Serialize)]
14pub struct EnrichedLoadout {
15 pub spec_id: i32,
16 pub spec_name: String,
17 pub class_name: String,
18 pub talents: Vec<EnrichedTalent>,
19 pub summary: TalentPointSummary,
20}
21
22#[derive(Clone, Debug, Serialize)]
24pub struct EnrichedTalent {
25 pub node_id: i32,
26 pub trait_node_entry_id: i32,
27 pub spell_id: i32,
28 pub override_spell_id: i32,
30 pub replaces_spell_id: i32,
32 pub name: String,
33 pub ranks: i32,
34 pub max_ranks: i32,
35 pub tree_index: i32,
36 pub sub_tree_id: i32,
37 pub tree_section: String,
38 #[serde(skip_serializing_if = "Option::is_none")]
39 pub choice_index: Option<u8>,
40 pub effect_overrides: Vec<EnrichedTalentEffectOverride>,
41}
42
43#[derive(Clone, Debug, Serialize)]
45pub struct EnrichedTalentEffectOverride {
47 pub effect_index: i32,
48 pub operation: wowlab_types::data::TraitEffectOperation,
49 pub value: f64,
50}
51
52#[derive(Clone, Debug, Serialize)]
54pub struct TalentPointSummary {
55 pub total_points: i32,
56 pub class_points: i32,
57 pub spec_points: i32,
58 pub hero_points: i32,
59}
60
61#[must_use]
63pub fn enrich_loadout(data: &TraitTreeWithSelections) -> EnrichedLoadout {
64 let mut node_by_id =
65 FastMap::with_capacity_and_hasher(data.tree.nodes.len(), FastBuildHasher::default());
66
67 for node in &data.tree.nodes {
68 node_by_id.entry(node.id).or_insert(node);
69 }
70
71 let sub_tree_names: FastMap<i32, &str> = data
72 .tree
73 .sub_trees
74 .iter()
75 .map(|st| (st.id, st.name.as_str()))
76 .collect();
77
78 let active_sub_trees = active_sub_tree_ids(data, &node_by_id);
79
80 let mut talents = Vec::with_capacity(data.selections.len());
81 let mut class_points = 0i32;
82 let mut spec_points = 0i32;
83 let mut hero_points = 0i32;
84
85 for sel in &data.selections {
86 if !sel.selected || sel.ranks_purchased <= 0 {
87 continue;
88 }
89
90 let Some(node) = node_by_id.get(&sel.node_id) else {
91 continue;
92 };
93
94 if node.sub_tree_id > 0
95 && active_sub_trees
96 .as_ref()
97 .is_some_and(|ids| !ids.contains(&node.sub_tree_id))
98 {
99 continue;
100 }
101
102 let allocations = allocate_entry_ranks(node, sel);
103
104 if allocations.is_empty() {
105 continue;
106 }
107
108 let tree_section = match TraitTreeIndex::from_dbc(node.tree_index) {
109 Some(TraitTreeIndex::Class) => {
110 class_points += sel.ranks_purchased;
111
112 "class".to_string()
113 }
114 Some(TraitTreeIndex::Specialization) => {
115 spec_points += sel.ranks_purchased;
116
117 "spec".to_string()
118 }
119 Some(TraitTreeIndex::Selection) => "selection".to_string(),
120 _ => {
121 hero_points += sel.ranks_purchased;
122
123 sub_tree_names
124 .get(&node.sub_tree_id)
125 .map_or_else(|| "hero".to_string(), |s| (*s).to_string())
126 }
127 };
128
129 for allocation in allocations {
130 let EntryAllocation {
131 entry,
132 ranks,
133 max_ranks,
134 } = allocation;
135
136 talents.push(EnrichedTalent {
137 node_id: sel.node_id,
138 trait_node_entry_id: entry.id,
139 spell_id: entry.spell_id,
140 override_spell_id: entry.override_spell_id,
141 replaces_spell_id: entry.replaces_spell_id,
142 name: entry.name.clone(),
144 ranks,
145 max_ranks,
146 tree_index: node.tree_index,
147 sub_tree_id: if TraitTreeIndex::from_dbc(node.tree_index)
148 == Some(TraitTreeIndex::Selection)
149 {
150 entry.sub_tree_id
151 } else {
152 node.sub_tree_id
153 },
154 tree_section: tree_section.clone(),
155 choice_index: sel.choice_index,
156 effect_overrides: effect_overrides(entry, ranks),
157 });
158 }
159 }
160
161 EnrichedLoadout {
162 spec_id: data.tree.spec_id,
163 spec_name: data.tree.spec_name.clone(),
164 class_name: data.tree.class_name.clone(),
165 talents,
166 summary: TalentPointSummary {
167 total_points: class_points + spec_points + hero_points,
168 class_points,
169 spec_points,
170 hero_points,
171 },
172 }
173}
174
175fn effect_overrides(entry: &TraitNodeEntry, ranks: i32) -> Vec<EnrichedTalentEffectOverride> {
176 let rank = f64::from(ranks);
177
178 entry
179 .effect_points
180 .iter()
181 .filter_map(|effect| {
182 effect
183 .rank_values
184 .iter()
185 .find(|(point_rank, _)| (*point_rank - rank).abs() < f64::EPSILON)
186 .map(|&(_, value)| EnrichedTalentEffectOverride {
187 effect_index: effect.effect_index,
188 operation: effect.operation,
189 value,
190 })
191 })
192 .collect()
193}
194
195fn active_sub_tree_ids(
196 data: &TraitTreeWithSelections,
197 node_by_id: &FastMap<i32, &TraitNode>,
198) -> Option<FastSet<i32>> {
199 let mut active =
200 FastSet::with_capacity_and_hasher(data.tree.sub_trees.len(), FastBuildHasher::default());
201 let mut has_selected_picker = false;
202
203 for sel in &data.selections {
204 if !sel.selected {
205 continue;
206 }
207
208 let Some(node) = node_by_id.get(&sel.node_id) else {
209 continue;
210 };
211
212 if TraitTreeIndex::from_dbc(node.tree_index) != Some(TraitTreeIndex::Selection) {
213 continue;
214 }
215
216 has_selected_picker = true;
217
218 let entry_idx = sel.choice_index.unwrap_or(0) as usize;
219 let Some(entry) = node.entries.get(entry_idx) else {
220 continue;
221 };
222
223 if entry.sub_tree_id > 0
224 && data
225 .tree
226 .sub_trees
227 .iter()
228 .any(|sub_tree| sub_tree.id == entry.sub_tree_id)
229 {
230 active.insert(entry.sub_tree_id);
231 }
232 }
233
234 has_selected_picker.then_some(active)
235}
236
237struct EntryAllocation<'a> {
238 entry: &'a TraitNodeEntry,
239 ranks: i32,
240 max_ranks: i32,
241}
242
243fn allocate_entry_ranks<'a>(node: &'a TraitNode, sel: &TraitSelection) -> Vec<EntryAllocation<'a>> {
244 if TraitNodeType::from_dbc(node.node_type) == Some(TraitNodeType::Tiered)
245 && node.entries.iter().any(|e| e.max_ranks > 0)
246 {
247 let mut allocations = Vec::with_capacity(node.entries.len());
248 let mut remaining = sel.ranks_purchased;
249
250 for entry in &node.entries {
251 if remaining <= 0 {
252 break;
253 }
254
255 let allocated = remaining.min(entry.max_ranks);
256
257 if allocated > 0 {
258 allocations.push(EntryAllocation {
259 entry,
260 ranks: allocated,
261 max_ranks: entry.max_ranks,
262 });
263 remaining -= allocated;
264 }
265 }
266
267 return allocations;
268 }
269
270 let entry_idx = sel.choice_index.unwrap_or(0) as usize;
271 let entry = node.entries.get(entry_idx).or_else(|| node.entries.first());
272
273 entry
274 .map(|entry| EntryAllocation {
275 entry,
276 ranks: sel.ranks_purchased,
277 max_ranks: node.max_ranks,
278 })
279 .into_iter()
280 .collect()
281}
282
283impl EnrichedLoadout {
284 #[must_use]
287 pub fn talent_selections(&self) -> Vec<(u32, u8)> {
288 self.talents
289 .iter()
290 .filter(|t| t.spell_id > 0)
291 .filter_map(|t| Some((u32::try_from(t.spell_id).ok()?, u8::try_from(t.ranks).ok()?)))
292 .collect()
293 }
294 }
296
297#[cfg(test)]
298#[path = "enrich_tests.rs"]
299mod tests;