Skip to main content

warcraft_rs/utils/
tree.rs

1//! Tree structure rendering utilities for file format visualization
2
3use console::Style;
4use std::collections::HashMap;
5
6/// Represents a node in a tree structure
7#[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/// Types of nodes in the tree
18#[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)] // Reserved for future use
29    Reference,
30    #[allow(dead_code)]
31    Property,
32    #[allow(dead_code)]
33    Data,
34}
35
36/// External file reference
37#[derive(Debug, Clone)]
38pub struct ExternalRef {
39    pub path: String,
40    pub ref_type: RefType,
41    pub exists: Option<bool>,
42}
43
44/// Types of external references
45#[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/// Options for tree rendering
59#[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    /// Create a new tree node
85    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    /// Add a child node
97    pub fn add_child(mut self, child: TreeNode) -> Self {
98        self.children.push(child);
99        self
100    }
101
102    /// Set the size of this node
103    pub fn with_size(mut self, size: u64) -> Self {
104        self.size = Some(size);
105        self
106    }
107
108    /// Add metadata
109    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    /// Add external reference
115    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    /// Get emoji icon for reference type
127    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    /// Get color style based on existence
142    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    /// Get emoji icon for node type
157    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    /// Get color style for node type
172    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
191/// Render a tree structure to string
192pub 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
198/// Render a single node and its children
199fn render_node(
200    node: &TreeNode,
201    output: &mut String,
202    prefix: &str,
203    is_last: bool,
204    depth: usize,
205    options: &TreeOptions,
206) {
207    // Check depth limit
208    if let Some(max_depth) = options.max_depth
209        && depth > max_depth
210    {
211        return;
212    }
213
214    // Node icon and name
215    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    // Add size if available
234    if let Some(size) = node.size {
235        line.push_str(&format!(" ({})", format_bytes(size)));
236    }
237
238    // Add metadata if enabled and available
239    if options.show_metadata && !node.metadata.is_empty() && options.compact {
240        // Show key metadata inline
241        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    // Show detailed metadata on separate lines (non-compact mode)
256    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    // Show external references
278    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    // Render children
301    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
322/// Format bytes in human-readable format
323fn 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
340/// Detect reference type from file extension
341pub 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}