pipecircuit-visualization 0.0.5

use crate::{DiagramNode, StateMachineDiagram};

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum DiagramTarget {
    TextTree,
    AnsiTextTree,
    MermaidFlowchart,
    GraphvizDot,
}

pub(super) fn render(diagram: &StateMachineDiagram, target: DiagramTarget) -> String {
    match target {
        DiagramTarget::TextTree => text_tree(diagram, false),
        DiagramTarget::AnsiTextTree => text_tree(diagram, true),
        DiagramTarget::MermaidFlowchart => mermaid(diagram),
        DiagramTarget::GraphvizDot => graphviz(diagram),
    }
}

fn text_tree(diagram: &StateMachineDiagram, ansi: bool) -> String {
    let mut rendering = String::new();
    push_role_marker(&mut rendering, "●", ansi);
    rendering.push_str(" state machine  ");
    push_role_marker(&mut rendering, "▸", ansi);
    rendering.push_str(" step input  ");
    push_role_marker(&mut rendering, "◆", ansi);
    rendering.push_str(" outcome\n\n");
    push_role_marker(&mut rendering, "●", ansi);
    rendering.push(' ');
    rendering.push_str(&diagram.name);
    rendering.push('\n');
    for (index, step) in diagram.steps.iter().enumerate() {
        let final_step = index == diagram.steps.len().saturating_sub(1);
        rendering.push_str(if final_step {
            "└── "
        } else {
            "├── "
        });
        push_role_marker(&mut rendering, "▸", ansi);
        rendering.push(' ');
        rendering.push_str(&step.input);
        rendering.push('\n');
        let prefix = if final_step { "    " } else { "│   " };
        render_text_outcome(&mut rendering, &step.outcome, prefix, ansi);
    }
    rendering
}

fn push_role_marker(rendering: &mut String, marker: &str, ansi: bool) {
    if ansi {
        rendering.push_str("\u{1b}[38;5;114m");
    }
    rendering.push_str(marker);
    if ansi {
        rendering.push_str("\u{1b}[0m");
    }
}

fn render_text_outcome(rendering: &mut String, outcome: &DiagramNode, prefix: &str, ansi: bool) {
    rendering.push_str(prefix);
    rendering.push_str("└── ");
    push_role_marker(rendering, "◆", ansi);
    rendering.push(' ');
    rendering.push_str(&outcome.label);
    rendering.push('\n');
    let mut child_prefix = String::from(prefix);
    child_prefix.push_str("    ");
    for (index, child) in outcome.children.iter().enumerate() {
        render_text_node(
            rendering,
            child,
            &child_prefix,
            index == outcome.children.len().saturating_sub(1),
        );
    }
}

fn render_text_node(rendering: &mut String, node: &DiagramNode, prefix: &str, final_node: bool) {
    rendering.push_str(prefix);
    rendering.push_str(if final_node {
        "└── "
    } else {
        "├── "
    });
    rendering.push_str(&node.label);
    rendering.push('\n');
    let mut child_prefix = String::from(prefix);
    child_prefix.push_str(if final_node { "    " } else { "│   " });
    for (index, child) in node.children.iter().enumerate() {
        render_text_node(
            rendering,
            child,
            &child_prefix,
            index == node.children.len().saturating_sub(1),
        );
    }
}

struct FlatGraph {
    nodes: Vec<String>,
    edges: Vec<(usize, usize)>,
}

impl FlatGraph {
    fn from_diagram(diagram: &StateMachineDiagram) -> Self {
        let mut graph = Self {
            nodes: vec![diagram.name.clone()],
            edges: Vec::new(),
        };
        for step in &diagram.steps {
            let input = graph.push_node(step.input.clone());
            graph.edges.push((0, input));
            graph.push_tree(input, &step.outcome);
        }
        graph
    }

    fn push_node(&mut self, label: String) -> usize {
        let identifier = self.nodes.len();
        self.nodes.push(label);
        identifier
    }

    fn push_tree(&mut self, parent: usize, node: &DiagramNode) {
        let identifier = self.push_node(node.label.clone());
        self.edges.push((parent, identifier));
        for child in &node.children {
            self.push_tree(identifier, child);
        }
    }
}

fn mermaid(diagram: &StateMachineDiagram) -> String {
    let graph = FlatGraph::from_diagram(diagram);
    let mut rendering = String::from("flowchart LR\n");
    for (identifier, label) in graph.nodes.iter().enumerate() {
        rendering.push_str(&format!("    n{identifier}[\"{}\"]\n", escape_label(label)));
    }
    for (from, to) in graph.edges {
        rendering.push_str(&format!("    n{from} --> n{to}\n"));
    }
    rendering
}

fn graphviz(diagram: &StateMachineDiagram) -> String {
    let graph = FlatGraph::from_diagram(diagram);
    let mut rendering = format!("digraph \"{}\" {{\n", escape_label(&diagram.name));
    rendering.push_str("    rankdir=LR;\n");
    for (identifier, label) in graph.nodes.iter().enumerate() {
        rendering.push_str(&format!(
            "    n{identifier} [label=\"{}\"];\n",
            escape_label(label)
        ));
    }
    for (from, to) in graph.edges {
        rendering.push_str(&format!("    n{from} -> n{to};\n"));
    }
    rendering.push_str("}\n");
    rendering
}

fn escape_label(label: &str) -> String {
    label.replace('"', "&quot;")
}