//! Formatters for relationship visualization
//! "Making code relationships beautiful" - Trisha from Accounting
use crate::relations::{FileRelation, RelationAnalyzer, RelationType};
use anyhow::Result;
use std::collections::{HashMap, HashSet};
use std::io::Write;
use std::path::{Path, PathBuf};
/// Format relationships as Mermaid diagram
pub struct MermaidRelationFormatter;
impl MermaidRelationFormatter {
pub fn format<W: Write>(
&self,
writer: &mut W,
analyzer: &RelationAnalyzer,
root_path: &Path,
) -> Result<()> {
writeln!(writer, "```mermaid")?;
writeln!(writer, "graph TD")?;
writeln!(writer, " %% Smart Tree Relationship Map 🌟")?;
writeln!(writer, " %% Generated by st --relations --mode mermaid")?;
writeln!(writer)?;
// Collect unique files
let mut files = HashSet::new();
let mut file_ids = HashMap::new();
let mut id_counter = 0;
for rel in analyzer.get_relations() {
files.insert(&rel.source);
files.insert(&rel.target);
}
// Generate node definitions with styling
writeln!(writer, " %% File nodes")?;
for file in files {
let relative = file.strip_prefix(root_path).unwrap_or(file);
let display_name = relative.to_string_lossy();
let id = format!("F{}", id_counter);
file_ids.insert(file, id.clone());
id_counter += 1;
// Style based on file type
let style = if display_name.contains("test") {
"fill:#90EE90,stroke:#228B22,stroke-width:2px" // Light green for tests
} else if display_name.ends_with(".rs") {
"fill:#FFE4B5,stroke:#FF8C00,stroke-width:2px" // Moccasin for Rust
} else if display_name.ends_with(".py") {
"fill:#87CEEB,stroke:#4682B4,stroke-width:2px" // Sky blue for Python
} else {
"fill:#F0F0F0,stroke:#696969,stroke-width:1px"
};
writeln!(
writer,
" {}[\"{}\"]\n style {} {}",
id, display_name, id, style
)?;
}
writeln!(writer)?;
writeln!(writer, " %% Relationships")?;
// Generate relationships with labels
for rel in analyzer.get_relations() {
let source_id = &file_ids[&rel.source];
let target_id = &file_ids[&rel.target];
let (arrow, label) = match rel.relation_type {
RelationType::Imports => ("-->", "imports"),
RelationType::FunctionCall => ("-.->", "calls"),
RelationType::TypeUsage => ("-.->", "uses"),
RelationType::TestedBy => ("==>", "tested by"),
RelationType::Exports => ("<-->", "exports"),
RelationType::Coupled => ("<===>", "coupled!"),
};
if rel.items.is_empty() {
writeln!(
writer,
" {} {}|{}| {}",
source_id, arrow, label, target_id
)?;
} else {
let items = rel.items.join(", ");
writeln!(
writer,
" {} {}|{}: {}| {}",
source_id, arrow, label, items, target_id
)?;
}
}
writeln!(writer)?;
writeln!(writer, " %% Legend")?;
writeln!(writer, " subgraph Legend")?;
writeln!(writer, " L1[Imports] --> L2[Target]")?;
writeln!(writer, " L3[Caller] -.-> L4[Function]")?;
writeln!(writer, " L5[Source] ==> L6[Tests]")?;
writeln!(writer, " L7[Tightly] <==> L8[Coupled]")?;
writeln!(writer, " end")?;
writeln!(writer, "```")?;
Ok(())
}
}
/// Format relationships as DOT/GraphViz
pub struct DotRelationFormatter;
impl DotRelationFormatter {
pub fn format<W: Write>(
&self,
writer: &mut W,
analyzer: &RelationAnalyzer,
root_path: &Path,
) -> Result<()> {
writeln!(writer, "digraph CodeRelations {{")?;
writeln!(writer, " // Smart Tree Relationship Graph")?;
writeln!(writer, " rankdir=LR;")?;
writeln!(writer, " node [shape=box, style=filled];")?;
writeln!(writer)?;
// Collect unique files
let mut files = HashSet::new();
for rel in analyzer.get_relations() {
files.insert(&rel.source);
files.insert(&rel.target);
}
// Node definitions
writeln!(writer, " // Nodes")?;
for file in &files {
let relative = file.strip_prefix(root_path).unwrap_or(file);
let display_name = relative.to_string_lossy();
let color = if display_name.contains("test") {
"lightgreen"
} else if display_name.ends_with(".rs") {
"lightyellow"
} else if display_name.ends_with(".py") {
"lightblue"
} else {
"lightgray"
};
writeln!(
writer,
" \"{}\" [fillcolor=\"{}\"];",
display_name, color
)?;
}
writeln!(writer)?;
writeln!(writer, " // Edges")?;
for rel in analyzer.get_relations() {
let source = rel.source.strip_prefix(root_path).unwrap_or(&rel.source);
let target = rel.target.strip_prefix(root_path).unwrap_or(&rel.target);
let style = match rel.relation_type {
RelationType::Imports => "solid",
RelationType::FunctionCall => "dashed",
RelationType::TypeUsage => "dotted",
RelationType::TestedBy => "bold",
RelationType::Exports => "solid",
RelationType::Coupled => "bold",
};
let color = match rel.relation_type {
RelationType::TestedBy => "green",
RelationType::Coupled => "red",
_ => "black",
};
writeln!(
writer,
" \"{}\" -> \"{}\" [style={}, color={}, label=\"{:?}\"];",
source.to_string_lossy(),
target.to_string_lossy(),
style,
color,
rel.relation_type
)?;
}
writeln!(writer, "}}")?;
Ok(())
}
}
/// Format relationships in compressed AI-friendly format
pub struct CompressedRelationFormatter;
impl CompressedRelationFormatter {
pub fn format<W: Write>(
&self,
writer: &mut W,
analyzer: &RelationAnalyzer,
root_path: &Path,
) -> Result<()> {
writeln!(writer, "RELATIONS_V1:")?;
// Create file index
let mut files = HashSet::new();
for rel in analyzer.get_relations() {
files.insert(&rel.source);
files.insert(&rel.target);
}
let mut file_index: HashMap<&PathBuf, usize> = HashMap::new();
writeln!(writer, "FILES:")?;
for (idx, file) in files.iter().enumerate() {
let relative = file.strip_prefix(root_path).unwrap_or(file);
writeln!(writer, "{:x}:{}", idx, relative.to_string_lossy())?;
file_index.insert(file, idx);
}
writeln!(writer, "RELS:")?;
// Format: source_idx,target_idx,type,strength[:items]
for rel in analyzer.get_relations() {
let source_idx = file_index[&rel.source];
let target_idx = file_index[&rel.target];
let type_code = match rel.relation_type {
RelationType::Imports => 'I',
RelationType::FunctionCall => 'F',
RelationType::TypeUsage => 'T',
RelationType::TestedBy => 'X',
RelationType::Exports => 'E',
RelationType::Coupled => 'C',
};
if rel.items.is_empty() {
writeln!(
writer,
"{:x},{:x},{},{}",
source_idx, target_idx, type_code, rel.strength
)?;
} else {
writeln!(
writer,
"{:x},{:x},{},{}:{}",
source_idx,
target_idx,
type_code,
rel.strength,
rel.items.join(",")
)?;
}
}
writeln!(writer, "END_RELATIONS")?;
Ok(())
}
}
/// Format relationships as a text summary
pub struct TextRelationFormatter;
impl TextRelationFormatter {
pub fn format<W: Write>(
&self,
writer: &mut W,
analyzer: &RelationAnalyzer,
root_path: &Path,
) -> Result<()> {
writeln!(writer, "🔗 Code Relationship Analysis")?;
writeln!(writer, "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")?;
writeln!(writer)?;
// Group by file
let mut file_relations: HashMap<&PathBuf, Vec<&FileRelation>> = HashMap::new();
for rel in analyzer.get_relations() {
file_relations.entry(&rel.source).or_default().push(rel);
}
// Calculate total files before moving file_relations
let total_files = file_relations.len();
for (file, relations) in file_relations {
let relative = file.strip_prefix(root_path).unwrap_or(file);
writeln!(writer, "📄 {}", relative.to_string_lossy())?;
// Group by type
for rel_type in &[
RelationType::Imports,
RelationType::FunctionCall,
RelationType::TypeUsage,
RelationType::TestedBy,
RelationType::Exports,
RelationType::Coupled,
] {
let typed_rels: Vec<_> = relations
.iter()
.filter(|r| &r.relation_type == rel_type)
.collect();
if !typed_rels.is_empty() {
let emoji = match rel_type {
RelationType::Imports => "├─→ imports:",
RelationType::FunctionCall => "├─→ calls:",
RelationType::TypeUsage => "├─→ uses types from:",
RelationType::TestedBy => "├─→ tested by:",
RelationType::Exports => "├─→ exports to:",
RelationType::Coupled => "├─⚠️ tightly coupled with:",
};
writeln!(writer, " {}", emoji)?;
for rel in typed_rels {
let target = rel.target.strip_prefix(root_path).unwrap_or(&rel.target);
if rel.items.is_empty() {
writeln!(writer, " • {}", target.to_string_lossy())?;
} else {
writeln!(
writer,
" • {} ({})",
target.to_string_lossy(),
rel.items.join(", ")
)?;
}
}
}
}
writeln!(writer)?;
}
// Summary statistics
writeln!(writer, "📊 Summary")?;
writeln!(writer, "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")?;
let total_relations = analyzer.get_relations().len();
let coupled_count = analyzer
.get_relations()
.iter()
.filter(|r| r.relation_type == RelationType::Coupled)
.count();
let tested_count = analyzer
.get_relations()
.iter()
.filter(|r| r.relation_type == RelationType::TestedBy)
.count();
writeln!(writer, "Total files analyzed: {}", total_files)?;
writeln!(writer, "Total relationships: {}", total_relations)?;
writeln!(writer, "Tightly coupled pairs: {}", coupled_count)?;
writeln!(writer, "Files with tests: {}", tested_count)?;
Ok(())
}
}