mermaid-go/pkg/renderer/state.go

// Package renderer provides rendering functionality for state diagrams
package renderer

import (
	"fmt"
	"strings"

	"mermaid-go/pkg/ast"
)

// StateRenderer renders state diagrams back to mermaid syntax
type StateRenderer struct{}

// NewStateRenderer creates a new state renderer
func NewStateRenderer() *StateRenderer {
	return &StateRenderer{}
}

// Render renders a state diagram to mermaid syntax
func (r *StateRenderer) Render(diagram *ast.StateDiagram) (string, error) {
	var builder strings.Builder

	// Start with diagram declaration
	builder.WriteString("stateDiagram-v2\n")

	// Add title if present
	if diagram.Title != nil {
		builder.WriteString(fmt.Sprintf("    title %s\n", *diagram.Title))
	}

	// Add direction if present
	if diagram.Direction != "" {
		builder.WriteString(fmt.Sprintf("    direction %s\n", diagram.Direction))
	}

	// Render states
	for _, state := range diagram.States {
		builder.WriteString("    state ")
		builder.WriteString(state.ID)

		// Add alias if different from ID
		if state.Label != state.ID {
			builder.WriteString(" as ")
			if strings.Contains(state.Label, " ") {
				builder.WriteString(fmt.Sprintf("\"%s\"", state.Label))
			} else {
				builder.WriteString(state.Label)
			}
		}

		// Add description or special type
		if state.Description != nil {
			builder.WriteString(" : ")
			builder.WriteString(*state.Description)
		} else {
			switch state.Type {
			case ast.StateTypeFork:
				builder.WriteString(" : <<fork>>")
			case ast.StateTypeJoin:
				builder.WriteString(" : <<join>>")
			case ast.StateTypeChoice:
				builder.WriteString(" : <<choice>>")
			case ast.StateTypeHistory:
				builder.WriteString(" : <<history>>")
			case ast.StateTypeDeepHistory:
				builder.WriteString(" : <<deepHistory>>")
			}
		}

		builder.WriteString("\n")

		// Render composite state body if it has sub-states
		if len(state.SubStates) > 0 {
			builder.WriteString("    state ")
			builder.WriteString(state.ID)
			builder.WriteString(" {\n")

			for _, subState := range state.SubStates {
				builder.WriteString("        state ")
				builder.WriteString(subState.ID)
				if subState.Label != subState.ID {
					builder.WriteString(" as ")
					builder.WriteString(subState.Label)
				}
				if subState.Description != nil {
					builder.WriteString(" : ")
					builder.WriteString(*subState.Description)
				}
				builder.WriteString("\n")
			}

			builder.WriteString("    }\n")
		}

		// Render note if present
		if state.Note != nil {
			builder.WriteString("    note ")
			builder.WriteString(string(state.Note.Position))
			builder.WriteString(" ")
			builder.WriteString(state.ID)
			builder.WriteString(" : ")
			builder.WriteString(state.Note.Text)
			builder.WriteString("\n")
		}

		// Render state actions
		if state.EntryAction != nil {
			builder.WriteString("    ")
			builder.WriteString(state.ID)
			builder.WriteString(" : entry ")
			builder.WriteString(*state.EntryAction)
			builder.WriteString("\n")
		}
		if state.ExitAction != nil {
			builder.WriteString("    ")
			builder.WriteString(state.ID)
			builder.WriteString(" : exit ")
			builder.WriteString(*state.ExitAction)
			builder.WriteString("\n")
		}
		if state.DoAction != nil {
			builder.WriteString("    ")
			builder.WriteString(state.ID)
			builder.WriteString(" : do ")
			builder.WriteString(*state.DoAction)
			builder.WriteString("\n")
		}
	}

	// Render transitions
	for _, transition := range diagram.Transitions {
		builder.WriteString("    ")
		builder.WriteString(transition.From)
		builder.WriteString(" --> ")
		builder.WriteString(transition.To)

		// Add transition decorations: label, [guard], /action
		var decorations []string
		if transition.Label != nil {
			decorations = append(decorations, *transition.Label)
		}
		if transition.Condition != nil {
			decorations = append(decorations, fmt.Sprintf("[%s]", *transition.Condition))
		}
		if transition.Action != nil {
			decorations = append(decorations, fmt.Sprintf("/%s", *transition.Action))
		}

		if len(decorations) > 0 {
			builder.WriteString(" : ")
			builder.WriteString(strings.Join(decorations, " "))
		}
		builder.WriteString("\n")
	}

	return builder.String(), nil
}