mermaid-go/pkg/parser/class.go

// Package parser provides class diagram parsing based on classDiagram.jison
package parser

import (
	"fmt"
	"strings"

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

// ClassParser implements class diagram parsing following classDiagram.jison
type ClassParser struct {
	tokens  []lexer.Token
	current int
	diagram *ast.ClassDiagram
}

// NewClassParser creates a new class parser
func NewClassParser() *ClassParser {
	return &ClassParser{
		diagram: ast.NewClassDiagram(),
	}
}

// Parse parses class diagram syntax
func (p *ClassParser) Parse(input string) (*ast.ClassDiagram, 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.NewClassDiagram()

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

	return p.diagram, nil
}

// parseDocument parses the class diagram document
func (p *ClassParser) parseDocument() error {
	// Expect classDiagram
	if !p.check(lexer.TokenID) || p.peek().Value != "classDiagram" {
		return p.error("expected 'classDiagram'")
	}
	p.advance()

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

	return nil
}

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

	switch {
	case p.check(lexer.TokenNewline):
		p.advance() // Skip newlines
		return nil
	case p.check(lexer.TokenClass) || p.checkKeyword("class"):
		return p.parseClass()
	case p.checkKeyword("direction"):
		return p.parseDirection()
	case p.checkKeyword("link"):
		return p.parseLink()
	case p.checkKeyword("click"):
		return p.parseClick()
	case p.checkKeyword("note"):
		return p.parseNote()
	case p.checkKeyword("classDef"):
		return p.parseClassDef()
	case p.check(lexer.TokenID):
		// Try to parse as class definition or relation
		return p.parseClassOrRelation()
	default:
		token := p.peek()
		return p.error(fmt.Sprintf("unexpected token: %s", token.Value))
	}
}

// parseClass parses class statements
func (p *ClassParser) parseClass() error {
	// Consume 'class' token
	if p.check(lexer.TokenClass) {
		p.advance() // consume TokenClass
	} else {
		p.advance() // consume 'class' keyword
	}

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

	className := p.advance().Value

	class := &ast.ClassNode{
		ID:          className,
		Label:       className,
		Type:        ast.ClassTypeClass,
		Members:     make([]*ast.ClassMember, 0),
		Methods:     make([]*ast.ClassMethod, 0),
		Annotations: make([]string, 0),
		CssClasses:  make([]string, 0),
	}

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

	p.diagram.Classes[className] = class
	return nil
}

// parseClassBody parses the contents of a class body
func (p *ClassParser) parseClassBody(class *ast.ClassNode) error {
	for !p.check(lexer.TokenCloseBrace) && !p.isAtEnd() {
		if p.check(lexer.TokenNewline) {
			p.advance()
			continue
		}

		// Parse member or method
		visibility := ast.VisibilityPublic // default
		if p.checkVisibility() {
			switch p.peek().Value {
			case "+":
				visibility = ast.VisibilityPublic
			case "-":
				visibility = ast.VisibilityPrivate
			case "#":
				visibility = ast.VisibilityProtected
			case "~":
				visibility = ast.VisibilityPackage
			}
			p.advance()
		}

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

		name := p.advance().Value
		var memberType string

		// Check for type annotation
		if p.check(lexer.TokenColon) {
			p.advance() // consume ':'
			if p.check(lexer.TokenID) {
				memberType = p.advance().Value
			}
		}

		// Check if it's a method (has parentheses)
		if p.check(lexer.TokenOpenParen) {
			p.advance() // consume '('

			method := &ast.ClassMethod{
				Name:       name,
				Type:       memberType,
				Parameters: make([]string, 0),
				Visibility: visibility,
			}

			// Parse parameters
			for !p.check(lexer.TokenCloseParen) && !p.isAtEnd() {
				if p.check(lexer.TokenID) {
					param := p.advance().Value
					method.Parameters = append(method.Parameters, param)
				}
				if p.check(lexer.TokenComma) {
					p.advance()
				}
			}

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

			class.Methods = append(class.Methods, method)
		} else {
			// It's a member
			member := &ast.ClassMember{
				Name:       name,
				Type:       memberType,
				Visibility: visibility,
			}
			class.Members = append(class.Members, member)
		}
	}

	return nil
}

// parseClassOrRelation parses either a class definition or relationship
func (p *ClassParser) parseClassOrRelation() error {
	className := p.advance().Value

	// Ensure class exists
	p.ensureClass(className)

	// Check for relationship operators
	if p.checkRelation() {
		return p.parseRelation(className)
	}

	// Check for class body
	if p.check(lexer.TokenOpenBrace) {
		class := p.diagram.Classes[className]
		p.advance() // consume '{'
		err := p.parseClassBody(class)
		if err != nil {
			return err
		}
		if !p.check(lexer.TokenCloseBrace) {
			return p.error("expected '}'")
		}
		p.advance() // consume '}'
	}

	return nil
}

// parseRelation parses class relationships
func (p *ClassParser) parseRelation(fromClass string) error {
	relationType := p.parseRelationType()
	if relationType == "" {
		return p.error("expected relationship operator")
	}

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

	// Ensure target class exists
	p.ensureClass(toClass)

	relation := &ast.ClassRelation{
		From: fromClass,
		To:   toClass,
		Type: relationType,
	}

	// 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, " "))
			relation.Label = &label
		}
	}

	p.diagram.Relations = append(p.diagram.Relations, relation)
	return nil
}

// parseRelationType parses relationship type tokens
func (p *ClassParser) parseRelationType() ast.ClassRelationType {
	token := p.peek()

	// Check for direct arrow tokens
	if p.check(lexer.TokenArrowSolid) {
		p.advance() // consume -->
		return ast.RelationAssociation
	}
	if p.check(lexer.TokenArrowDotted) {
		p.advance() // consume -.->
		return ast.RelationDependency
	}

	// Check for inheritance: --|>
	if token.Value == "--" && p.checkNext(lexer.TokenPipe) && p.checkAt(2, lexer.TokenCloseAngle) {
		p.advance() // --
		p.advance() // |
		p.advance() // >
		return ast.RelationInheritance
	}

	// Check for composition: --*
	if token.Value == "--" && p.checkNextValue("*") {
		p.advance() // --
		p.advance() // *
		return ast.RelationComposition
	}

	// Check for aggregation: --o
	if token.Value == "--" && p.checkNextValue("o") {
		p.advance() // --
		p.advance() // o
		return ast.RelationAggregation
	}

	// Check for association: -->
	if token.Value == "--" && p.checkNext(lexer.TokenCloseAngle) {
		p.advance() // --
		p.advance() // >
		return ast.RelationAssociation
	}

	// Check for realization: ..|>
	if token.Value == ".." && p.checkNext(lexer.TokenPipe) && p.checkAt(2, lexer.TokenCloseAngle) {
		p.advance() // ..
		p.advance() // |
		p.advance() // >
		return ast.RelationRealization
	}

	// Check for dependency: ..>
	if token.Value == ".." && p.checkNext(lexer.TokenCloseAngle) {
		p.advance() // ..
		p.advance() // >
		return ast.RelationDependency
	}

	return ""
}

// parseDirection parses direction statements
func (p *ClassParser) 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
}

// parseLink, parseClick, parseNote, parseClassDef - placeholder implementations
func (p *ClassParser) parseLink() error {
	return p.skipToNextStatement()
}

func (p *ClassParser) parseClick() error {
	return p.skipToNextStatement()
}

func (p *ClassParser) parseNote() error {
	return p.skipToNextStatement()
}

func (p *ClassParser) parseClassDef() error {
	return p.skipToNextStatement()
}

// ensureClass ensures a class exists, creating it if needed
func (p *ClassParser) ensureClass(id string) {
	if _, exists := p.diagram.Classes[id]; !exists {
		class := &ast.ClassNode{
			ID:          id,
			Label:       id,
			Type:        ast.ClassTypeClass,
			Members:     make([]*ast.ClassMember, 0),
			Methods:     make([]*ast.ClassMethod, 0),
			Annotations: make([]string, 0),
			CssClasses:  make([]string, 0),
		}
		p.diagram.Classes[id] = class
	}
}

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

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

func (p *ClassParser) checkAt(offset int, tokenType lexer.TokenType) bool {
	if p.current+offset >= len(p.tokens) {
		return false
	}
	return p.tokens[p.current+offset].Type == tokenType
}

func (p *ClassParser) checkNextValue(value string) bool {
	if p.current+1 >= len(p.tokens) {
		return false
	}
	return p.tokens[p.current+1].Value == value
}

func (p *ClassParser) 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 *ClassParser) checkVisibility() bool {
	if p.isAtEnd() {
		return false
	}
	token := p.peek()
	return token.Value == "+" || token.Value == "-" || token.Value == "#" || token.Value == "~"
}

func (p *ClassParser) checkRelation() bool {
	token := p.peek()
	// Check for various relation operators
	return token.Value == "--" || token.Value == ".." ||
		p.check(lexer.TokenArrowSolid) || p.check(lexer.TokenArrowDotted) ||
		token.Value == "--|>" || token.Value == "--*" || token.Value == "--o"
}

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

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

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

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

func (p *ClassParser) 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 *ClassParser) skipToNextStatement() error {
	for !p.isAtEnd() && !p.check(lexer.TokenNewline) {
		p.advance()
	}
	if p.check(lexer.TokenNewline) {
		p.advance()
	}
	return nil
}