1use schemars::JsonSchema;
2use serde::{Deserialize, Serialize};
3use serde_json::Value;
4use sqlx::{PgPool, Row};
5use wowlab_types::{
6 data::{ItemEffect, SpellEffect},
7 sim::{FastMap, FastSet},
8};
9
10const MAX_IDS: usize = 50;
11
12#[derive(Debug, Default, Deserialize, JsonSchema)]
13#[schemars(
14 description = "Resolve spell and/or item effects. Provide at least one of spell_ids or item_ids (max 50 total)."
15)]
16pub(super) struct ResolveEffectsInput {
17 #[serde(default)]
18 #[schemars(description = "Spell IDs to resolve effects for.")]
19 pub spell_ids: Option<Vec<i32>>,
20 #[serde(default)]
21 #[schemars(description = "Item IDs to resolve effects for.")]
22 pub item_ids: Option<Vec<i32>>,
23}
24
25#[derive(Debug, JsonSchema, Serialize)]
26pub(super) struct SpellWithEffects {
27 pub id: i32,
28 pub name: String,
29 pub school_mask: i32,
30 pub duration: i32,
31 pub effects: Vec<SpellEffect>,
32 pub trigger_spells: Vec<SpellWithEffects>,
33}
34
35#[derive(Debug, JsonSchema, Serialize)]
36pub(super) struct ItemWithEffects {
37 pub id: i32,
38 pub name: String,
39 pub effects: Vec<ItemEffect>,
40 pub effect_spells: Vec<SpellWithEffects>,
41}
42
43#[derive(Debug, JsonSchema, Serialize)]
44pub(super) struct ResolveEffectsOutput {
45 pub spells: Vec<SpellWithEffects>,
46 pub items: Vec<ItemWithEffects>,
47}
48
49pub(super) async fn execute(
50 db: &PgPool,
51 input: ResolveEffectsInput,
52) -> Result<ResolveEffectsOutput, super::ToolError> {
53 let spell_ids = input.spell_ids.unwrap_or_default();
54 let item_ids = input.item_ids.unwrap_or_default();
55
56 if spell_ids.is_empty() && item_ids.is_empty() {
57 return Err(super::ToolError::invalid_params(
58 "Provide at least one spell_id or item_id",
59 ));
60 }
61
62 if spell_ids.len() + item_ids.len() > MAX_IDS {
63 return Err(super::ToolError::invalid_params(format!(
64 "Max {MAX_IDS} total IDs per request"
65 )));
66 }
67
68 let (spell_rows, item_rows) = futures::future::try_join(
69 async {
70 if spell_ids.is_empty() {
71 Ok(vec![])
72 } else {
73 fetch_spells(db, &spell_ids).await
74 }
75 },
76 async {
77 if item_ids.is_empty() {
78 Ok(vec![])
79 } else {
80 fetch_items(db, &item_ids).await
81 }
82 },
83 )
84 .await?;
85
86 let trigger_rows = fetch_related_spells(db, &spell_rows, &item_rows).await?;
87
88 Ok(assemble_output(&spell_rows, &item_rows, &trigger_rows))
89}
90
91async fn fetch_related_spells(
92 db: &PgPool,
93 spell_rows: &[SpellRow],
94 item_rows: &[ItemRow],
95) -> Result<Vec<SpellRow>, super::ToolError> {
96 let mut visited: FastSet<i32> = spell_rows.iter().map(|row| row.id).collect();
97 let mut pending = collect_related_ids(spell_rows, item_rows, &visited);
98 let mut related = Vec::new();
99
100 while !pending.is_empty() {
101 let ids = pending.drain().collect::<Vec<_>>();
102
103 visited.extend(ids.iter().copied());
104
105 let rows = fetch_spells(db, &ids).await?;
106
107 pending = collect_related_ids(&rows, &[], &visited);
108 related.extend(rows);
109 }
110
111 Ok(related)
112}
113
114fn collect_related_ids(
115 spell_rows: &[SpellRow],
116 item_rows: &[ItemRow],
117 visited: &FastSet<i32>,
118) -> FastSet<i32> {
119 let mut related_ids = FastSet::default();
120
121 for row in spell_rows {
122 for eff in &row.effects {
123 if eff.trigger_spell != 0 && !visited.contains(&eff.trigger_spell) {
124 related_ids.insert(eff.trigger_spell);
125 }
126 }
127 }
128
129 for row in item_rows {
130 for eff in &row.effects {
131 if eff.spell_id != 0 && !visited.contains(&eff.spell_id) {
132 related_ids.insert(eff.spell_id);
133 }
134 }
135 }
136
137 related_ids
138}
139
140fn assemble_output(
141 spell_rows: &[SpellRow],
142 item_rows: &[ItemRow],
143 trigger_rows: &[SpellRow],
144) -> ResolveEffectsOutput {
145 let mut trigger_map: FastMap<i32, &SpellRow> = spell_rows.iter().map(|r| (r.id, r)).collect();
146
147 for r in trigger_rows {
148 trigger_map.entry(r.id).or_insert(r);
149 }
150
151 let spells = spell_rows
152 .iter()
153 .map(|row| to_spell_with_effects(row, &trigger_map))
154 .collect();
155
156 let items = item_rows
157 .iter()
158 .map(|row| {
159 let mut seen = FastSet::default();
160 let unique_effects = row
161 .effects
162 .iter()
163 .filter(|e| e.spell_id != 0 && seen.insert(e.spell_id))
164 .filter_map(|e| trigger_map.get(&e.spell_id));
165 let effect_spells = unique_effects
166 .map(|row| to_spell_with_effects(row, &trigger_map))
167 .collect();
168
169 ItemWithEffects {
170 id: row.id,
171 name: row.name.clone(),
172 effects: row.effects.clone(),
173 effect_spells,
174 }
175 })
176 .collect();
177
178 ResolveEffectsOutput { spells, items }
179}
180
181fn to_spell_with_effects(
182 row: &SpellRow,
183 trigger_map: &FastMap<i32, &SpellRow>,
184) -> SpellWithEffects {
185 let mut ancestors = FastSet::default();
186
187 ancestors.insert(row.id);
188
189 let mut stack = vec![SpellTreeFrame::new(row, trigger_map)];
190
191 while !stack.is_empty() {
192 let next_trigger = stack.last_mut().and_then(SpellTreeFrame::next_trigger);
193
194 if let Some(trigger) = next_trigger {
195 if ancestors.insert(trigger.id) {
196 stack.push(SpellTreeFrame::new(trigger, trigger_map));
197 } else if let Some(frame) = stack.last_mut() {
198 frame.children.push(to_spell_leaf(trigger));
199 }
200
201 continue;
202 }
203
204 let Some(frame) = stack.pop() else {
205 break;
206 };
207
208 ancestors.remove(&frame.row.id);
209
210 let spell = to_spell_node(frame.row, frame.children);
211
212 if let Some(parent) = stack.last_mut() {
213 parent.children.push(spell);
214 } else {
215 return spell;
216 }
217 }
218
219 to_spell_leaf(row)
220}
221
222struct SpellTreeFrame<'a> {
223 row: &'a SpellRow,
224 triggers: Vec<&'a SpellRow>,
225 next: usize,
226 children: Vec<SpellWithEffects>,
227}
228
229impl<'a> SpellTreeFrame<'a> {
230 fn new(row: &'a SpellRow, trigger_map: &FastMap<i32, &'a SpellRow>) -> Self {
231 let triggers = unique_trigger_rows(row, trigger_map);
232 let children = Vec::with_capacity(triggers.len());
233
234 Self {
235 row,
236 triggers,
237 next: 0,
238 children,
239 }
240 }
241
242 fn next_trigger(&mut self) -> Option<&'a SpellRow> {
243 let trigger = self.triggers.get(self.next).copied();
244
245 self.next += usize::from(trigger.is_some());
246
247 trigger
248 }
249}
250
251fn unique_trigger_rows<'a>(
252 row: &SpellRow,
253 trigger_map: &FastMap<i32, &'a SpellRow>,
254) -> Vec<&'a SpellRow> {
255 let mut seen = FastSet::default();
256 let mut triggers = Vec::with_capacity(row.effects.len());
257
258 for effect in &row.effects {
259 if effect.trigger_spell == 0 || !seen.insert(effect.trigger_spell) {
260 continue;
261 }
262
263 if let Some(trigger) = trigger_map.get(&effect.trigger_spell) {
264 triggers.push(*trigger);
265 }
266 }
267
268 triggers
269}
270
271fn to_spell_node(row: &SpellRow, trigger_spells: Vec<SpellWithEffects>) -> SpellWithEffects {
272 SpellWithEffects {
273 id: row.id,
274 name: row.name.clone(),
275 school_mask: row.school_mask,
276 duration: row.duration,
277 effects: row.effects.clone(),
278 trigger_spells,
279 }
280}
281
282fn to_spell_leaf(row: &SpellRow) -> SpellWithEffects {
283 SpellWithEffects {
284 id: row.id,
285 name: row.name.clone(),
286 school_mask: row.school_mask,
287 duration: row.duration,
288 effects: row.effects.clone(),
289 trigger_spells: vec![],
290 }
291}
292
293struct SpellRow {
295 id: i32,
296 name: String,
297 school_mask: i32,
298 duration: i32,
299 effects: Vec<SpellEffect>,
300}
301
302struct ItemRow {
303 id: i32,
304 name: String,
305 effects: Vec<ItemEffect>,
306}
307
308async fn fetch_spells(db: &PgPool, ids: &[i32]) -> Result<Vec<SpellRow>, super::ToolError> {
309 let rows = sqlx::query(
310 "SELECT id, name, school_mask, duration, effects FROM game.spells WHERE id = ANY($1)",
311 )
312 .bind(ids)
313 .fetch_all(db)
314 .await
315 .map_err(|error| super::ToolError::database("Failed to fetch spells", error))?;
316
317 rows.iter()
318 .map(|r| {
319 let effects_json: Value = r
320 .try_get("effects")
321 .unwrap_or_else(|_| Value::Array(vec![]));
322 let effects: Vec<SpellEffect> =
323 serde_json::from_value(effects_json).unwrap_or_default();
324
325 Ok(SpellRow {
326 id: r.try_get("id").map_err(|error| {
327 super::ToolError::database("Failed to decode spell ID", error)
328 })?,
329 name: r.try_get("name").map_err(|error| {
330 super::ToolError::database("Failed to decode spell name", error)
331 })?,
332 school_mask: r.try_get("school_mask").unwrap_or(0),
333 duration: r.try_get("duration").unwrap_or(0),
334 effects,
335 })
336 })
337 .collect()
338}
339
340async fn fetch_items(db: &PgPool, ids: &[i32]) -> Result<Vec<ItemRow>, super::ToolError> {
341 let rows = sqlx::query("SELECT id, name, effects FROM game.items WHERE id = ANY($1)")
342 .bind(ids)
343 .fetch_all(db)
344 .await
345 .map_err(|error| super::ToolError::database("Failed to fetch items", error))?;
346
347 rows.iter()
348 .map(|r| {
349 let effects_json: Value = r
350 .try_get("effects")
351 .unwrap_or_else(|_| Value::Array(vec![]));
352 let effects: Vec<ItemEffect> = serde_json::from_value(effects_json).unwrap_or_default();
353
354 Ok(ItemRow {
355 id: r.try_get("id").map_err(|error| {
356 super::ToolError::database("Failed to decode item ID", error)
357 })?,
358 name: r.try_get("name").map_err(|error| {
359 super::ToolError::database("Failed to decode item name", error)
360 })?,
361 effects,
362 })
363 })
364 .collect()
365}
366
367pub(super) async fn handle(
368 db: &PgPool,
369 params: ResolveEffectsInput,
370) -> Result<rmcp::model::CallToolResult, rmcp::ErrorData> {
371 Ok(crate::mcp::tool_result(
372 crate::telemetry::McpTool::ResolveEffects,
373 execute(db, params).await,
374 ))
375}
376
377#[cfg(test)]
378mod tests {
379 use googletest::prelude::*;
380 use serde_json::json;
381
382 use super::*;
383
384 fn spell_row(id: i32, trigger_ids: &[i32]) -> SpellRow {
385 SpellRow {
386 id,
387 name: format!("Spell {id}"),
388 school_mask: 1,
389 duration: 0,
390 effects: trigger_ids
391 .iter()
392 .map(|trigger_spell| SpellEffect {
393 trigger_spell: *trigger_spell,
394 ..SpellEffect::default()
395 })
396 .collect(),
397 }
398 }
399
400 #[gtest]
401 fn recursive_effect_output_conforms_to_advertised_schema() -> Result<()> {
402 let leaf = SpellWithEffects {
403 id: 228_597,
404 name: "Frostbolt".to_owned(),
405 school_mask: 16,
406 duration: 0,
407 effects: Vec::new(),
408 trigger_spells: Vec::new(),
409 };
410 let output = ResolveEffectsOutput {
411 spells: vec![SpellWithEffects {
412 id: 116,
413 name: "Frostbolt".to_owned(),
414 school_mask: 16,
415 duration: 0,
416 effects: Vec::new(),
417 trigger_spells: vec![leaf],
418 }],
419 items: Vec::new(),
420 };
421 let serialized = serde_json::to_value(output)?;
422 let schema = Value::Object(
423 crate::mcp::output_schema_for::<ResolveEffectsOutput>()
424 .as_ref()
425 .clone(),
426 );
427 let validator = jsonschema::validator_for(&schema).or_fail()?;
428 let errors = validator
429 .iter_errors(&serialized)
430 .map(|error| error.to_string())
431 .collect::<Vec<_>>();
432
433 verify_that!(serialized.pointer("/items"), some(eq(&json!([]))))?;
434 verify_that!(
435 serialized.pointer("/spells/0/trigger_spells/0/trigger_spells"),
436 some(eq(&json!([])))
437 )?;
438
439 verify_that!(errors, is_empty())
440 }
441
442 #[gtest]
443 fn item_effect_spell_includes_its_trigger_chain() -> Result<()> {
444 let item_rows = vec![ItemRow {
445 id: 220_202,
446 name: "Item 220202".to_owned(),
447 effects: vec![ItemEffect {
448 spell_id: 444_958,
449 trigger_type: 0,
450 charges: 0,
451 cooldown: 0,
452 category_cooldown: 0,
453 }],
454 }];
455 let related_rows = vec![spell_row(444_958, &[451_199]), spell_row(451_199, &[])];
456
457 let output = assemble_output(&[], &item_rows, &related_rows);
458 let trigger_spells = &output.items[0].effect_spells[0].trigger_spells;
459
460 verify_that!(trigger_spells.len(), eq(1))?;
461 verify_that!(trigger_spells[0].id, eq(451_199))?;
462
463 verify_that!(&trigger_spells[0].trigger_spells, is_empty())
464 }
465
466 #[gtest]
467 fn cyclic_trigger_chain_terminates_with_leaf() -> Result<()> {
468 let spell_rows = vec![spell_row(1, &[2])];
469 let related_rows = vec![spell_row(2, &[1])];
470
471 let output = assemble_output(&spell_rows, &[], &related_rows);
472 let cycle_leaf = &output.spells[0].trigger_spells[0].trigger_spells;
473
474 verify_that!(cycle_leaf.len(), eq(1))?;
475 verify_that!(cycle_leaf[0].id, eq(1))?;
476
477 verify_that!(&cycle_leaf[0].trigger_spells, is_empty())
478 }
479}