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))
	}

	// Create a map to track which states need explicit declaration
	stateNeedsDeclaration := make(map[string]bool)

	// Check which states need explicit declaration
	for _, state := range diagram.States {
		needsDeclaration := false

		// State needs declaration if it has special type, alias, or description
		if state.Type != ast.StateTypeDefault || state.Label != state.ID || state.Description != nil {
			needsDeclaration = true
		}

		// State needs declaration if it has composite structure
		if len(state.SubStates) > 0 {
			needsDeclaration = true
		}

		stateNeedsDeclaration[state.ID] = needsDeclaration
	}

	// Separate transitions into external and internal
	var externalTransitions []*ast.StateTransition
	var internalTransitions = make(map[string][]*ast.StateTransition) // stateID -> transitions

	for _, transition := range diagram.Transitions {
		// Check if this is an internal transition (both states are within the same composite state)
		isInternal := false
		for _, state := range diagram.States {
			if len(state.SubStates) > 0 {
				fromInState := r.isStateInComposite(state, transition.From)
				toInState := r.isStateInComposite(state, transition.To)
				if fromInState && toInState {
					internalTransitions[state.ID] = append(internalTransitions[state.ID], transition)
					isInternal = true
					break
				}
			}
		}
		if !isInternal {
			externalTransitions = append(externalTransitions, transition)
		}
	}

	// Separate states into special states and others
	var specialStates []*ast.StateNode
	var otherStates []*ast.StateNode

	for _, state := range diagram.States {
		if stateNeedsDeclaration[state.ID] {
			// Check if this is a special state
			if state.Type == ast.StateTypeFork || state.Type == ast.StateTypeJoin ||
				state.Type == ast.StateTypeChoice || state.Type == ast.StateTypeHistory ||
				state.Type == ast.StateTypeDeepHistory {
				specialStates = append(specialStates, state)
			} else {
				otherStates = append(otherStates, state)
			}
		}
	}

	// Render special states first (sorted by type priority)
	r.sortStatesByTransitionOrder(specialStates, diagram.Transitions)
	for _, state := range specialStates {
		r.renderStateWithInternalTransitions(&builder, state, internalTransitions[state.ID])
	}

	// Render [*] --> state transitions
	for _, transition := range externalTransitions {
		if transition.From == "[*]" {
			r.renderTransition(&builder, transition)
		}
	}

	// Render other states that need explicit declaration (composite states, etc.)
	for _, state := range otherStates {
		r.renderStateWithInternalTransitions(&builder, state, internalTransitions[state.ID])
	}

	// Render state actions for states that don't need explicit declaration
	for _, state := range diagram.States {
		if !stateNeedsDeclaration[state.ID] {
			r.renderStateActions(&builder, state)
		}
	}

	// Render other external transitions (state --> [*] and state --> state)
	for _, transition := range externalTransitions {
		if transition.From != "[*]" {
			r.renderTransition(&builder, transition)
		}
	}

	return builder.String(), nil
}

// sortStatesByTransitionOrder sorts states by their order in the transition chain
func (r *StateRenderer) sortStatesByTransitionOrder(states []*ast.StateNode, transitions []*ast.StateTransition) {
	// For special states, sort by type priority: join -> fork -> choice -> others
	typePriority := map[ast.StateType]int{
		ast.StateTypeJoin:        0,
		ast.StateTypeFork:        1,
		ast.StateTypeChoice:      2,
		ast.StateTypeHistory:     3,
		ast.StateTypeDeepHistory: 4,
		ast.StateTypeDefault:     5,
	}

	// Sort states by type priority
	for i := 0; i < len(states); i++ {
		for j := i + 1; j < len(states); j++ {
			priorityI := typePriority[states[i].Type]
			priorityJ := typePriority[states[j].Type]
			if priorityI > priorityJ {
				states[i], states[j] = states[j], states[i]
			}
		}
	}
}

// renderTransition renders a single transition
func (r *StateRenderer) renderTransition(builder *strings.Builder, transition *ast.StateTransition) {
	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")
}

// renderStateActions renders state actions without explicit state declaration
func (r *StateRenderer) renderStateActions(builder *strings.Builder, state *ast.StateNode) {
	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")
	}
}

// renderState renders a state with explicit declaration
func (r *StateRenderer) renderState(builder *strings.Builder, state *ast.StateNode) {
	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
	r.renderStateActions(builder, state)
}

// isStateInComposite checks if a state ID is within a composite state
func (r *StateRenderer) isStateInComposite(compositeState *ast.StateNode, stateID string) bool {
	// Don't consider the composite state itself as internal
	if stateID == compositeState.ID {
		return false
	}

	// Check if it's a direct substate
	if _, exists := compositeState.SubStates[stateID]; exists {
		return true
	}

	// Check if it's [*] (start/end states are considered internal to composite states)
	if stateID == "[*]" {
		return true
	}

	// Check if it's within any nested composite state
	for _, subState := range compositeState.SubStates {
		if len(subState.SubStates) > 0 {
			// This is a nested composite state, check recursively
			if r.isStateInComposite(subState, stateID) {
				return true
			}
		}
	}

	return false
}

// renderStateWithInternalTransitions renders a state with its internal transitions
func (r *StateRenderer) renderStateWithInternalTransitions(builder *strings.Builder, state *ast.StateNode, internalTransitions []*ast.StateTransition) {
	// For composite states, only render the composite structure
	if len(state.SubStates) > 0 {
		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 transitions that belong to this composite state but not to nested composite states
		for _, transition := range internalTransitions {
			// Check if this transition belongs to a nested composite state
			belongsToNested := false
			for _, subState := range state.SubStates {
				if len(subState.SubStates) > 0 {
					// This is a nested composite state
					if r.isStateInComposite(subState, transition.From) && r.isStateInComposite(subState, transition.To) {
						belongsToNested = true
						break
					}
				}
			}

			if !belongsToNested {
				builder.WriteString("        ")
				builder.WriteString(transition.From)
				builder.WriteString(" --> ")
				builder.WriteString(transition.To)

				// Add transition decorations
				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")
			}
		}

		// Render nested composite states (in original order)
		// Create ordered slice from map to ensure consistent ordering
		var nestedStates []*ast.StateNode
		for _, subState := range state.SubStates {
			if len(subState.SubStates) > 0 {
				nestedStates = append(nestedStates, subState)
			}
		}

		// Sort nested states by their ID to ensure consistent ordering
		for i := 0; i < len(nestedStates); i++ {
			for j := i + 1; j < len(nestedStates); j++ {
				if nestedStates[i].ID > nestedStates[j].ID {
					nestedStates[i], nestedStates[j] = nestedStates[j], nestedStates[i]
				}
			}
		}

		for _, subState := range nestedStates {
			// This is a nested composite state, render it
			builder.WriteString("        state ")
			builder.WriteString(subState.ID)

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

			// Add description or special type
			if subState.Description != nil {
				builder.WriteString(" : ")
				builder.WriteString(*subState.Description)
			} else {
				switch subState.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 nested state transitions
			for _, nestedTransition := range internalTransitions {
				// Check if this transition is within the nested state
				if r.isStateInComposite(subState, nestedTransition.From) && r.isStateInComposite(subState, nestedTransition.To) {
					builder.WriteString("            ")
					builder.WriteString(nestedTransition.From)
					builder.WriteString(" --> ")
					builder.WriteString(nestedTransition.To)

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

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

			builder.WriteString("        }\n")
		}

		builder.WriteString("    }\n")
	} else {
		// For non-composite states, render the state declaration
		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 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
		r.renderStateActions(builder, state)
	}
}