1use crate::trace::{CallNode, CallTree, TraceDirection};
2use crate::tree::{NodeType, ReferenceTree, TreeNode};
3use crate::{CodeReference, SearchResult};
4use colored::*;
5use regex::RegexBuilder;
6
7pub struct TreeFormatter {
9 max_width: usize,
10 search_query: String,
11 simple_format: bool,
12}
13
14impl TreeFormatter {
15 pub fn new() -> Self {
17 Self {
18 max_width: 80,
19 search_query: String::new(),
20 simple_format: false,
21 }
22 }
23
24 pub fn with_width(max_width: usize) -> Self {
26 Self {
27 max_width,
28 search_query: String::new(),
29 simple_format: false,
30 }
31 }
32
33 pub fn with_search_query(mut self, query: String) -> Self {
35 self.search_query = query;
36 self
37 }
38
39 pub fn with_simple_format(mut self, simple: bool) -> Self {
41 self.simple_format = simple;
42 self
43 }
44
45 pub fn format_result(&self, result: &SearchResult) -> String {
47 if self.simple_format {
48 return self.format_result_simple(result);
49 }
50
51 let mut output = String::new();
52
53 if !result.translation_entries.is_empty() {
55 output.push_str(&format!("{}\n", "=== Translation Files ===".bold()));
56 for entry in &result.translation_entries {
57 output.push_str(&format!(
58 "{}:{}:{}: {}\n",
59 entry.file.display(),
60 entry.line,
61 entry.key.yellow().bold(),
62 format!("\"{}\"", entry.value).green().bold()
63 ));
64 }
65 output.push('\n');
66 }
67
68 if !result.code_references.is_empty() {
70 output.push_str(&format!("{}\n", "=== Code References ===".bold()));
71
72 let grouped_refs = self.group_code_references_by_file(&result.code_references);
74
75 for (file_path, refs) in &grouped_refs {
76 let mut sorted_refs = refs.clone();
78 sorted_refs.sort_by_key(|r| r.line);
79
80 let formatted_output =
82 self.format_code_references_with_context(file_path, &sorted_refs);
83 output.push_str(&formatted_output);
84 }
85 }
86
87 output
88 }
89
90 fn format_result_simple(&self, result: &SearchResult) -> String {
92 let mut output = String::new();
93
94 for entry in &result.translation_entries {
96 let escaped_key = self.escape_simple_content(&entry.key);
97 let escaped_value = self.escape_simple_content(&entry.value);
98 let content = format!("{}: {}", escaped_key, escaped_value);
99 output.push_str(&format!(
100 "{}:{}:{}\n",
101 self.escape_simple_path(&entry.file.display().to_string()),
102 entry.line,
103 content
104 ));
105 }
106
107 for code_ref in &result.code_references {
109 let escaped_content = self.escape_simple_content(code_ref.context.trim());
110 output.push_str(&format!(
111 "{}:{}:{}\n",
112 self.escape_simple_path(&code_ref.file.display().to_string()),
113 code_ref.line,
114 escaped_content
115 ));
116 }
117
118 output
119 }
120
121 fn escape_simple_path(&self, path: &str) -> String {
123 path.replace(':', "\\:")
126 }
127
128 fn escape_simple_content(&self, content: &str) -> String {
130 let clean_content = self.strip_ansi_codes(content);
132
133 let single_line = clean_content.replace(['\n', '\r'], " ");
135
136 let trimmed = single_line.trim();
138
139 let normalized = regex::Regex::new(r"\s+").unwrap().replace_all(trimmed, " ");
141
142 normalized.to_string()
143 }
144
145 fn strip_ansi_codes(&self, text: &str) -> String {
147 let complete_ansi = regex::Regex::new(r"\x1b\[[0-9;]*[a-zA-Z]").unwrap();
150 let mut result = complete_ansi.replace_all(text, "").to_string();
151
152 let incomplete_ansi = regex::Regex::new(r"\x1b\[[0-9;]*$").unwrap();
154 result = incomplete_ansi.replace_all(&result, "").to_string();
155
156 let remaining_ansi = regex::Regex::new(r"\x1b\[").unwrap();
158 result = remaining_ansi.replace_all(&result, "").to_string();
159
160 result
161 }
162
163 fn group_code_references_by_file(
165 &self,
166 code_refs: &[CodeReference],
167 ) -> std::collections::HashMap<std::path::PathBuf, Vec<CodeReference>> {
168 use std::collections::HashMap;
169
170 let mut grouped: HashMap<std::path::PathBuf, Vec<CodeReference>> = HashMap::new();
171 for code_ref in code_refs {
172 grouped
173 .entry(code_ref.file.clone())
174 .or_default()
175 .push(code_ref.clone());
176 }
177 grouped
178 }
179
180 fn format_code_references_with_context(
182 &self,
183 file_path: &std::path::Path,
184 refs: &[CodeReference],
185 ) -> String {
186 let mut output = String::new();
187
188 if refs.is_empty() {
189 return output;
190 }
191
192 let mut all_lines: Vec<(usize, String, bool)> = Vec::new(); for code_ref in refs {
196 for (i, context_line) in code_ref.context_before.iter().enumerate() {
198 let line_num = code_ref.line - code_ref.context_before.len() + i;
199 all_lines.push((line_num, context_line.clone(), false));
200 }
201
202 let highlighted_context =
204 self.highlight_key_in_context(&code_ref.context, &code_ref.key_path);
205 all_lines.push((code_ref.line, highlighted_context, true));
206
207 for (i, context_line) in code_ref.context_after.iter().enumerate() {
209 let line_num = code_ref.line + 1 + i;
210 all_lines.push((line_num, context_line.clone(), false));
211 }
212 }
213
214 all_lines.sort_by_key(|(line_num, _, _)| *line_num);
216 all_lines.dedup_by_key(|(line_num, _, _)| *line_num);
217
218 for (line_num, content, is_match) in all_lines {
220 let separator = if is_match { ":" } else { "-" };
221 output.push_str(&format!(
222 "{}{}{}:{}\n",
223 file_path.display(),
224 separator,
225 line_num,
226 content
227 ));
228 }
229
230 output.push('\n'); output
232 }
233
234 fn highlight_key_in_context(&self, context: &str, key: &str) -> String {
236 let escaped_key = regex::escape(key);
238
239 let re = match RegexBuilder::new(&escaped_key)
241 .case_insensitive(true)
242 .build()
243 {
244 Ok(r) => r,
245 Err(_) => return context.to_string(), };
247
248 let result = re.replace_all(context, |caps: ®ex::Captures| caps[0].bold().to_string());
250
251 result.to_string()
252 }
253
254 pub fn format(&self, tree: &ReferenceTree) -> String {
256 if self.simple_format {
257 return self.format_tree_simple(tree);
258 }
259
260 let mut output = String::new();
261 self.format_node(&tree.root, &mut output, "", true, true);
262 output
263 }
264
265 fn format_tree_simple(&self, tree: &ReferenceTree) -> String {
267 let mut output = String::new();
268 self.collect_simple_entries(&tree.root, &mut output);
269 output
270 }
271
272 fn collect_simple_entries(&self, node: &TreeNode, output: &mut String) {
274 if let Some(location) = &node.location {
276 let content = match node.node_type {
277 NodeType::Translation => {
278 let key = &node.content;
279 let value = node.metadata.as_deref().unwrap_or("");
280 format!("{}: {}", key, value)
281 }
282 NodeType::CodeRef => node.content.trim().to_string(),
283 _ => node.content.clone(),
284 };
285
286 let escaped_content = self.escape_simple_content(&content);
287 output.push_str(&format!(
288 "{}:{}:{}\n",
289 self.escape_simple_path(&location.file.display().to_string()),
290 location.line,
291 escaped_content
292 ));
293 }
294
295 for child in &node.children {
297 self.collect_simple_entries(child, output);
298 }
299 }
300
301 pub fn format_trace_tree(&self, tree: &CallTree, direction: TraceDirection) -> String {
302 match direction {
303 TraceDirection::Forward => self.format_forward_tree(tree),
304 TraceDirection::Backward => self.format_backward_tree(tree),
305 }
306 }
307
308 fn format_forward_tree(&self, tree: &CallTree) -> String {
309 let mut output = String::new();
310 Self::format_call_node(&tree.root, &mut output, "", true, true);
311 output
312 }
313
314 fn format_backward_tree(&self, tree: &CallTree) -> String {
315 let mut output = String::new();
316 let mut paths = Vec::new();
327 Self::collect_backward_paths(&tree.root, vec![], &mut paths);
328
329 for path in paths {
330 let mut display_path = path.clone();
342 display_path.reverse();
343
344 let mut chain = display_path
345 .iter()
346 .map(|node| {
347 format!(
348 "{} ({}:{})",
349 node.def.name.bold(),
350 node.def.file.display(),
351 node.def.line
352 )
353 })
354 .collect::<Vec<_>>()
355 .join(" -> ");
356
357 if let Some(first) = display_path.first() {
359 if first.truncated {
360 chain = format!("{} -> {}", "[depth limit reached]".red(), chain);
361 }
362 }
363
364 output.push_str(&chain);
365 output.push('\n');
366 }
367
368 if output.is_empty() {
369 output.push_str(&format!(
371 "{} (No incoming calls found)\n",
372 tree.root.def.name
373 ));
374 }
375
376 output
377 }
378
379 fn collect_backward_paths<'a>(
380 node: &'a CallNode,
381 mut current_path: Vec<&'a CallNode>,
382 paths: &mut Vec<Vec<&'a CallNode>>,
383 ) {
384 current_path.push(node);
385
386 if node.children.is_empty() {
387 if node.truncated {
391 }
395 paths.push(current_path);
396 } else {
397 for child in &node.children {
398 Self::collect_backward_paths(child, current_path.clone(), paths);
399 }
400 }
401 }
402
403 fn format_call_node(
404 node: &CallNode,
405 output: &mut String,
406 prefix: &str,
407 is_last: bool,
408 is_root: bool,
409 ) {
410 if !is_root {
411 output.push_str(prefix);
412 output.push_str(if is_last { "└─> " } else { "├─> " });
413 }
414
415 let content = format!(
416 "{} ({}:{})",
417 node.def.name.bold(),
418 node.def.file.display(),
419 node.def.line
420 );
421 output.push_str(&content);
422
423 if node.truncated {
424 output.push_str(&" [depth limit reached]".red().to_string());
425 }
426
427 output.push('\n');
428
429 let child_count = node.children.len();
430 for (i, child) in node.children.iter().enumerate() {
431 let is_last_child = i == child_count - 1;
432 let child_prefix = if is_root {
433 String::new()
434 } else {
435 format!("{}{} ", prefix, if is_last { " " } else { "│" })
436 };
437 Self::format_call_node(child, output, &child_prefix, is_last_child, false);
438 }
439 }
440
441 fn format_node(
443 &self,
444 node: &TreeNode,
445 output: &mut String,
446 prefix: &str,
447 is_last: bool,
448 is_root: bool,
449 ) {
450 if !is_root {
452 output.push_str(prefix);
453 output.push_str(if is_last { "└─> " } else { "├─> " });
454 }
455
456 let content = self.format_content(node);
458 output.push_str(&content);
459
460 if let Some(location) = &node.location {
462 let location_str = format!(" ({}:{})", location.file.display(), location.line);
463 output.push_str(&location_str);
464 }
465
466 output.push('\n');
467
468 let child_count = node.children.len();
470 for (i, child) in node.children.iter().enumerate() {
471 let is_last_child = i == child_count - 1;
472 let child_prefix = if is_root {
473 String::new()
474 } else {
475 format!("{}{} ", prefix, if is_last { " " } else { "│" })
476 };
477
478 self.format_node(child, output, &child_prefix, is_last_child, false);
479 }
480 }
481
482 fn format_content(&self, node: &TreeNode) -> String {
484 match node.node_type {
485 NodeType::Root => {
486 format!("'{}' (search query)", node.content)
487 }
488 NodeType::Translation => {
489 let key = &node.content;
490 let value = node.metadata.as_deref().unwrap_or("");
491
492 let available_width = self.max_width.saturating_sub(key.len()).saturating_sub(10);
494 let width = if available_width < 10 {
495 10
496 } else {
497 available_width
498 };
499 let truncated_value = self.truncate(value, width);
500
501 let highlighted_value = if !self.search_query.is_empty() {
503 self.highlight_key_in_context(&truncated_value, &self.search_query)
504 } else {
505 truncated_value
506 };
507
508 format!("{}: '{}'", key.yellow().bold(), highlighted_value)
509 }
510 NodeType::KeyPath => {
511 format!("Key: {}", node.content)
512 }
513 NodeType::CodeRef => {
514 let available_width = self.max_width.saturating_sub(10); let width = if available_width < 100 {
518 200 } else {
520 available_width.max(200) };
522 let truncated = self.truncate(node.content.trim(), width);
523
524 if let Some(key) = &node.metadata {
526 self.highlight_key_in_context(&truncated, key)
527 } else {
528 truncated
529 }
530 }
531 }
532 }
533
534 fn truncate(&self, s: &str, max_len: usize) -> String {
536 if s.chars().count() <= max_len {
537 s.to_string()
538 } else {
539 let truncated: String = s.chars().take(max_len.saturating_sub(3)).collect();
540 format!("{}...", truncated)
541 }
542 }
543}
544
545impl Default for TreeFormatter {
546 fn default() -> Self {
547 Self::new()
548 }
549}
550
551#[cfg(test)]
552mod tests {
553 use super::*;
554 use crate::tree::{Location, TreeNode};
555 use std::path::PathBuf;
556
557 #[test]
558 fn test_formatter_creation() {
559 let formatter = TreeFormatter::new();
560 assert_eq!(formatter.max_width, 80);
561 }
562
563 #[test]
564 fn test_formatter_with_custom_width() {
565 let formatter = TreeFormatter::with_width(120);
566 assert_eq!(formatter.max_width, 120);
567 }
568
569 #[test]
570 fn test_format_empty_tree() {
571 let tree = ReferenceTree::with_search_text("test".to_string());
572 let formatter = TreeFormatter::new();
573 let output = formatter.format(&tree);
574
575 assert!(output.contains("'test'"));
576 assert!(output.contains("search query"));
577 }
578
579 #[test]
580 fn test_format_tree_with_translation() {
581 let mut root = TreeNode::new(NodeType::Root, "add new".to_string());
582 let mut translation = TreeNode::with_location(
583 NodeType::Translation,
584 "invoice.labels.add_new".to_string(),
585 Location::new(PathBuf::from("en.yml"), 4),
586 );
587 translation.metadata = Some("add new".to_string());
588 root.add_child(translation);
589
590 let tree = ReferenceTree::new(root);
591 let formatter = TreeFormatter::new();
592 let output = formatter.format(&tree);
593
594 assert!(output.contains("'add new'"));
595 assert!(output.contains("invoice.labels.add_new"));
596 assert!(output.contains("en.yml:4"));
597 assert!(output.contains("└─>") || output.contains("├─>"));
598 }
599
600 #[test]
601 fn test_format_complete_tree() {
602 let mut root = TreeNode::new(NodeType::Root, "add new".to_string());
603
604 let mut translation = TreeNode::with_location(
605 NodeType::Translation,
606 "invoice.labels.add_new".to_string(),
607 Location::new(PathBuf::from("en.yml"), 4),
608 );
609 translation.metadata = Some("add new".to_string());
610
611 let mut key_path = TreeNode::new(NodeType::KeyPath, "invoice.labels.add_new".to_string());
612
613 let code_ref = TreeNode::with_location(
614 NodeType::CodeRef,
615 "I18n.t('invoice.labels.add_new')".to_string(),
616 Location::new(PathBuf::from("invoices.ts"), 14),
617 );
618
619 key_path.add_child(code_ref);
620 translation.add_child(key_path);
621 root.add_child(translation);
622
623 let tree = ReferenceTree::new(root);
624 let formatter = TreeFormatter::new();
625 let output = formatter.format(&tree);
626
627 assert!(output.contains("'add new'"));
629 assert!(output.contains("invoice.labels.add_new"));
630 assert!(output.contains("Key:"));
631 assert!(output.contains("I18n.t"));
632 assert!(output.contains("en.yml:4"));
633 assert!(output.contains("invoices.ts:14"));
634 }
635
636 #[test]
637 fn test_format_multiple_children() {
638 let mut root = TreeNode::new(NodeType::Root, "test".to_string());
639
640 let mut child1 = TreeNode::with_location(
641 NodeType::Translation,
642 "key1".to_string(),
643 Location::new(PathBuf::from("file1.yml"), 1),
644 );
645 child1.metadata = Some("value1".to_string());
646
647 let mut child2 = TreeNode::with_location(
648 NodeType::Translation,
649 "key2".to_string(),
650 Location::new(PathBuf::from("file2.yml"), 2),
651 );
652 child2.metadata = Some("value2".to_string());
653
654 root.add_child(child1);
655 root.add_child(child2);
656
657 let tree = ReferenceTree::new(root);
658 let formatter = TreeFormatter::new();
659 let output = formatter.format(&tree);
660
661 assert!(output.contains("key1"));
663 assert!(output.contains("key2"));
664 assert!(output.contains("file1.yml:1"));
665 assert!(output.contains("file2.yml:2"));
666
667 assert!(output.contains("├─>"));
669 assert!(output.contains("└─>"));
670 }
671
672 #[test]
673 fn test_truncate_long_content() {
674 let formatter = TreeFormatter::with_width(50);
675 let long_string = "a".repeat(100);
676 let truncated = formatter.truncate(&long_string, 20);
677
678 assert!(truncated.len() <= 20);
679 assert!(truncated.ends_with("..."));
680 }
681
682 #[test]
683 fn test_truncate_short_content() {
684 let formatter = TreeFormatter::new();
685 let short_string = "short";
686 let result = formatter.truncate(short_string, 20);
687
688 assert_eq!(result, "short");
689 }
690
691 #[test]
692 fn test_format_content_root() {
693 let formatter = TreeFormatter::new();
694 let node = TreeNode::new(NodeType::Root, "test query".to_string());
695 let content = formatter.format_content(&node);
696
697 assert!(content.contains("test query"));
698 assert!(content.contains("search query"));
699 }
700
701 #[test]
702 fn test_format_content_key_path() {
703 let formatter = TreeFormatter::new();
704 let node = TreeNode::new(NodeType::KeyPath, "invoice.labels.add_new".to_string());
705 let content = formatter.format_content(&node);
706
707 assert!(content.contains("Key:"));
708 assert!(content.contains("invoice.labels.add_new"));
709 }
710
711 #[test]
712 fn test_format_content_code_ref() {
713 let formatter = TreeFormatter::new();
714 let node = TreeNode::new(
715 NodeType::CodeRef,
716 " I18n.t('invoice.labels.add_new') ".to_string(),
717 );
718 let content = formatter.format_content(&node);
719
720 assert!(content.contains("I18n.t"));
721 assert!(!content.starts_with(" "));
723 }
724
725 #[test]
726 fn test_format_deep_nesting() {
727 let mut root = TreeNode::new(NodeType::Root, "test".to_string());
728 let mut level1 = TreeNode::new(NodeType::Translation, "level1".to_string());
729 let mut level2 = TreeNode::new(NodeType::KeyPath, "level2".to_string());
730 let level3 = TreeNode::new(NodeType::CodeRef, "level3".to_string());
731
732 level2.add_child(level3);
733 level1.add_child(level2);
734 root.add_child(level1);
735
736 let tree = ReferenceTree::new(root);
737 let formatter = TreeFormatter::new();
738 let output = formatter.format(&tree);
739
740 let lines: Vec<&str> = output.lines().collect();
742 assert!(lines.len() >= 4);
743
744 assert!(lines[2].starts_with(' ') || lines[2].starts_with('│'));
746 }
747
748 #[test]
749 fn test_highlight_case_insensitive_lowercase() {
750 colored::control::set_override(true); let formatter = TreeFormatter::new();
752 let context = "const value = pmfc.getData();";
753 let key = "PMFC";
754 let result = formatter.highlight_key_in_context(context, key);
755
756 assert!(result.contains("pmfc"));
758 assert_ne!(result, context);
761 }
762
763 #[test]
764 fn test_highlight_case_insensitive_uppercase() {
765 colored::control::set_override(true); let formatter = TreeFormatter::new();
767 let context = "const value = PMFC.getData();";
768 let key = "pmfc";
769 let result = formatter.highlight_key_in_context(context, key);
770
771 assert!(result.contains("PMFC"));
773 assert_ne!(result, context);
774 }
775
776 #[test]
777 fn test_highlight_case_insensitive_mixed() {
778 colored::control::set_override(true); let formatter = TreeFormatter::new();
780 let context = "const a = PmFc.get(); const b = pmfc.set();";
781 let key = "PMFC";
782 let result = formatter.highlight_key_in_context(context, key);
783
784 assert!(result.contains("PmFc"));
786 assert!(result.contains("pmfc"));
787 assert_ne!(result, context);
788 }
789
790 #[test]
791 fn test_highlight_with_special_regex_chars() {
792 colored::control::set_override(true); let formatter = TreeFormatter::new();
794 let context = "price: $19.99";
795 let key = "$19.99";
796 let result = formatter.highlight_key_in_context(context, key);
797
798 assert!(result.contains("$19.99"));
800 assert_ne!(result, context);
801 }
802
803 #[test]
804 fn test_highlight_exact_match_still_works() {
805 colored::control::set_override(true); let formatter = TreeFormatter::new();
807 let context = "I18n.t('invoice.labels.add_new')";
808 let key = "invoice.labels.add_new";
809 let result = formatter.highlight_key_in_context(context, key);
810
811 assert!(result.contains("invoice.labels.add_new"));
813 assert_ne!(result, context);
814 }
815}