mermaid-go/pkg/renderer/flowchart.go

// Package renderer provides rendering functionality to convert AST back to Mermaid syntax.
// Based on the rendering patterns from mermaid.js
package renderer

import (
	"fmt"
	"strings"

	"mermaid-go/pkg/ast"
)

// Renderer interface for converting diagrams back to mermaid syntax
type Renderer interface {
	Render(diagram ast.Diagram) string
}

// FlowchartRenderer renders flowchart diagrams back to mermaid syntax
type FlowchartRenderer struct{}

// NewFlowchartRenderer creates a new flowchart renderer
func NewFlowchartRenderer() *FlowchartRenderer {
	return &FlowchartRenderer{}
}

// Render converts a flowchart back to mermaid syntax
func (r *FlowchartRenderer) Render(flowchart *ast.Flowchart) (string, error) {

	var builder strings.Builder

	// Write graph declaration with direction
	direction := flowchart.Direction
	if direction == "" {
		direction = "TD" // Default direction
	}
	builder.WriteString(fmt.Sprintf("graph %s\n", direction))

	// Collect all vertices that are referenced in edges for proper node definitions
	referencedVertices := make(map[string]bool)
	for _, edge := range flowchart.Edges {
		referencedVertices[edge.Start] = true
		referencedVertices[edge.End] = true
	}

	// Render edges (which implicitly define vertices)
	for _, edge := range flowchart.Edges {
		line := r.renderEdge(edge, flowchart.Vertices)
		builder.WriteString("    ")
		builder.WriteString(line)
		builder.WriteString("\n")
	}

	// Render standalone vertices (not connected to any edges)
	for id, vertex := range flowchart.Vertices {
		if !referencedVertices[id] {
			line := r.renderStandaloneVertex(vertex)
			if line != "" {
				builder.WriteString("    ")
				builder.WriteString(line)
				builder.WriteString("\n")
			}
		}
	}

	// Render subgraphs
	for _, subGraph := range flowchart.SubGraphs {
		r.renderSubGraph(&builder, subGraph)
	}

	// Render class definitions
	for _, class := range flowchart.Classes {
		r.renderClassDef(&builder, class)
	}

	// Render class assignments
	classAssignments := make(map[string][]string)
	for _, vertex := range flowchart.Vertices {
		for _, className := range vertex.Classes {
			if classAssignments[className] == nil {
				classAssignments[className] = make([]string, 0)
			}
			classAssignments[className] = append(classAssignments[className], vertex.ID)
		}
	}

	for className, nodeIDs := range classAssignments {
		if len(nodeIDs) > 0 {
			builder.WriteString(fmt.Sprintf("    class %s %s\n", strings.Join(nodeIDs, ","), className))
		}
	}

	// Render individual node styles
	for _, vertex := range flowchart.Vertices {
		if len(vertex.Styles) > 0 {
			styles := strings.Join(vertex.Styles, " ")
			builder.WriteString(fmt.Sprintf("    style %s %s\n", vertex.ID, styles))
		}
	}

	// Render click events
	for _, vertex := range flowchart.Vertices {
		if vertex.OnClick != nil {
			if vertex.OnClick.Link != nil {
				if vertex.OnClick.Target != nil {
					builder.WriteString(fmt.Sprintf("    click %s \"%s\" \"%s\"\n",
						vertex.ID, *vertex.OnClick.Link, *vertex.OnClick.Target))
				} else {
					builder.WriteString(fmt.Sprintf("    click %s \"%s\"\n",
						vertex.ID, *vertex.OnClick.Link))
				}
			} else if vertex.OnClick.Callback != nil {
				builder.WriteString(fmt.Sprintf("    click %s %s\n",
					vertex.ID, *vertex.OnClick.Callback))
			}
		}
	}

	// Render tooltips
	for nodeID, tooltip := range flowchart.Tooltips {
		builder.WriteString(fmt.Sprintf("    %s --> tooltip[\"%s\"]\n", nodeID, tooltip))
	}

	return builder.String(), nil
}

// renderEdge renders an edge with its connected vertices
func (r *FlowchartRenderer) renderEdge(edge *ast.FlowEdge, vertices map[string]*ast.FlowVertex) string {
	startPart := r.renderVertexReference(edge.Start, vertices)
	arrow := r.renderArrow(edge)
	endPart := r.renderVertexReference(edge.End, vertices)

	return fmt.Sprintf("%s %s %s", startPart, arrow, endPart)
}

// renderVertexReference renders a vertex with its shape and label
func (r *FlowchartRenderer) renderVertexReference(vertexID string, vertices map[string]*ast.FlowVertex) string {
	vertex := vertices[vertexID]
	if vertex == nil {
		return vertexID
	}

	return r.renderVertexWithShape(vertex)
}

// renderVertexWithShape renders a vertex with its shape delimiters
func (r *FlowchartRenderer) renderVertexWithShape(vertex *ast.FlowVertex) string {
	text := vertex.ID
	if vertex.Text != nil && *vertex.Text != "" {
		text = *vertex.Text
	}

	// Determine shape based on vertex type
	vertexType := ast.VertexTypeRect // default
	if vertex.Type != nil {
		vertexType = *vertex.Type
	}

	switch vertexType {
	case ast.VertexTypeRect, ast.VertexTypeSquare:
		return fmt.Sprintf("%s[%s]", vertex.ID, text)
	case ast.VertexTypeRound:
		return fmt.Sprintf("%s(%s)", vertex.ID, text)
	case ast.VertexTypeCircle:
		return fmt.Sprintf("%s((%s))", vertex.ID, text)
	case ast.VertexTypeDiamond:
		return fmt.Sprintf("%s{%s}", vertex.ID, text)
	case ast.VertexTypeFlag:
		return fmt.Sprintf("%s>%s]", vertex.ID, text)
	case ast.VertexTypeLeanRight:
		return fmt.Sprintf("%s[/%s/]", vertex.ID, text)
	case ast.VertexTypeLeanLeft:
		return fmt.Sprintf("%s[\\%s\\]", vertex.ID, text)
	case ast.VertexTypeStadium:
		return fmt.Sprintf("%s([%s])", vertex.ID, text)
	case ast.VertexTypeCylinder:
		return fmt.Sprintf("%s[(%s)]", vertex.ID, text)
	case ast.VertexTypeSubroutine:
		return fmt.Sprintf("%s[[%s]]", vertex.ID, text)
	case ast.VertexTypeHexagon:
		return fmt.Sprintf("%s{{%s}}", vertex.ID, text)
	case ast.VertexTypeOdd:
		return fmt.Sprintf("%s>%s]", vertex.ID, text)
	case ast.VertexTypeTrapezoid:
		return fmt.Sprintf("%s[/%s/]", vertex.ID, text)
	case ast.VertexTypeInvTrapezoid:
		return fmt.Sprintf("%s[\\%s\\]", vertex.ID, text)
	default:
		return fmt.Sprintf("%s[%s]", vertex.ID, text)
	}
}

// renderArrow renders the arrow part of an edge with optional label
func (r *FlowchartRenderer) renderArrow(edge *ast.FlowEdge) string {
	// Build arrow based on stroke and type
	arrow := r.buildArrowString(edge)

	// Add label if present
	if edge.Text != "" {
		return fmt.Sprintf("%s|%s|%s", r.getArrowStart(edge), edge.Text, r.getArrowEnd(edge))
	}

	return arrow
}

// buildArrowString creates the arrow string based on edge properties
func (r *FlowchartRenderer) buildArrowString(edge *ast.FlowEdge) string {
	stroke := ast.StrokeNormal
	if edge.Stroke != nil {
		stroke = *edge.Stroke
	}

	edgeType := "arrow_point"
	if edge.Type != nil {
		edgeType = *edge.Type
	}

	switch stroke {
	case ast.StrokeNormal:
		switch edgeType {
		case "arrow_point":
			return "-->"
		case "arrow_cross":
			return "--x"
		case "arrow_circle":
			return "--o"
		case "arrow_open":
			return "---"
		default:
			return "-->"
		}
	case ast.StrokeThick:
		return "==>"
	case ast.StrokeDotted:
		return "-.->"

	case ast.StrokeInvisible:
		return "~~~"
	default:
		return "-->"
	}
}

// getArrowStart returns the start part of arrow for labeled edges
func (r *FlowchartRenderer) getArrowStart(edge *ast.FlowEdge) string {
	stroke := ast.StrokeNormal
	if edge.Stroke != nil {
		stroke = *edge.Stroke
	}

	switch stroke {
	case ast.StrokeThick:
		return "=="
	case ast.StrokeDotted:
		return "-."
	case ast.StrokeInvisible:
		return "~~"
	default:
		return "--"
	}
}

// getArrowEnd returns the end part of arrow for labeled edges
func (r *FlowchartRenderer) getArrowEnd(edge *ast.FlowEdge) string {
	edgeType := "arrow_point"
	if edge.Type != nil {
		edgeType = *edge.Type
	}

	switch edgeType {
	case "arrow_point":
		return ">"
	case "arrow_cross":
		return "x"
	case "arrow_circle":
		return "o"
	case "arrow_open":
		return ""
	default:
		return ">"
	}
}

// renderStandaloneVertex renders vertices not connected to any edges
func (r *FlowchartRenderer) renderStandaloneVertex(vertex *ast.FlowVertex) string {
	// Only render if vertex has explicit shape/text definition
	if vertex.Text != nil || vertex.Type != nil {
		return r.renderVertexWithShape(vertex)
	}
	return ""
}

// renderSubGraph renders a subgraph definition
func (r *FlowchartRenderer) renderSubGraph(builder *strings.Builder, subGraph *ast.FlowSubGraph) {
	builder.WriteString(fmt.Sprintf("    subgraph %s", subGraph.ID))
	if subGraph.Title != "" {
		builder.WriteString(fmt.Sprintf("[%s]", subGraph.Title))
	}
	builder.WriteString("\n")

	// Render direction if specified
	if subGraph.Dir != nil && *subGraph.Dir != "" {
		builder.WriteString(fmt.Sprintf("        direction %s\n", *subGraph.Dir))
	}

	// Render subgraph nodes
	for _, nodeID := range subGraph.Nodes {
		builder.WriteString(fmt.Sprintf("        %s\n", nodeID))
	}

	builder.WriteString("    end\n")
}

// renderClassDef renders a class definition
func (r *FlowchartRenderer) renderClassDef(builder *strings.Builder, class *ast.FlowClass) {
	if len(class.Styles) > 0 {
		styles := strings.Join(class.Styles, " ")
		builder.WriteString(fmt.Sprintf("    classDef %s %s\n", class.ID, styles))
	}
}