st/formatters/

relations.rs

//! 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(())
    }
}