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.checkComment():
		return p.parseComment()
	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()
		}

		// Check for annotations like <<interface>>, <<abstract>>
		if p.check(lexer.TokenOpenAngle) {
			// Look ahead to see if it's << (double angle)
			if p.checkNext(lexer.TokenOpenAngle) {
				p.advance() // consume first <
				p.advance() // consume second <

				annotation := ""
				for !p.check(lexer.TokenCloseAngle) && !p.isAtEnd() {
					if p.check(lexer.TokenID) {
						annotation += p.advance().Value
					} else {
						annotation += p.advance().Value
					}
				}

				// Consume closing >>
				if p.check(lexer.TokenCloseAngle) {
					p.advance() // consume first >
					if p.check(lexer.TokenCloseAngle) {
						p.advance() // consume second >
						class.Annotations = append(class.Annotations, annotation)

						// Update class type based on annotation
						switch strings.ToLower(annotation) {
						case "interface":
							class.Type = ast.ClassTypeInterface
						case "abstract":
							class.Type = ast.ClassTypeAbstract
						case "enumeration", "enum":
							class.Type = ast.ClassTypeEnum
						}
					}
				}
				continue
			}
		}

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

		name := p.advance().Value

		// Check if it's a method (has parentheses)
		if p.check(lexer.TokenOpenParen) {
			method, err := p.parseMethod(name, visibility)
			if err != nil {
				return err
			}
			class.Methods = append(class.Methods, method)
		} else {
			// It's a member
			member, err := p.parseMember(name, visibility)
			if err != nil {
				return err
			}
			class.Members = append(class.Members, member)
		}
	}

	return nil
}

// parseMethod parses a method with parameters and return type
func (p *ClassParser) parseMethod(name string, visibility ast.MemberVisibility) (*ast.ClassMethod, error) {
	method := &ast.ClassMethod{
		Name:       name,
		Parameters: make([]string, 0),
		Visibility: visibility,
	}

	// Parse opening parenthesis
	if !p.check(lexer.TokenOpenParen) {
		return nil, p.error("expected '('")
	}
	p.advance() // consume '('

	// Parse parameters
	for !p.check(lexer.TokenCloseParen) && !p.isAtEnd() {
		if p.check(lexer.TokenID) {
			param := p.advance().Value

			// Check for parameter type
			if p.check(lexer.TokenID) {
				paramType := p.advance().Value
				method.Parameters = append(method.Parameters, param+" "+paramType)
			} else {
				method.Parameters = append(method.Parameters, param)
			}
		}

		// Skip commas
		if p.check(lexer.TokenComma) {
			p.advance()
		}
	}

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

	// Check for return type
	if p.check(lexer.TokenID) {
		returnType := p.advance().Value
		method.Type = returnType
	}

	return method, nil
}

// parseMember parses a class member/field
func (p *ClassParser) parseMember(name string, visibility ast.MemberVisibility) (*ast.ClassMember, error) {
	member := &ast.ClassMember{
		Name:       name,
		Visibility: visibility,
	}

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

	return member, 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'

	// Check for direction tokens or ID
	var direction string
	if p.check(lexer.TokenTD) || p.check(lexer.TokenTB) || p.check(lexer.TokenBT) ||
		p.check(lexer.TokenRL) || p.check(lexer.TokenLR) {
		direction = p.advance().Value
	} else if p.check(lexer.TokenID) {
		direction = p.advance().Value
	} else {
		return p.error("expected direction value (TD, TB, BT, RL, LR)")
	}

	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) 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 == "~") ||
		(token.Type == lexer.TokenPlus || token.Type == lexer.TokenMinus ||
		 token.Type == lexer.TokenHash || token.Type == lexer.TokenTilde)
}

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
}

// checkComment checks if current token sequence is a comment (%%)
func (p *ClassParser) checkComment() bool {
	return p.check(lexer.TokenPercent) && p.checkNext(lexer.TokenPercent)
}

// parseComment parses comment statements
func (p *ClassParser) parseComment() error {
	p.advance() // consume first %
	p.advance() // consume second %

	// Skip everything until newline
	for !p.check(lexer.TokenNewline) && !p.isAtEnd() {
		p.advance()
	}
	if p.check(lexer.TokenNewline) {
		p.advance()
	}
	return nil
}

// Enhanced parseNote to support class-specific notes
func (p *ClassParser) parseNote() error {
	p.advance() // consume 'note'

	var note *ast.ClassNote

	// Check if it's "note for ClassName"
	if p.checkKeyword("for") {
		p.advance() // consume 'for'
		if !p.check(lexer.TokenID) {
			return p.error("expected class name after 'note for'")
		}
		className := p.advance().Value

		// Parse note text
		text := ""
		for !p.check(lexer.TokenNewline) && !p.isAtEnd() {
			token := p.advance()
			if token.Type == lexer.TokenString {
				// Remove quotes from string tokens
				val := token.Value
				text += val[1 : len(val)-1] // Remove surrounding quotes
			} else {
				text += token.Value + " "
			}
		}

		note = &ast.ClassNote{
			ForClass: &className,
			Text:     strings.TrimSpace(text),
		}
	} else {
		// General note for the whole diagram
		text := ""
		for !p.check(lexer.TokenNewline) && !p.isAtEnd() {
			token := p.advance()
			if token.Type == lexer.TokenString {
				// Remove quotes from string tokens
				val := token.Value
				text += val[1 : len(val)-1] // Remove surrounding quotes
			} else {
				text += token.Value + " "
			}
		}

		note = &ast.ClassNote{
			Text: strings.TrimSpace(text),
		}
	}

	p.diagram.Notes = append(p.diagram.Notes, note)
	return nil
}

// parseGeneric parses generic type parameters with ~Type~
func (p *ClassParser) parseGeneric() (*ast.Generic, error) {
	if !p.check(lexer.TokenTilde) {
		return nil, p.error("expected '~' for generic type")
	}
	p.advance() // consume ~

	if !p.check(lexer.TokenID) {
		return nil, p.error("expected generic type name")
	}

	typeName := p.advance().Value
	generic := &ast.Generic{
		Name:      typeName,
		Arguments: make([]*ast.Generic, 0),
	}

	// Check for nested generics like List<String>
	if p.check(lexer.TokenOpenAngle) {
		p.advance() // consume <

		for !p.check(lexer.TokenCloseAngle) && !p.isAtEnd() {
			if p.check(lexer.TokenTilde) {
				nestedGeneric, err := p.parseGeneric()
				if err != nil {
					return nil, err
				}
				generic.Arguments = append(generic.Arguments, nestedGeneric)
			} else if p.check(lexer.TokenID) {
				// Simple type argument
				argType := p.advance().Value
				generic.Arguments = append(generic.Arguments, &ast.Generic{
					Name: argType,
				})
			}

			if p.check(lexer.TokenComma) {
				p.advance() // consume comma
			}
		}

		if !p.check(lexer.TokenCloseAngle) {
			return nil, p.error("expected '>' to close generic")
		}
		p.advance() // consume >
	}

	if !p.check(lexer.TokenTilde) {
		return nil, p.error("expected closing '~' for generic type")
	}
	p.advance() // consume closing ~

	return generic, nil
}