archlens/crates/adapters/mermaid/src/mermaid_renderer.rs

use std::collections::{HashMap, HashSet};

use archlens_domain::{
    CodeElement, CodeGraph, DiagramLevel, DomainError, ModuleName, RelationshipKind, RenderOutput,
    RenderedFile, Visibility, ports::DiagramRenderer,
};

pub struct MermaidRenderer {
    level: DiagramLevel,
}

impl Default for MermaidRenderer {
    fn default() -> Self {
        Self::new()
    }
}

impl MermaidRenderer {
    pub fn new() -> Self {
        Self {
            level: DiagramLevel::Type,
        }
    }

    pub fn with_level(level: DiagramLevel) -> Self {
        Self { level }
    }

    fn format_element_name(element: &CodeElement) -> String {
        let name = element.name();
        if element.generics().is_empty() {
            name.to_string()
        } else {
            format!("{}~{}~", name, element.generics().join(", "))
        }
    }

    fn format_visibility(visibility: Visibility) -> &'static str {
        match visibility {
            Visibility::Public => "public",
            Visibility::Private => "private",
            Visibility::Internal => "internal",
        }
    }

    fn render_class_diagram(&self, graph: &CodeGraph) -> String {
        let mut lines = vec!["classDiagram".to_string()];

        let (grouped, ungrouped) = graph.elements_by_module();
        let has_namespaces = !grouped.is_empty();

        let mut seen: HashSet<String> = HashSet::new();
        let mut deferred_members: Vec<String> = Vec::new();

        for element in &ungrouped {
            if seen.insert(element.name().to_string()) {
                Self::push_class_lines(&mut lines, &mut deferred_members, element, "    ", false);
            }
        }

        if has_namespaces {
            for (namespace, elements) in &grouped {
                lines.push(format!("    namespace {namespace} {{"));
                let mut ns_seen: HashSet<String> = HashSet::new();
                for element in elements {
                    if ns_seen.insert(element.name().to_string()) {
                        Self::push_class_lines(
                            &mut lines,
                            &mut deferred_members,
                            element,
                            "        ",
                            true,
                        );
                    }
                }
                lines.push("    }".to_string());
            }
        }

        lines.extend(deferred_members);

        let mut rel_seen: HashSet<String> = HashSet::new();
        for rel in graph.relationships() {
            if rel.kind() == RelationshipKind::Import {
                continue;
            }
            let arrow = match rel.kind() {
                RelationshipKind::Inheritance => "<|--",
                RelationshipKind::Composition => "-->",
                RelationshipKind::Import => "..>",
            };
            let key = format!("{} {} {}", rel.source(), arrow, rel.target());
            if rel_seen.insert(key.clone()) {
                lines.push(format!("    {key}"));
            }
        }

        lines.join("\n")
    }

    fn push_class_lines(
        lines: &mut Vec<String>,
        deferred: &mut Vec<String>,
        element: &CodeElement,
        indent: &str,
        in_namespace: bool,
    ) {
        lines.push(format!(
            "{indent}class {}",
            Self::format_element_name(element)
        ));

        let member_target = if in_namespace { deferred } else { lines };
        if element.visibility() != Visibility::Public {
            member_target.push(format!(
                "    <<{}>> {}",
                Self::format_visibility(element.visibility()),
                element.name()
            ));
        }
        let name = element.name();
        for field in element.fields() {
            member_target.push(format!("    {name} : {field}"));
        }
        for method in element.methods() {
            member_target.push(format!("    {name} : {method}"));
        }
    }

    fn render_module_flowchart(&self, graph: &CodeGraph) -> String {
        let mut lines = vec!["graph TD".to_string()];

        let mut name_to_modules: HashMap<&str, HashSet<&str>> = HashMap::new();
        let mut file_to_module: HashMap<String, String> = HashMap::new();
        let mut modules: HashSet<String> = HashSet::new();

        for element in graph.elements() {
            if let Some(module) = element.module() {
                name_to_modules
                    .entry(element.name())
                    .or_default()
                    .insert(module.as_str());
                modules.insert(module.as_str().to_string());

                let file_stem = std::path::Path::new(element.file_path().as_str())
                    .file_stem()
                    .and_then(|s| s.to_str())
                    .unwrap_or("");
                if !file_stem.is_empty() {
                    file_to_module.insert(file_stem.to_string(), module.as_str().to_string());
                }
            }
        }

        for module in &modules {
            lines.push(format!("    {module}[{module}]"));
        }

        let mut module_edges: HashSet<(String, String)> = HashSet::new();
        for rel in graph.relationships() {
            match rel.kind() {
                RelationshipKind::Import => {
                    let source_mod = file_to_module.get(rel.source());
                    let target_top = rel.target().split('.').next().unwrap_or("");
                    let target_mod = ModuleName::capitalize(target_top);

                    if let Some(src) = source_mod
                        && modules.contains(&target_mod)
                        && *src != target_mod
                    {
                        module_edges.insert((src.clone(), target_mod));
                    }
                }
                _ => {
                    if modules.contains(rel.source())
                        && modules.contains(rel.target())
                        && rel.source() != rel.target()
                    {
                        module_edges.insert((rel.source().to_string(), rel.target().to_string()));
                        continue;
                    }

                    let src_mods = name_to_modules.get(rel.source());
                    let tgt_mods = name_to_modules.get(rel.target());

                    if let (Some(src_set), Some(tgt_set)) = (src_mods, tgt_mods) {
                        for src_mod in src_set {
                            if tgt_set.contains(src_mod) {
                                continue;
                            }
                            for tgt_mod in tgt_set {
                                if src_mod != tgt_mod {
                                    module_edges.insert((src_mod.to_string(), tgt_mod.to_string()));
                                }
                            }
                        }
                    }
                }
            }
        }

        for (source, target) in &module_edges {
            lines.push(format!("    {source} --> {target}"));
        }

        lines.join("\n")
    }

    fn render_project_flowchart(&self, graph: &CodeGraph) -> String {
        let mut lines = vec!["graph TD".to_string()];

        let (grouped, ungrouped) = graph.elements_by_module();

        for element in &ungrouped {
            let id = Self::sanitize_id(element.name());
            lines.push(format!("    {id}[{}]", element.name()));
        }

        for (group, elements) in &grouped {
            lines.push(format!("    subgraph {group}"));
            for element in elements {
                let id = Self::sanitize_id(element.name());
                lines.push(format!("        {id}[{}]", element.name()));
            }
            lines.push("    end".to_string());
        }

        for rel in graph.relationships() {
            let source_id = Self::sanitize_id(rel.source());
            let target_id = Self::sanitize_id(rel.target());
            lines.push(format!("    {source_id} --> {target_id}"));
        }

        lines.join("\n")
    }

    fn sanitize_id(name: &str) -> String {
        name.replace(['-', '.'], "_")
    }
}

impl DiagramRenderer for MermaidRenderer {
    fn render(&self, graph: &CodeGraph) -> Result<RenderOutput, DomainError> {
        let content = match self.level {
            DiagramLevel::Type => self.render_class_diagram(graph),
            DiagramLevel::Module => self.render_module_flowchart(graph),
            DiagramLevel::Project => self.render_project_flowchart(graph),
        };
        let file = RenderedFile::new("diagram.mmd", &content)?;
        Ok(RenderOutput::single(file))
    }
}