mermaid-go/pkg/parser/flowchart.go

// Package parser provides syntax analysis for Mermaid diagrams.
// Based on the grammar rules from flow.jison in mermaid.js
package parser

import (
	"fmt"
	_ "strconv"
	_ "strings"

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

// Parser implements recursive descent parsing for Mermaid flowcharts
// Following the grammar structure from flow.jison
type Parser struct {
	tokens  []lexer.Token
	current int
	flowDB  *FlowDB
}

// FlowDB manages the state during parsing, mirroring mermaid.js FlowDB
type FlowDB struct {
	vertexCounter      int
	vertices           map[string]*ast.FlowVertex
	edges              []*ast.FlowEdge
	classes            map[string]*ast.FlowClass
	subGraphs          []*ast.FlowSubGraph
	subGraphLookup     map[string]*ast.FlowSubGraph
	tooltips           map[string]string
	direction          string
	version            string
	defaultStyle       []string
	defaultInterpolate string
}

// NewFlowDB creates a new flow database
func NewFlowDB() *FlowDB {
	return &FlowDB{
		vertices:       make(map[string]*ast.FlowVertex),
		edges:          make([]*ast.FlowEdge, 0),
		classes:        make(map[string]*ast.FlowClass),
		subGraphs:      make([]*ast.FlowSubGraph, 0),
		subGraphLookup: make(map[string]*ast.FlowSubGraph),
		tooltips:       make(map[string]string),
		version:        "gen-2",
	}
}

// NewParser creates a new parser
func NewParser() *Parser {
	return &Parser{
		flowDB: NewFlowDB(),
	}
}

// NewFlowchartParser creates a new flowchart parser (alias for NewParser)
func NewFlowchartParser() *Parser {
	return NewParser()
}

// Parse parses the input string and returns a flowchart diagram
func (p *Parser) Parse(input string) (ast.Diagram, error) {
	// Tokenize
	l := lexer.NewLexer(input)
	tokens, err := l.Tokenize()
	if err != nil {
		return nil, fmt.Errorf("lexical analysis failed: %w", err)
	}

	// Filter out whitespace and comments
	p.tokens = lexer.FilterTokens(tokens)
	p.current = 0

	// Reset parser state
	p.flowDB = NewFlowDB()

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

	// Build final flowchart
	return p.buildFlowchart(), nil
}

// parseDocument implements the top-level grammar rule
// document: graphStatement | document graphStatement
func (p *Parser) parseDocument() error {
	for !p.isAtEnd() {
		if err := p.parseStatement(); err != nil {
			return err
		}
	}
	return nil
}

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

	token := p.peek()
	switch token.Type {
	case lexer.TokenGraph:
		return p.parseGraphStatement()
	case lexer.TokenSubgraph:
		return p.parseSubgraphStatement()
	case lexer.TokenClass:
		return p.parseClassStatement()
	case lexer.TokenClassDef:
		return p.parseClassDefStatement()
	case lexer.TokenStyle:
		return p.parseStyleStatement()
	case lexer.TokenLinkStyle:
		return p.parseLinkStyleStatement()
	case lexer.TokenClick:
		return p.parseClickStatement()
	case lexer.TokenNewline:
		p.advance() // Skip newlines
		return nil
	case lexer.TokenEOF:
		return nil
	default:
		// Try to parse as edge statement
		return p.parseEdgeStatement()
	}
}

// parseGraphStatement: GRAPH dir? (NL graphStatementList)?
func (p *Parser) parseGraphStatement() error {
	if !p.check(lexer.TokenGraph) {
		return p.error("expected 'graph'")
	}
	p.advance()

	// Optional direction
	if p.checkDirection() {
		dir := p.advance()
		p.flowDB.direction = dir.Value
	}

	// Optional newline
	if p.check(lexer.TokenNewline) {
		p.advance()
	}

	return nil
}

// parseSubgraphStatement handles subgraph definitions
func (p *Parser) parseSubgraphStatement() error {
	if !p.check(lexer.TokenSubgraph) {
		return p.error("expected 'subgraph'")
	}
	p.advance()

	// TODO: Implement subgraph parsing based on flow.jison
	// For now, skip to end
	for !p.check(lexer.TokenEnd) && !p.isAtEnd() {
		p.advance()
	}
	if p.check(lexer.TokenEnd) {
		p.advance()
	}

	return nil
}

// parseClassStatement handles class assignments
func (p *Parser) parseClassStatement() error {
	if !p.check(lexer.TokenClass) {
		return p.error("expected 'class'")
	}
	p.advance()

	// Skip implementation for now
	return p.skipToNextStatement()
}

// parseClassDefStatement handles class definitions
func (p *Parser) parseClassDefStatement() error {
	if !p.check(lexer.TokenClassDef) {
		return p.error("expected 'classDef'")
	}
	p.advance()

	// Skip implementation for now
	return p.skipToNextStatement()
}

// parseStyleStatement handles style definitions
func (p *Parser) parseStyleStatement() error {
	if !p.check(lexer.TokenStyle) {
		return p.error("expected 'style'")
	}
	p.advance()

	// Skip implementation for now
	return p.skipToNextStatement()
}

// parseLinkStyleStatement handles link style definitions
func (p *Parser) parseLinkStyleStatement() error {
	if !p.check(lexer.TokenLinkStyle) {
		return p.error("expected 'linkStyle'")
	}
	p.advance()

	// Skip implementation for now
	return p.skipToNextStatement()
}

// parseClickStatement handles click event definitions
func (p *Parser) parseClickStatement() error {
	if !p.check(lexer.TokenClick) {
		return p.error("expected 'click'")
	}
	p.advance()

	// Skip implementation for now
	return p.skipToNextStatement()
}

// parseEdgeStatement parses edge definitions
// This is the core parsing logic for flowchart connections
func (p *Parser) parseEdgeStatement() error {
	// Parse start vertex
	startVertex, err := p.parseVertex()
	if err != nil {
		return err
	}

	// Parse edge
	edge, err := p.parseEdge()
	if err != nil {
		return err
	}

	// Parse end vertex
	endVertex, err := p.parseVertex()
	if err != nil {
		return err
	}

	// Create edge in flowDB
	return p.addEdge(startVertex, endVertex, edge)
}

// parseVertex parses vertex definitions with shapes
// Examples: A[Text], B(Text), C{Text}, etc.
func (p *Parser) parseVertex() (*VertexInfo, error) {
	if !p.check(lexer.TokenID) {
		return nil, p.error("expected vertex identifier")
	}

	id := p.advance().Value
	vertex := &VertexInfo{ID: id}

	// Check for shape definition
	if p.checkShapeStart() {
		shape, text, err := p.parseShape()
		if err != nil {
			return nil, err
		}
		vertex.Shape = shape
		vertex.Text = text

		// Add vertex to flowDB
		p.addVertex(id, text, shape)
	}

	return vertex, nil
}

// VertexInfo holds parsed vertex information
type VertexInfo struct {
	ID    string
	Text  string
	Shape ast.FlowVertexTypeParam
}

// EdgeInfo holds parsed edge information
type EdgeInfo struct {
	Type   string
	Text   string
	Length int
	Stroke ast.FlowEdgeStroke
}

// parseShape parses shape definitions [text], (text), {text}, etc.
func (p *Parser) parseShape() (ast.FlowVertexTypeParam, string, error) {
	startToken := p.peek()
	var shape ast.FlowVertexTypeParam
	var endToken lexer.TokenType

	switch startToken.Type {
	case lexer.TokenOpenBracket:
		shape = ast.VertexTypeRect
		endToken = lexer.TokenCloseBracket
	case lexer.TokenOpenParen:
		if p.checkNext(lexer.TokenOpenParen) { // ((text))
			shape = ast.VertexTypeCircle
			p.advance() // skip first (
			p.advance() // skip second (
			endToken = lexer.TokenCloseParen
		} else { // (text)
			shape = ast.VertexTypeRound
			endToken = lexer.TokenCloseParen
		}
	case lexer.TokenOpenBrace:
		shape = ast.VertexTypeDiamond
		endToken = lexer.TokenCloseBrace
	case lexer.TokenOpenDoubleParen:
		shape = ast.VertexTypeCircle
		endToken = lexer.TokenCloseDoubleParen
	default:
		return "", "", p.error("expected shape delimiter")
	}

	if shape != ast.VertexTypeCircle || startToken.Type != lexer.TokenOpenDoubleParen {
		p.advance() // consume opening delimiter
	}

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

	if !p.check(endToken) {
		return "", "", p.error(fmt.Sprintf("expected closing delimiter"))
	}
	p.advance() // consume closing delimiter

	// Handle double paren closing
	if shape == ast.VertexTypeCircle && endToken == lexer.TokenCloseParen {
		if !p.check(lexer.TokenCloseParen) {
			return "", "", p.error("expected second closing parenthesis")
		}
		p.advance()
	}

	return shape, text, nil
}

// parseEdge parses edge definitions with arrows and labels
func (p *Parser) parseEdge() (*EdgeInfo, error) {
	edge := &EdgeInfo{
		Stroke: ast.StrokeNormal,
		Length: 1,
	}

	// Parse edge label if present (|text|)
	if p.check(lexer.TokenPipe) {
		p.advance() // consume |

		// Collect text until next |
		text := ""
		for !p.check(lexer.TokenPipe) && !p.isAtEnd() {
			text += p.advance().Value
		}

		if !p.check(lexer.TokenPipe) {
			return nil, p.error("expected closing pipe for edge label")
		}
		p.advance() // consume closing |

		edge.Text = text
	}

	// Parse arrow type
	if !p.checkArrow() {
		return nil, p.error("expected arrow")
	}

	arrow := p.advance()
	edge.Type, edge.Stroke = p.parseArrowType(arrow.Value)

	return edge, nil
}

// parseArrowType extracts type and stroke from arrow token
func (p *Parser) parseArrowType(arrow string) (string, ast.FlowEdgeStroke) {
	switch arrow {
	case "-->":
		return "arrow_point", ast.StrokeNormal
	case "-.->":
		return "arrow_point", ast.StrokeDotted
	case "==>":
		return "arrow_point", ast.StrokeThick
	case "--x":
		return "arrow_cross", ast.StrokeNormal
	case "--o":
		return "arrow_circle", ast.StrokeNormal
	case "---":
		return "arrow_open", ast.StrokeNormal
	default:
		return "arrow_point", ast.StrokeNormal
	}
}

// FlowDB manipulation methods (mirroring mermaid.js FlowDB)

// addVertex adds a vertex to the flow database
func (p *Parser) addVertex(id, text string, vertexType ast.FlowVertexTypeParam) {
	vertex := p.flowDB.vertices[id]
	if vertex == nil {
		vertex = &ast.FlowVertex{
			ID:        id,
			LabelType: "text",
			DomID:     fmt.Sprintf("flowchart-%s-%d", id, p.flowDB.vertexCounter),
			Styles:    make([]string, 0),
			Classes:   make([]string, 0),
		}
		p.flowDB.vertices[id] = vertex
		p.flowDB.vertexCounter++
	}

	if text != "" {
		vertex.Text = &text
	}
	if vertexType != "" {
		vertex.Type = &vertexType
	}
}

// addEdge adds an edge to the flow database
func (p *Parser) addEdge(start, end *VertexInfo, edge *EdgeInfo) error {
	// Ensure vertices exist
	p.addVertex(start.ID, start.Text, start.Shape)
	p.addVertex(end.ID, end.Text, end.Shape)

	// Create edge
	flowEdge := &ast.FlowEdge{
		Start:           start.ID,
		End:             end.ID,
		Text:            edge.Text,
		LabelType:       "text",
		Classes:         make([]string, 0),
		IsUserDefinedID: false,
	}

	if edge.Type != "" {
		flowEdge.Type = &edge.Type
	}
	if edge.Stroke != "" {
		flowEdge.Stroke = &edge.Stroke
	}

	// Generate edge ID
	edgeID := fmt.Sprintf("L-%s-%s-%d", start.ID, end.ID, len(p.flowDB.edges))
	flowEdge.ID = edgeID

	p.flowDB.edges = append(p.flowDB.edges, flowEdge)
	return nil
}

// buildFlowchart creates the final flowchart from flowDB state
func (p *Parser) buildFlowchart() *ast.Flowchart {
	flowchart := ast.NewFlowchart()
	flowchart.Direction = p.flowDB.direction
	flowchart.Vertices = p.flowDB.vertices
	flowchart.Edges = p.flowDB.edges
	flowchart.Classes = p.flowDB.classes
	flowchart.SubGraphs = p.flowDB.subGraphs
	flowchart.SubGraphLookup = p.flowDB.subGraphLookup
	flowchart.Tooltips = p.flowDB.tooltips
	flowchart.Version = p.flowDB.version
	return flowchart
}

// Helper methods

// check returns true if current token matches the given type
func (p *Parser) check(tokenType lexer.TokenType) bool {
	if p.isAtEnd() {
		return false
	}
	return p.peek().Type == tokenType
}

// checkNext returns true if next token matches the given type
func (p *Parser) checkNext(tokenType lexer.TokenType) bool {
	if p.current+1 >= len(p.tokens) {
		return false
	}
	return p.tokens[p.current+1].Type == tokenType
}

// checkDirection returns true if current token is a direction
func (p *Parser) checkDirection() bool {
	if p.isAtEnd() {
		return false
	}
	tokenType := p.peek().Type
	return tokenType == lexer.TokenTD || tokenType == lexer.TokenTB ||
		tokenType == lexer.TokenBT || tokenType == lexer.TokenRL ||
		tokenType == lexer.TokenLR
}

// checkShapeStart returns true if current token starts a shape
func (p *Parser) checkShapeStart() bool {
	if p.isAtEnd() {
		return false
	}
	tokenType := p.peek().Type
	return tokenType == lexer.TokenOpenBracket || tokenType == lexer.TokenOpenParen ||
		tokenType == lexer.TokenOpenBrace || tokenType == lexer.TokenOpenDoubleParen
}

// checkArrow returns true if current token is an arrow
func (p *Parser) checkArrow() bool {
	if p.isAtEnd() {
		return false
	}
	tokenType := p.peek().Type
	return tokenType == lexer.TokenArrowSolid || tokenType == lexer.TokenArrowDotted ||
		tokenType == lexer.TokenArrowThick || tokenType == lexer.TokenArrowOpen ||
		tokenType == lexer.TokenArrowPoint || tokenType == lexer.TokenArrowCross ||
		tokenType == lexer.TokenArrowCircle
}

// advance consumes and returns the current token
func (p *Parser) advance() lexer.Token {
	if !p.isAtEnd() {
		p.current++
	}
	return p.previous()
}

// isAtEnd returns true if we've reached the end of tokens
func (p *Parser) isAtEnd() bool {
	return p.current >= len(p.tokens) || p.peek().Type == lexer.TokenEOF
}

// peek returns the current token without advancing
func (p *Parser) peek() lexer.Token {
	if p.current >= len(p.tokens) {
		return lexer.Token{Type: lexer.TokenEOF}
	}
	return p.tokens[p.current]
}

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

// error creates a parsing error
func (p *Parser) 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)
}

// skipToNextStatement skips tokens until next statement
func (p *Parser) skipToNextStatement() error {
	for !p.isAtEnd() && !p.check(lexer.TokenNewline) {
		p.advance()
	}
	if p.check(lexer.TokenNewline) {
		p.advance()
	}
	return nil
}