mermaid-go/pkg/parser/state.go

// Package parser provides state diagram parsing based on stateDiagram.jison
package parser

import (
	"fmt"
	"strings"

	"mermaid-go/pkg/ast"
	"mermaid-go/pkg/lexer"
)

// StateParser implements state diagram parsing following stateDiagram.jison
type StateParser struct {
	tokens  []lexer.Token
	current int
	diagram *ast.StateDiagram
}

// NewStateParser creates a new state parser
func NewStateParser() *StateParser {
	return &StateParser{
		diagram: ast.NewStateDiagram(),
	}
}

// Parse parses state diagram syntax
func (p *StateParser) Parse(input string) (*ast.StateDiagram, error) {
	// Tokenize
	l := lexer.NewLexer(input)
	tokens, err := l.Tokenize()
	if err != nil {
		return nil, fmt.Errorf("lexical analysis failed: %w", err)
	}

	// Filter tokens
	p.tokens = lexer.FilterTokens(tokens)
	p.current = 0
	p.diagram = ast.NewStateDiagram()

	// Parse document
	err = p.parseDocument()
	if err != nil {
		return nil, fmt.Errorf("syntax analysis failed: %w", err)
	}

	return p.diagram, nil
}

// parseDocument parses the state diagram document
func (p *StateParser) parseDocument() error {
	// Expect stateDiagram or stateDiagram-v2
	if !p.check(lexer.TokenID) ||
		(p.peek().Value != "stateDiagram" && p.peek().Value != "stateDiagram-v2") {
		return p.error("expected 'stateDiagram' or 'stateDiagram-v2'")
	}
	p.advance()

	// Parse statements
	for !p.isAtEnd() {
		if err := p.parseStatement(); err != nil {
			return err
		}
	}

	return nil
}

// parseStatement parses individual state diagram statements
func (p *StateParser) parseStatement() error {
	if p.isAtEnd() {
		return nil
	}

	switch {
	case p.check(lexer.TokenNewline):
		p.advance() // Skip newlines
		return nil
	case p.checkKeyword("direction"):
		return p.parseDirection()
	case p.checkKeyword("note"):
		return p.parseNote()
	case p.checkKeyword("state"):
		return p.parseState()
	case p.check(lexer.TokenOpenBracket):
		// Handle [*] start/end states
		return p.parseStartEndState()
	case p.check(lexer.TokenID):
		// Try to parse as state or transition
		return p.parseStateOrTransition()
	default:
		token := p.peek()
		return p.error(fmt.Sprintf("unexpected token: %s", token.Value))
	}
}

// parseState parses state declarations
func (p *StateParser) parseState() error {
	p.advance() // consume 'state'

	if !p.check(lexer.TokenID) {
		return p.error("expected state name")
	}

	stateName := p.advance().Value

	state := &ast.StateNode{
		ID:         stateName,
		Label:      stateName,
		Type:       ast.StateTypeDefault,
		SubStates:  make(map[string]*ast.StateNode),
		CssClasses: make([]string, 0),
	}

	// Check for 'as' alias
	if p.checkKeyword("as") {
		p.advance() // consume 'as'
		if !p.check(lexer.TokenID) && !p.check(lexer.TokenString) {
			return p.error("expected state label after 'as'")
		}
		label := p.advance().Value
		if strings.HasPrefix(label, "\"") && strings.HasSuffix(label, "\"") {
			label = label[1 : len(label)-1] // Remove quotes
		}
		state.Label = label
	}

	// Check for state body (composite state)
	if p.check(lexer.TokenOpenBrace) {
		p.advance() // consume '{'
		err := p.parseStateBody(state)
		if err != nil {
			return err
		}
		if !p.check(lexer.TokenCloseBrace) {
			return p.error("expected '}'")
		}
		p.advance() // consume '}'
	}

	// Check for special state types
	if p.check(lexer.TokenColon) {
		p.advance() // consume ':'
		if p.checkKeyword("<<fork>>") {
			p.advance()
			state.Type = ast.StateTypeFork
		} else if p.checkKeyword("<<join>>") {
			p.advance()
			state.Type = ast.StateTypeJoin
		} else if p.checkKeyword("<<choice>>") {
			p.advance()
			state.Type = ast.StateTypeChoice
		} else if p.checkKeyword("<<history>>") {
			p.advance()
			state.Type = ast.StateTypeHistory
		} else if p.checkKeyword("<<deepHistory>>") {
			p.advance()
			state.Type = ast.StateTypeDeepHistory
		} else {
			// Parse description
			var descParts []string
			for !p.check(lexer.TokenNewline) && !p.isAtEnd() {
				descParts = append(descParts, p.advance().Value)
			}
			if len(descParts) > 0 {
				desc := strings.TrimSpace(strings.Join(descParts, " "))
				state.Description = &desc
			}
		}
	}

	p.diagram.States[stateName] = state
	return nil
}

// parseStateBody parses the contents of a composite state
func (p *StateParser) parseStateBody(parentState *ast.StateNode) error {
	for !p.check(lexer.TokenCloseBrace) && !p.isAtEnd() {
		if p.check(lexer.TokenNewline) {
			p.advance()
			continue
		}

		// Parse sub-statements (simplified for now)
		if err := p.parseStatement(); err != nil {
			return err
		}
	}
	return nil
}

// parseStartEndState parses [*] --> state or state --> [*] transitions
func (p *StateParser) parseStartEndState() error {
	if !p.check(lexer.TokenOpenBracket) {
		return p.error("expected '['")
	}
	p.advance() // consume '['

	if !p.check(lexer.TokenMult) {
		return p.error("expected '*'")
	}
	p.advance() // consume '*'

	if !p.check(lexer.TokenCloseBracket) {
		return p.error("expected ']'")
	}
	p.advance() // consume ']'

	// Parse arrow
	if !p.checkArrow() {
		return p.error("expected transition arrow")
	}
	p.parseArrow()

	if !p.check(lexer.TokenID) {
		return p.error("expected target state")
	}
	targetState := p.advance().Value

	// Ensure target state exists
	p.ensureState(targetState)

	// Create transition
	transition := &ast.StateTransition{
		From: "[*]",
		To:   targetState,
	}

	// Check for label
	if p.check(lexer.TokenColon) {
		p.advance() // consume ':'
		var labelParts []string
		for !p.check(lexer.TokenNewline) && !p.isAtEnd() {
			labelParts = append(labelParts, p.advance().Value)
		}
		if len(labelParts) > 0 {
			label := strings.TrimSpace(strings.Join(labelParts, " "))
			transition.Label = &label
		}
	}

	p.diagram.Transitions = append(p.diagram.Transitions, transition)

	// Set start state if it's the first [*] transition
	if p.diagram.StartState == nil {
		start := "[*]"
		p.diagram.StartState = &start
	}

	return nil
}

// parseStateOrTransition parses either a state definition or transition
func (p *StateParser) parseStateOrTransition() error {
	stateName := p.advance().Value

	// Ensure state exists
	p.ensureState(stateName)

	// Check for transition arrow
	if p.checkArrow() {
		return p.parseTransition(stateName)
	}

	// Check for colon (description or special type)
	if p.check(lexer.TokenColon) {
		p.advance() // consume ':'
		var descParts []string
		for !p.check(lexer.TokenNewline) && !p.isAtEnd() {
			descParts = append(descParts, p.advance().Value)
		}
		if len(descParts) > 0 {
			desc := strings.TrimSpace(strings.Join(descParts, " "))
			state := p.diagram.States[stateName]
			state.Description = &desc
		}
	}

	return nil
}

// parseTransition parses state transitions
func (p *StateParser) parseTransition(fromState string) error {
	p.parseArrow()

	var toState string
	if p.check(lexer.TokenOpenBracket) {
		// Handle --> [*] end state
		p.advance() // consume '['
		if !p.check(lexer.TokenMult) {
			return p.error("expected '*'")
		}
		p.advance() // consume '*'
		if !p.check(lexer.TokenCloseBracket) {
			return p.error("expected ']'")
		}
		p.advance() // consume ']'
		toState = "[*]"

		// Add to end states if not already there
		found := false
		for _, endState := range p.diagram.EndStates {
			if endState == "[*]" {
				found = true
				break
			}
		}
		if !found {
			p.diagram.EndStates = append(p.diagram.EndStates, "[*]")
		}
	} else if p.check(lexer.TokenID) {
		toState = p.advance().Value
		p.ensureState(toState)
	} else {
		return p.error("expected target state")
	}

	transition := &ast.StateTransition{
		From: fromState,
		To:   toState,
	}

	// Check for label
	if p.check(lexer.TokenColon) {
		p.advance() // consume ':'
		var labelParts []string
		for !p.check(lexer.TokenNewline) && !p.isAtEnd() {
			labelParts = append(labelParts, p.advance().Value)
		}
		if len(labelParts) > 0 {
			label := strings.TrimSpace(strings.Join(labelParts, " "))
			transition.Label = &label
		}
	}

	p.diagram.Transitions = append(p.diagram.Transitions, transition)
	return nil
}

// parseArrow parses transition arrows
func (p *StateParser) parseArrow() string {
	token := p.peek()

	if token.Value == "-->" {
		p.advance()
		return "-->"
	} else if token.Value == "--" && p.checkNext(lexer.TokenCloseAngle) {
		p.advance() // consume '--'
		p.advance() // consume '>'
		return "-->"
	}

	// Default
	p.advance()
	return "-->"
}

// parseDirection parses direction statements
func (p *StateParser) parseDirection() error {
	p.advance() // consume 'direction'

	if !p.check(lexer.TokenID) {
		return p.error("expected direction value")
	}

	direction := p.advance().Value
	p.diagram.Direction = direction
	return nil
}

// parseNote parses note statements - placeholder
func (p *StateParser) parseNote() error {
	return p.skipToNextStatement()
}

// ensureState ensures a state exists, creating it if needed
func (p *StateParser) ensureState(id string) {
	if _, exists := p.diagram.States[id]; !exists {
		state := &ast.StateNode{
			ID:         id,
			Label:      id,
			Type:       ast.StateTypeDefault,
			SubStates:  make(map[string]*ast.StateNode),
			CssClasses: make([]string, 0),
		}
		p.diagram.States[id] = state
	}
}

// Helper methods
func (p *StateParser) check(tokenType lexer.TokenType) bool {
	if p.isAtEnd() {
		return false
	}
	return p.peek().Type == tokenType
}

func (p *StateParser) checkNext(tokenType lexer.TokenType) bool {
	if p.current+1 >= len(p.tokens) {
		return false
	}
	return p.tokens[p.current+1].Type == tokenType
}

func (p *StateParser) checkKeyword(keyword string) bool {
	if p.isAtEnd() {
		return false
	}
	token := p.peek()
	return token.Type == lexer.TokenID && strings.ToLower(token.Value) == strings.ToLower(keyword)
}

func (p *StateParser) checkArrow() bool {
	token := p.peek()
	return token.Value == "-->" || token.Value == "--"
}

func (p *StateParser) advance() lexer.Token {
	if !p.isAtEnd() {
		p.current++
	}
	return p.previous()
}

func (p *StateParser) isAtEnd() bool {
	return p.current >= len(p.tokens) || p.peek().Type == lexer.TokenEOF
}

func (p *StateParser) peek() lexer.Token {
	if p.current >= len(p.tokens) {
		return lexer.Token{Type: lexer.TokenEOF}
	}
	return p.tokens[p.current]
}

func (p *StateParser) previous() lexer.Token {
	if p.current <= 0 {
		return lexer.Token{Type: lexer.TokenEOF}
	}
	return p.tokens[p.current-1]
}

func (p *StateParser) error(message string) error {
	token := p.peek()
	return fmt.Errorf("parse error at line %d, column %d: %s (got %s)",
		token.Line, token.Column, message, token.Type)
}

func (p *StateParser) skipToNextStatement() error {
	for !p.isAtEnd() && !p.check(lexer.TokenNewline) {
		p.advance()
	}
	if p.check(lexer.TokenNewline) {
		p.advance()
	}
	return nil
}