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 transitions
	for _, transition := range diagram.Transitions {
		builder.WriteString("    ")

		// Handle start state
		if transition.From == "[*]" {
			builder.WriteString("[*]")
		} else {
			builder.WriteString(transition.From)
		}

		builder.WriteString(" --> ")

		// Handle end state
		if transition.To == "[*]" {
			builder.WriteString("[*]")
		} else {
			builder.WriteString(transition.To)
		}

		// Add label if present
		if transition.Label != nil {
			builder.WriteString(" : ")
			builder.WriteString(*transition.Label)
		}

		// Add condition if present
		if transition.Condition != nil {
			builder.WriteString(" [")
			builder.WriteString(*transition.Condition)
			builder.WriteString("]")
		}

		// Add action if present
		if transition.Action != nil {
			builder.WriteString(" / ")
			builder.WriteString(*transition.Action)
		}

		builder.WriteString("\n")
	}

	return builder.String(), nil
}