1use console::Style;
4use std::collections::HashMap;
5
6#[derive(Debug, Clone)]
8pub struct TreeNode {
9 pub name: String,
10 pub node_type: NodeType,
11 pub size: Option<u64>,
12 pub children: Vec<TreeNode>,
13 pub metadata: HashMap<String, String>,
14 pub external_refs: Vec<ExternalRef>,
15}
16
17#[derive(Debug, Clone, PartialEq)]
19pub enum NodeType {
20 Root,
21 Header,
22 #[allow(dead_code)]
23 Chunk,
24 #[allow(dead_code)]
25 Table,
26 File,
27 Directory,
28 #[allow(dead_code)] Reference,
30 #[allow(dead_code)]
31 Property,
32 #[allow(dead_code)]
33 Data,
34}
35
36#[derive(Debug, Clone)]
38pub struct ExternalRef {
39 pub path: String,
40 pub ref_type: RefType,
41 pub exists: Option<bool>,
42}
43
44#[derive(Debug, Clone, PartialEq)]
46pub enum RefType {
47 Texture,
48 Model,
49 Animation,
50 Map,
51 Database,
52 Sound,
53 Script,
54 Archive,
55 Unknown,
56}
57
58#[derive(Debug, Clone)]
60pub struct TreeOptions {
61 pub max_depth: Option<usize>,
62 pub show_external_refs: bool,
63 pub no_color: bool,
64 pub show_metadata: bool,
65 pub compact: bool,
66 #[cfg_attr(not(feature = "wmo"), allow(dead_code))]
67 pub verbose: bool,
68}
69
70impl Default for TreeOptions {
71 fn default() -> Self {
72 Self {
73 max_depth: None,
74 show_external_refs: true,
75 no_color: false,
76 show_metadata: true,
77 compact: false,
78 verbose: false,
79 }
80 }
81}
82
83impl TreeNode {
84 pub fn new(name: String, node_type: NodeType) -> Self {
86 Self {
87 name,
88 node_type,
89 size: None,
90 children: Vec::new(),
91 metadata: HashMap::new(),
92 external_refs: Vec::new(),
93 }
94 }
95
96 pub fn add_child(mut self, child: TreeNode) -> Self {
98 self.children.push(child);
99 self
100 }
101
102 pub fn with_size(mut self, size: u64) -> Self {
104 self.size = Some(size);
105 self
106 }
107
108 pub fn with_metadata(mut self, key: &str, value: &str) -> Self {
110 self.metadata.insert(key.to_string(), value.to_string());
111 self
112 }
113
114 pub fn with_external_ref(mut self, path: &str, ref_type: RefType) -> Self {
116 self.external_refs.push(ExternalRef {
117 path: path.to_string(),
118 ref_type,
119 exists: None,
120 });
121 self
122 }
123}
124
125impl ExternalRef {
126 pub fn icon(&self) -> &'static str {
128 match self.ref_type {
129 RefType::Texture => "πΌοΈ",
130 RefType::Model => "ποΈ",
131 RefType::Animation => "π½οΈ",
132 RefType::Map => "πΊοΈ",
133 RefType::Database => "π",
134 RefType::Sound => "π",
135 RefType::Script => "π",
136 RefType::Archive => "π¦",
137 RefType::Unknown => "π",
138 }
139 }
140
141 pub fn style(&self, no_color: bool) -> Style {
143 if no_color {
144 Style::new()
145 } else {
146 match self.exists {
147 Some(true) => Style::new().green(),
148 Some(false) => Style::new().red(),
149 None => Style::new().yellow(),
150 }
151 }
152 }
153}
154
155impl NodeType {
156 pub fn icon(&self) -> &'static str {
158 match self {
159 NodeType::Root => "π",
160 NodeType::Header => "π",
161 NodeType::Chunk => "π¦",
162 NodeType::Table => "π",
163 NodeType::File => "π",
164 NodeType::Directory => "π",
165 NodeType::Reference => "π",
166 NodeType::Property => "π·οΈ",
167 NodeType::Data => "πΎ",
168 }
169 }
170
171 pub fn style(&self, no_color: bool) -> Style {
173 if no_color {
174 Style::new()
175 } else {
176 match self {
177 NodeType::Root => Style::new().bold().cyan(),
178 NodeType::Header => Style::new().bold().yellow(),
179 NodeType::Chunk => Style::new().blue(),
180 NodeType::Table => Style::new().magenta(),
181 NodeType::File => Style::new().green(),
182 NodeType::Directory => Style::new().cyan(),
183 NodeType::Reference => Style::new().yellow(),
184 NodeType::Property => Style::new().dim(),
185 NodeType::Data => Style::new().white(),
186 }
187 }
188 }
189}
190
191pub fn render_tree(root: &TreeNode, options: &TreeOptions) -> String {
193 let mut output = String::new();
194 render_node(root, &mut output, "", true, 0, options);
195 output
196}
197
198fn render_node(
200 node: &TreeNode,
201 output: &mut String,
202 prefix: &str,
203 is_last: bool,
204 depth: usize,
205 options: &TreeOptions,
206) {
207 if let Some(max_depth) = options.max_depth
209 && depth > max_depth
210 {
211 return;
212 }
213
214 let icon = node.node_type.icon();
216 let style = node.node_type.style(options.no_color);
217 let connector = if depth == 0 {
218 ""
219 } else if is_last {
220 "βββ "
221 } else {
222 "βββ "
223 };
224
225 let mut line = format!(
226 "{}{}{} {}",
227 prefix,
228 connector,
229 icon,
230 style.apply_to(&node.name)
231 );
232
233 if let Some(size) = node.size {
235 line.push_str(&format!(" ({})", format_bytes(size)));
236 }
237
238 if options.show_metadata && !node.metadata.is_empty() && options.compact {
240 let mut meta_parts = Vec::new();
242 for (key, value) in &node.metadata {
243 if ["version", "count", "flags", "type"].contains(&key.as_str()) {
244 meta_parts.push(format!("{key}:{value}"));
245 }
246 }
247 if !meta_parts.is_empty() {
248 line.push_str(&format!(" [{}]", meta_parts.join(", ")));
249 }
250 }
251
252 output.push_str(&line);
253 output.push('\n');
254
255 if options.show_metadata && !options.compact && !node.metadata.is_empty() {
257 let child_prefix = if depth == 0 {
258 ""
259 } else if is_last {
260 " "
261 } else {
262 "β "
263 };
264 let meta_prefix = format!("{prefix}{child_prefix} ");
265
266 for (key, value) in &node.metadata {
267 let meta_style = Style::new().dim();
268 output.push_str(&format!(
269 "{}π·οΈ {}: {}\n",
270 meta_prefix,
271 meta_style.apply_to(key),
272 value
273 ));
274 }
275 }
276
277 if options.show_external_refs && !node.external_refs.is_empty() {
279 let child_prefix = if depth == 0 {
280 ""
281 } else if is_last {
282 " "
283 } else {
284 "β "
285 };
286 let ref_prefix = format!("{prefix}{child_prefix} ");
287
288 for ext_ref in &node.external_refs {
289 let icon = ext_ref.icon();
290 let style = ext_ref.style(options.no_color);
291 output.push_str(&format!(
292 "{}βββ {} {}\n",
293 ref_prefix,
294 icon,
295 style.apply_to(&ext_ref.path)
296 ));
297 }
298 }
299
300 if !node.children.is_empty() {
302 let new_prefix = if depth == 0 {
303 String::new()
304 } else {
305 format!("{}{}", prefix, if is_last { " " } else { "β " })
306 };
307
308 for (i, child) in node.children.iter().enumerate() {
309 let is_last_child = i == node.children.len() - 1;
310 render_node(
311 child,
312 output,
313 &new_prefix,
314 is_last_child,
315 depth + 1,
316 options,
317 );
318 }
319 }
320}
321
322fn format_bytes(bytes: u64) -> String {
324 const UNITS: &[&str] = &["B", "KB", "MB", "GB"];
325 let mut size = bytes as f64;
326 let mut unit_index = 0;
327
328 while size >= 1024.0 && unit_index < UNITS.len() - 1 {
329 size /= 1024.0;
330 unit_index += 1;
331 }
332
333 if unit_index == 0 {
334 format!("{} {}", bytes, UNITS[unit_index])
335 } else {
336 format!("{:.1} {}", size, UNITS[unit_index])
337 }
338}
339
340pub fn detect_ref_type(path: &str) -> RefType {
342 let path_lower = path.to_lowercase();
343
344 if path_lower.ends_with(".blp") {
345 RefType::Texture
346 } else if path_lower.ends_with(".m2") || path_lower.ends_with(".mdx") {
347 RefType::Model
348 } else if path_lower.ends_with(".anim") || path_lower.ends_with(".bone") {
349 RefType::Animation
350 } else if path_lower.ends_with(".wdt")
351 || path_lower.ends_with(".adt")
352 || path_lower.ends_with(".wdl")
353 {
354 RefType::Map
355 } else if path_lower.ends_with(".dbc") || path_lower.ends_with(".db2") {
356 RefType::Database
357 } else if path_lower.ends_with(".wav") || path_lower.ends_with(".mp3") {
358 RefType::Sound
359 } else if path_lower.ends_with(".lua") || path_lower.ends_with(".xml") {
360 RefType::Script
361 } else if path_lower.ends_with(".mpq") {
362 RefType::Archive
363 } else {
364 RefType::Unknown
365 }
366}
367
368#[cfg(test)]
369mod tests {
370 use super::*;
371
372 #[test]
373 fn test_tree_rendering() {
374 let root = TreeNode::new("test.mpq".to_string(), NodeType::Root)
375 .with_size(1024)
376 .with_metadata("version", "v2")
377 .add_child(
378 TreeNode::new("Header".to_string(), NodeType::Header)
379 .with_size(32)
380 .with_metadata("format", "MPQ v2"),
381 )
382 .add_child(
383 TreeNode::new("Files".to_string(), NodeType::Directory).add_child(
384 TreeNode::new("texture.blp".to_string(), NodeType::File)
385 .with_size(2048)
386 .with_external_ref("Interface/Icons/texture.blp", RefType::Texture),
387 ),
388 );
389
390 let options = TreeOptions::default();
391 let output = render_tree(&root, &options);
392
393 assert!(output.contains("test.mpq"));
394 assert!(output.contains("Header"));
395 assert!(output.contains("Files"));
396 assert!(output.contains("texture.blp"));
397 }
398
399 #[test]
400 fn test_ref_type_detection() {
401 assert_eq!(detect_ref_type("texture.blp"), RefType::Texture);
402 assert_eq!(detect_ref_type("model.m2"), RefType::Model);
403 assert_eq!(detect_ref_type("data.dbc"), RefType::Database);
404 assert_eq!(detect_ref_type("archive.mpq"), RefType::Archive);
405 }
406}