Chapter 9: Advanced Parsing & Edge Cases: Nested Structures and Complexities

// src/ast/types.rs

use crate::utils::span::Span;
use std::collections::HashMap;

/// Represents the overall structure of a Mermaid diagram.
#[derive(Debug, PartialEq, Clone)]
pub struct Ast {
    pub diagram_type: DiagramType,
    pub items: Vec<GraphItem>,
    pub span: Span,
}

/// Represents the type of Mermaid diagram.
#[derive(Debug, PartialEq, Clone)]
pub enum DiagramType {
    Flowchart(FlowchartDirection),
    Sequence,
    Class,
    // Add other diagram types as supported
    Unknown, // For diagrams not yet fully parsed or invalid
}

/// Represents the direction of a flowchart.
#[derive(Debug, PartialEq, Clone)]
pub enum FlowchartDirection {
    TopToBottom, // TD
    BottomToTop, // BT
    RightToLeft, // RL
    LeftToRight, // LR
}

/// Represents a node in a diagram.
#[derive(Debug, PartialEq, Clone)]
pub struct AstNode {
    pub id: String,
    pub label: Option<NodeLabel>,
    pub shape: NodeShape,
    pub span: Span,
}

/// Represents a label for a node, which can be multiline.
#[derive(Debug, PartialEq, Clone)]
pub struct NodeLabel {
    pub text: String, // Can contain '\n' or '<br/>' for multiline
    pub span: Span,
}

/// Represents the shape of a node.
#[derive(Debug, PartialEq, Clone)]
pub enum NodeShape {
    Rectangle,        // A[Label]
    RoundedRectangle, // A(Label)
    Stadium,          // A([Label])
    Subroutine,       // A[[Label]]
    Cylinder,         // A[(Label)]
    Circle,           // A((Label))
    Asymmetric,       // A>Label]
    Rhombus,          // A{Label}
    Hexagon,          // A{{Label}}
    Parallelogram,    // A[/Label/]
    ParallelogramAlt, // A[\Label\]
    Trapezoid,        // A[/Label\]
    TrapezoidAlt,     // A[\Label/]
    Database,         // A{{Label}} (often rendered differently, but same syntax as hexagon)
    Unknown,          // For malformed or unsupported shapes
}

/// Represents an edge (connection) between two nodes.
#[derive(Debug, PartialEq, Clone)]
pub struct AstEdge {
    pub source: String,
    pub target: String,
    pub label: Option<EdgeLabel>,
    pub arrow_type: ArrowType,
    pub line_type: LineType,
    pub span: Span,
}

/// Represents a label for an edge, which can be multiline.
#[derive(Debug, PartialEq, Clone)]
pub struct EdgeLabel {
    pub text: String, // Can contain '\n' or '<br/>' for multiline
    pub span: Span,
}

/// Represents the type of arrow for an edge.
#[derive(Debug, PartialEq, Clone)]
pub enum ArrowType {
    None,    // ---
    Normal,  // -->
    Open,    // --o
    Cross,   // --x
    BidirectionalNormal, // <-->
    BidirectionalOpen,   // o--o
    BidirectionalCross,  // x--x
    Thick,   // ===>
    Dot,     // -.->
    DotThick, // =.=>
    Unknown,
}

/// Represents the type of line for an edge.
#[derive(Debug, PartialEq, Clone)]
pub enum LineType {
    Solid,
    Dashed,
    Thick,
    Invisible, // For implicit relationships, e.g., class diagrams
    Unknown,
}

/// Represents a subgraph, which can contain its own items.
#[derive(Debug, PartialEq, Clone)]
pub struct AstSubgraph {
    pub id: String,
    pub label: Option<NodeLabel>, // Subgraph can have an optional label
    pub items: Vec<GraphItem>, // Recursive definition!
    pub span: Span,
}

/// Represents a comment in the diagram.
#[derive(Debug, PartialEq, Clone)]
pub struct AstComment {
    pub text: String,
    pub span: Span,
}

/// An enum to hold any top-level item in the graph.
/// This allows for nodes, edges, subgraphs, and comments to coexist.
#[derive(Debug, PartialEq, Clone)]
pub enum GraphItem {
    Node(AstNode),
    Edge(AstEdge),
    Subgraph(AstSubgraph),
    Comment(AstComment),
    // Future: Directive, Style, etc.
}

impl AstNode {
    pub fn new(id: String, label: Option<NodeLabel>, shape: NodeShape, span: Span) -> Self {
        AstNode { id, label, shape, span }
    }
}

impl AstEdge {
    pub fn new(source: String, target: String, label: Option<EdgeLabel>, arrow_type: ArrowType, line_type: LineType, span: Span) -> Self {
        AstEdge { source, target, label, arrow_type, line_type, span }
    }
}

impl AstSubgraph {
    pub fn new(id: String, label: Option<NodeLabel>, items: Vec<GraphItem>, span: Span) -> Self {
        AstSubgraph { id, label, items, span }
    }
}

impl AstComment {
    pub fn new(text: String, span: Span) -> Self {
        AstComment { text, span }
    }
}

impl NodeLabel {
    pub fn new(text: String, span: Span) -> Self {
        NodeLabel { text, span }
    }
}

impl EdgeLabel {
    pub fn new(text: String, span: Span) -> Self {
        EdgeLabel { text, span }
    }
}

// src/lexer/token.rs

use crate::utils::span::Span;

/// Represents a token produced by the lexer.
#[derive(Debug, PartialEq, Clone)]
pub struct Token {
    pub kind: TokenKind,
    pub lexeme: String,
    pub span: Span,
}

/// Enum representing the different kinds of tokens.
#[derive(Debug, PartialEq, Clone)]
pub enum TokenKind {
    // Keywords
    Graph,
    Flowchart,
    Sequence,
    Class,
    Subgraph, // NEW
    End,      // NEW

    // Directions
    TD,
    BT,
    RL,
    LR,

    // Identifiers and Literals
    Identifier,
    StringLiteral, // Quoted string
    NumberLiteral, // For future use, e.g., sequence numbers

    // Punctuation and Operators
    OpenParen,    // (
    CloseParen,   // )
    OpenBracket,  // [
    CloseBracket, // ]
    OpenBrace,    // {
    CloseBrace,   // }
    OpenAngle,    // <
    CloseAngle,   // >
    Colon,        // :
    Equals,       // =
    Comma,        // ,
    Semicolon,    // ;
    Newline,      // \n
    Whitespace,   // Space, tab (usually ignored, but sometimes relevant for formatting)

    // Arrows and Lines
    ArrowNormal,          // -->
    ArrowOpen,            // --o
    ArrowCross,           // --x
    ArrowBidirectionalNormal, // <-->
    ArrowBidirectionalOpen,   // o--o
    ArrowBidirectionalCross,  // x--x
    ArrowThick,           // ===>
    ArrowDot,             // -.->
    ArrowDotThick,        // =.=>

    LineSolid,            // ---
    LineDashed,           // --- (when no arrow)
    LineThick,            // ===
    LineDot,              // -.-

    // Special
    Comment,      // NEW: %% comment text
    Unknown,      // For unrecognised characters/sequences
    Eof,          // End of file
}

impl Token {
    pub fn new(kind: TokenKind, lexeme: String, span: Span) -> Self {
        Token { kind, lexeme, span }
    }
}

// src/lexer/mod.rs

use crate::lexer::token::{Token, TokenKind};
use crate::diagnostics::emitter::DiagnosticEmitter;
use crate::utils::span::Span;

pub struct Lexer<'a> {
    source: &'a str,
    chars: std::iter::Peekable<std::str::Chars<'a>>,
    start: usize,
    current: usize,
    line: usize,
    col: usize,
    emitter: &'a mut DiagnosticEmitter,
}

impl<'a> Lexer<'a> {
    pub fn new(source: &'a str, emitter: &'a mut DiagnosticEmitter) -> Self {
        Lexer {
            source,
            chars: source.chars().peekable(),
            start: 0,
            current: 0,
            line: 1,
            col: 1,
            emitter,
        }
    }

    pub fn next_token(&mut self) -> Token {
        self.skip_whitespace_and_comments(); // Add comment skipping here

        self.start = self.current;

        if self.is_at_end() {
            return self.make_token(TokenKind::Eof);
        }

        let c = self.advance().unwrap(); // We know it's not at end

        match c {
            '(' => self.make_token(TokenKind::OpenParen),
            ')' => self.make_token(TokenKind::CloseParen),
            '[' => self.make_token(TokenKind::OpenBracket),
            ']' => self.make_token(TokenKind::CloseBracket),
            '{' => self.make_token(TokenKind::OpenBrace),
            '}' => self.make_token(TokenKind::CloseBrace),
            '<' => {
                if self.peek() == Some('-') {
                    self.advance();
                    if self.peek() == Some('-') {
                        self.advance();
                        if self.peek() == Some('>') {
                            self.advance();
                            return self.make_token(TokenKind::ArrowBidirectionalNormal); // <-->
                        } else if self.peek() == Some('o') {
                            self.advance();
                            return self.make_token(TokenKind::ArrowBidirectionalOpen); // <--o (not official, but common variant)
                        } else if self.peek() == Some('x') {
                            self.advance();
                            return self.make_token(TokenKind::ArrowBidirectionalCross); // <--x (not official, but common variant)
                        }
                    }
                }
                self.make_token(TokenKind::OpenAngle) // Single <
            }
            '>' => self.make_token(TokenKind::CloseAngle),
            ':' => self.make_token(TokenKind::Colon),
            '=' => {
                if self.peek() == Some('=') {
                    self.advance();
                    if self.peek() == Some('=') {
                        self.advance();
                        if self.peek() == Some('>') {
                            self.advance();
                            return self.make_token(TokenKind::ArrowThick); // ===>
                        }
                        return self.make_token(TokenKind::LineThick); // ===
                    } else if self.peek() == Some('.') {
                        self.advance();
                        if self.peek() == Some('>') {
                            self.advance();
                            return self.make_token(TokenKind::ArrowDotThick); // =.=>
                        }
                    }
                }
                self.make_token(TokenKind::Equals) // Single =
            }
            ',' => self.make_token(TokenKind::Comma),
            ';' => self.make_token(TokenKind::Semicolon),
            '/' => {
                if self.peek() == Some('/') {
                    // This is a line comment in some languages, but Mermaid uses `%%`
                    // For now, treat `//` as an identifier start if not part of a specific pattern.
                    // Or we could define it as an Unknown token if it's truly invalid.
                    // For now, let's assume it might be part of a node ID if it's not a known symbol.
                    self.current -= 1; // Backtrack
                    self.start = self.current;
                    return self.identifier_or_keyword();
                }
                self.make_token(TokenKind::Unknown) // Single / might be part of shape, handle later
            }
            '%' => { // NEW: Handle comments
                if self.peek() == Some('%') {
                    self.advance(); // Consume the second '%'
                    self.skip_to_end_of_line();
                    let comment_text = &self.source[self.start + 2..self.current].trim_end(); // Exclude '%%' and trim
                    return self.make_token_with_lexeme(TokenKind::Comment, comment_text.to_string());
                }
                self.make_token(TokenKind::Unknown) // Single %
            }
            '-' => {
                if self.peek() == Some('.') {
                    self.advance();
                    if self.peek() == Some('>') {
                        self.advance();
                        return self.make_token(TokenKind::ArrowDot); // -.->
                    }
                }
                if self.peek() == Some('-') {
                    self.advance();
                    if self.peek() == Some('>') {
                        self.advance();
                        return self.make_token(TokenKind::ArrowNormal); // -->
                    } else if self.peek() == Some('o') {
                        self.advance();
                        return self.make_token(TokenKind::ArrowOpen); // --o
                    } else if self.peek() == Some('x') {
                        self.advance();
                        return self.make_token(TokenKind::ArrowCross); // --x
                    }
                    return self.make_token(TokenKind::LineSolid); // ---
                }
                self.make_token(TokenKind::Unknown) // Single -
            }
            '"' => self.string_literal(),
            '\n' => {
                self.line += 1;
                self.col = 1;
                self.make_token(TokenKind::Newline)
            }
            c if c.is_ascii_digit() => self.number_literal(),
            c if c.is_ascii_alphabetic() || c == '_' => self.identifier_or_keyword(),
            _ => {
                self.emitter.error(
                    format!("Unexpected character: '{}'", c),
                    self.current - 1,
                    self.current,
                    "E0001",
                    "Remove or correct the unexpected character.",
                );
                self.make_token(TokenKind::Unknown)
            }
        }
    }

    fn skip_whitespace_and_comments(&mut self) {
        loop {
            match self.peek() {
                Some(' ') | Some('\r') | Some('\t') => {
                    self.advance();
                    self.start = self.current; // Reset start to ignore skipped whitespace
                }
                Some('\n') => {
                    // Newlines are significant for parsing, so we don't skip them here.
                    // They are tokenized as `Newline`.
                    return;
                }
                Some('%') if self.peek_next() == Some('%') => {
                    self.advance(); // Consume first %
                    self.advance(); // Consume second %
                    self.skip_to_end_of_line();
                    self.start = self.current; // Reset start to ignore skipped comment
                }
                _ => return,
            }
        }
    }

    fn skip_to_end_of_line(&mut self) {
        while let Some(c) = self.peek() {
            if c == '\n' {
                break;
            }
            self.advance();
        }
    }

    fn string_literal(&mut self) -> Token {
        let mut content = String::new();
        let mut escaped = false;

        while let Some(c) = self.peek() {
            if c == '"' && !escaped {
                break; // End of string
            }
            if c == '\n' {
                // Mermaid allows multiline labels, but usually with <br/> or explicit newlines.
                // If a raw newline is encountered within a quoted string, it's part of the string.
                content.push(c);
                self.advance();
                self.line += 1;
                self.col = 1;
                continue;
            }

            if c == '\\' {
                self.advance(); // Consume the backslash
                match self.peek() {
                    Some('"') => {
                        content.push('"'); // Add the escaped quote
                        escaped = false; // Reset escaped state after consuming
                    }
                    Some('\\') => {
                        content.push('\\'); // Add the escaped backslash
                        escaped = false;
                    }
                    _ => {
                        // Invalid escape sequence, just add the backslash and the next char
                        self.emitter.warning(
                            format!("Invalid escape sequence '\\{}'", self.peek().unwrap_or(' ')),
                            self.current - 1,
                            self.current + 1,
                            "W0001",
                            "Consider removing the backslash or escaping correctly.",
                        );
                        content.push('\\');
                        if let Some(next_c) = self.peek() {
                            content.push(next_c);
                        }
                    }
                }
            } else {
                content.push(c);
            }
            self.advance();
            escaped = false; // Reset for next character
        }

        if self.is_at_end() || self.peek() != Some('"') {
            self.emitter.error(
                "Unterminated string literal.",
                self.start,
                self.current,
                "E0002",
                "Add a closing '\"' to terminate the string.",
            );
            return self.make_token_with_lexeme(TokenKind::StringLiteral, content); // Return partial string
        }

        self.advance(); // Consume the closing '"'
        self.make_token_with_lexeme(TokenKind::StringLiteral, content)
    }

    fn number_literal(&mut self) -> Token {
        while self.peek().map_or(false, |c| c.is_ascii_digit()) {
            self.advance();
        }
        // Mermaid mostly uses numbers for sequence diagram line numbers, etc.
        // For flowcharts, node IDs are alphanumeric.
        self.make_token(TokenKind::NumberLiteral)
    }

    fn identifier_or_keyword(&mut self) -> Token {
        while self.peek().map_or(false, |c| c.is_ascii_alphanumeric() || c == '_') {
            self.advance();
        }

        let text = &self.source[self.start..self.current];
        let kind = match text {
            "graph" => TokenKind::Graph,
            "flowchart" => TokenKind::Flowchart,
            "sequence" => TokenKind::Sequence,
            "classDiagram" => TokenKind::Class, // Assuming this as a keyword
            "TD" => TokenKind::TD,
            "BT" => TokenKind::BT,
            "RL" => TokenKind::RL,
            "LR" => TokenKind::LR,
            "subgraph" => TokenKind::Subgraph, // NEW
            "end" => TokenKind::End,           // NEW
            _ => TokenKind::Identifier,
        };
        self.make_token(kind)
    }

    fn advance(&mut self) -> Option<char> {
        self.current += 1;
        self.col += 1;
        self.chars.next()
    }

    fn peek(&self) -> Option<char> {
        self.chars.peek().copied()
    }

    fn peek_next(&self) -> Option<char> {
        let mut temp_chars = self.chars.clone();
        temp_chars.next(); // Consume current
        temp_chars.next().copied() // Peek next
    }

    fn is_at_end(&self) -> bool {
        self.current >= self.source.len()
    }

    fn make_token(&self, kind: TokenKind) -> Token {
        let lexeme = self.source[self.start..self.current].to_string();
        let span = Span::new(self.start, self.current, self.line, self.col - lexeme.len());
        Token::new(kind, lexeme, span)
    }

    fn make_token_with_lexeme(&self, kind: TokenKind, lexeme: String) -> Token {
        let span = Span::new(self.start, self.current, self.line, self.col - lexeme.len());
        Token::new(kind, lexeme, span)
    }
}

// src/parser/mod.rs

use crate::lexer::Lexer;
use crate::lexer::token::{Token, TokenKind};
use crate::ast::types::{Ast, DiagramType, FlowchartDirection, AstNode, NodeLabel, NodeShape, AstEdge, EdgeLabel, ArrowType, LineType, GraphItem, AstSubgraph, AstComment};
use crate::diagnostics::emitter::DiagnosticEmitter;
use crate::utils::span::Span;

pub struct Parser<'a> {
    lexer: Lexer<'a>,
    current_token: Token,
    previous_token: Token,
    emitter: &'a mut DiagnosticEmitter,
    errors_encountered: u32,
    // Stack to keep track of open subgraphs or other scopes
    scope_stack: Vec<String>,
}

impl<'a> Parser<'a> {
    pub fn new(lexer: Lexer<'a>, emitter: &'a mut DiagnosticEmitter) -> Self {
        let dummy_token = Token::new(TokenKind::Eof, "".to_string(), Span::new(0, 0, 0, 0));
        let mut parser = Parser {
            lexer,
            current_token: dummy_token.clone(),
            previous_token: dummy_token,
            emitter,
            errors_encountered: 0,
            scope_stack: Vec::new(),
        };
        parser.advance(); // Initialize current_token
        parser
    }

    pub fn parse(&mut self) -> Option<Ast> {
        let start_span = self.current_token.span;

        // 1. Parse diagram type
        let diagram_type = self.parse_diagram_type()?;

        // 2. Parse graph items (nodes, edges, subgraphs, comments)
        let mut items = Vec::new();
        while !self.is_at_end() {
            if let Some(item) = self.parse_graph_item() {
                items.push(item);
            } else {
                // If parse_graph_item returns None, it means an error occurred
                // and we should try to recover or stop.
                // For now, let's just break to avoid infinite loops on unrecoverable errors.
                // A more robust recovery would try to synchronize to a known token (e.g., 'graph', 'subgraph', 'node_id')
                break;
            }
        }

        if self.errors_encountered > 0 {
            None // Don't return AST if there were parsing errors
        } else {
            Some(Ast {
                diagram_type,
                items,
                span: start_span.merge(&self.previous_token.span),
            })
        }
    }

    fn parse_diagram_type(&mut self) -> Option<DiagramType> {
        let diagram_type_span_start = self.current_token.span.start;
        let diagram_type_lexeme = self.current_token.lexeme.clone();

        let kind = self.current_token.kind.clone();
        match kind {
            TokenKind::Graph | TokenKind::Flowchart => {
                self.expect_token_kind_and_advance(&[TokenKind::Graph, TokenKind::Flowchart], "Expected 'graph' or 'flowchart' declaration.")?;
                // Parse direction if present
                if let Some(direction) = self.parse_flowchart_direction() {
                    return Some(DiagramType::Flowchart(direction));
                }
                // Default to TD if no direction specified
                Some(DiagramType::Flowchart(FlowchartDirection::TopToBottom))
            }
            TokenKind::Sequence => {
                self.expect_token_kind_and_advance(&[TokenKind::Sequence], "Expected 'sequence' declaration.")?;
                Some(DiagramType::Sequence)
            }
            TokenKind::Class => { // Assuming 'classDiagram' is tokenized as Class
                self.expect_token_kind_and_advance(&[TokenKind::Class], "Expected 'classDiagram' declaration.")?;
                Some(DiagramType::Class)
            }
            _ => {
                self.emit_error(
                    format!("Expected a diagram type declaration (e.g., 'graph', 'flowchart', 'sequence', 'classDiagram'), found '{}'.", diagram_type_lexeme),
                    diagram_type_span_start,
                    self.current_token.span.end,
                    "E0101",
                    "Start your Mermaid diagram with a valid declaration.",
                );
                None
            }
        }
    }

    fn parse_flowchart_direction(&mut self) -> Option<FlowchartDirection> {
        match self.current_token.kind {
            TokenKind::TD => { self.advance(); Some(FlowchartDirection::TopToBottom) }
            TokenKind::BT => { self.advance(); Some(FlowchartDirection::BottomToTop) }
            TokenKind::RL => { self.advance(); Some(FlowchartDirection::RightToLeft) }
            TokenKind::LR => { self.advance(); Some(FlowchartDirection::LeftToRight) }
            _ => None, // No direction specified, or it's another token
        }
    }

    // NEW: Centralized parsing of graph items
    fn parse_graph_item(&mut self) -> Option<GraphItem> {
        match self.current_token.kind {
            TokenKind::Comment => { // NEW: Handle comments
                let comment_text = self.current_token.lexeme.clone();
                let span = self.current_token.span;
                self.advance(); // Consume the comment token
                Some(GraphItem::Comment(AstComment::new(comment_text, span)))
            }
            TokenKind::Subgraph => { // NEW: Handle subgraphs
                self.parse_subgraph().map(GraphItem::Subgraph)
            }
            TokenKind::Identifier => {
                // Could be a node declaration or the start of an edge
                // We need to peek to decide.
                let next_token_kind = self.lexer.peek_kind(); // Assuming lexer has peek_kind
                if next_token_kind.map_or(false, |k| k.is_arrow_or_line()) {
                    self.parse_edge().map(GraphItem::Edge)
                } else {
                    self.parse_node().map(GraphItem::Node)
                }
            }
            TokenKind::Newline => {
                self.advance(); // Consume newline, it's usually just separating items
                None // Newlines themselves are not graph items
            }
            TokenKind::Eof => None, // Reached end, no more items
            _ => {
                // Attempt error recovery: skip tokens until a known item start or EOF
                self.emit_error(
                    format!("Unexpected token '{}'. Expected node, edge, subgraph, or comment.", self.current_token.lexeme),
                    self.current_token.span.start,
                    self.current_token.span.end,
                    "E0102",
                    "Remove this token or correct the syntax to form a valid diagram element.",
                );
                self.synchronize();
                None
            }
        }
    }

    // NEW: Recursive subgraph parsing
    fn parse_subgraph(&mut self) -> Option<AstSubgraph> {
        let start_span = self.current_token.span;
        self.expect_token_kind_and_advance(&[TokenKind::Subgraph], "Expected 'subgraph' keyword.")?;

        let id_token = self.expect_token_kind_and_advance(&[TokenKind::Identifier], "Expected subgraph ID after 'subgraph' keyword.")?;
        let id = id_token.lexeme.clone();

        let label = if self.current_token.kind == TokenKind::OpenBracket {
            self.parse_node_label() // Subgraph labels use node label syntax
        } else {
            None
        };

        // Push subgraph ID onto scope stack
        self.scope_stack.push(id.clone());

        let mut items = Vec::new();
        // Parse items within the subgraph until 'end' or EOF
        while self.current_token.kind != TokenKind::End && !self.is_at_end() {
            if let Some(item) = self.parse_graph_item() {
                items.push(item);
            } else {
                // Error in subgraph item, try to recover
                if self.current_token.kind != TokenKind::End && !self.is_at_end() {
                    self.synchronize(); // Try to skip to a valid state
                }
            }
        }

        let end_span = self.current_token.span;
        self.expect_token_kind_and_advance(&[TokenKind::End], &format!("Expected 'end' to close subgraph '{}'.", id))?;

        // Pop subgraph ID from scope stack
        self.scope_stack.pop();

        Some(AstSubgraph::new(id, label, items, start_span.merge(&end_span)))
    }

    fn parse_node(&mut self) -> Option<AstNode> {
        let start_span = self.current_token.span;
        let id_token = self.expect_token_kind_and_advance(&[TokenKind::Identifier], "Expected node ID.")?;
        let id = id_token.lexeme.clone();

        let mut label: Option<NodeLabel> = None;
        let mut shape: NodeShape = NodeShape::Rectangle; // Default shape

        // Check for node label and shape
        match self.current_token.kind {
            TokenKind::OpenBracket => { // [Label] or [[Label]] or [/Label\] etc.
                label = self.parse_node_label(); // This handles different bracket types
                shape = self.determine_node_shape(&id_token.lexeme, &self.previous_token.kind, &self.current_token.kind);
            }
            TokenKind::OpenParen => { // (Label) or ((Label)) etc.
                label = self.parse_node_label();
                shape = self.determine_node_shape(&id_token.lexeme, &self.previous_token.kind, &self.current_token.kind);
            }
            TokenKind::OpenBrace => { // {Label} or {{Label}}
                label = self.parse_node_label();
                shape = self.determine_node_shape(&id_token.lexeme, &self.previous_token.kind, &self.current_token.kind);
            }
            TokenKind::OpenAngle => { // >Label]
                label = self.parse_node_label();
                shape = self.determine_node_shape(&id_token.lexeme, &self.previous_token.kind, &self.current_token.kind);
            }
            _ => { /* No label or shape, just a simple node ID */ }
        }

        Some(AstNode::new(id, label, shape, start_span.merge(&self.previous_token.span)))
    }

    fn parse_node_label(&mut self) -> Option<NodeLabel> {
        let start_span = self.current_token.span;
        let open_kind = self.current_token.kind.clone();
        let close_kind = match open_kind {
            TokenKind::OpenParen => TokenKind::CloseParen,
            TokenKind::OpenBracket => TokenKind::CloseBracket,
            TokenKind::OpenBrace => TokenKind::CloseBrace,
            TokenKind::OpenAngle => TokenKind::CloseBracket, // For asymmetric >Label]
            _ => {
                self.emit_error(
                    format!("Expected opening bracket for node label, found '{}'.", self.current_token.lexeme),
                    start_span.start,
                    start_span.end,
                    "E0103",
                    "Use '(', '[', '{', or '>' to open a node label.",
                );
                return None;
            }
        };

        self.advance(); // Consume opening bracket

        let label_text_token = self.expect_token_kind_and_advance(&[TokenKind::Identifier, TokenKind::StringLiteral], "Expected label text or quoted string.")?;
        let label_text = label_text_token.lexeme.clone();

        self.expect_token_kind_and_advance(&[close_kind], &format!("Expected closing bracket '{}' for node label.", self.token_kind_to_char(close_kind)))?;

        Some(NodeLabel::new(label_text, start_span.merge(&self.previous_token.span)))
    }

    fn parse_edge(&mut self) -> Option<AstEdge> {
        let start_span = self.current_token.span;
        let source_token = self.expect_token_kind_and_advance(&[TokenKind::Identifier], "Expected source node ID for an edge.")?;
        let source = source_token.lexeme.clone();

        let (arrow_type, line_type) = self.parse_arrow_and_line_type()?;
        let edge_label = self.parse_edge_label();

        let target_token = self.expect_token_kind_and_advance(&[TokenKind::Identifier], "Expected target node ID after arrow.")?;
        let target = target_token.lexeme.clone();

        Some(AstEdge::new(source, target, edge_label, arrow_type, line_type, start_span.merge(&target_token.span)))
    }

    fn parse_arrow_and_line_type(&mut self) -> Option<(ArrowType, LineType)> {
        let current_kind = self.current_token.kind.clone();
        let span = self.current_token.span;

        let (arrow_type, line_type) = match current_kind {
            TokenKind::ArrowNormal => (ArrowType::Normal, LineType::Solid),
            TokenKind::ArrowOpen => (ArrowType::Open, LineType::Solid),
            TokenKind::ArrowCross => (ArrowType::Cross, LineType::Solid),
            TokenKind::ArrowBidirectionalNormal => (ArrowType::BidirectionalNormal, LineType::Solid),
            TokenKind::ArrowBidirectionalOpen => (ArrowType::BidirectionalOpen, LineType::Solid),
            TokenKind::ArrowBidirectionalCross => (ArrowType::BidirectionalCross, LineType::Solid),
            TokenKind::ArrowThick => (ArrowType::Thick, LineType::Thick),
            TokenKind::ArrowDot => (ArrowType::Dot, LineType::Dashed),
            TokenKind::ArrowDotThick => (ArrowType::DotThick, LineType::Thick),
            TokenKind::LineSolid => (ArrowType::None, LineType::Solid),
            TokenKind::LineDashed => (ArrowType::None, LineType::Dashed), // Mermaid often uses --- for both solid line and dashed
            TokenKind::LineThick => (ArrowType::None, LineType::Thick),
            _ => {
                self.emit_error(
                    format!("Expected an arrow or line type, found '{}'.", self.current_token.lexeme),
                    span.start,
                    span.end,
                    "E0104",
                    "Use a valid Mermaid arrow or line syntax (e.g., '-->', '---', '===>').",
                );
                return None;
            }
        };
        self.advance(); // Consume the arrow/line token
        Some((arrow_type, line_type))
    }

    fn parse_edge_label(&mut self) -> Option<EdgeLabel> {
        if self.current_token.kind == TokenKind::Pipe { // |Label| syntax for edges
            let start_span = self.current_token.span;
            self.advance(); // Consume '|'

            let label_text_token = self.expect_token_kind_and_advance(&[TokenKind::Identifier, TokenKind::StringLiteral], "Expected label text for edge.")?;
            let label_text = label_text_token.lexeme.clone();

            self.expect_token_kind_and_advance(&[TokenKind::Pipe], "Expected closing '|' for edge label.")?;
            Some(EdgeLabel::new(label_text, start_span.merge(&self.previous_token.span)))
        } else {
            None
        }
    }


    // Helper to determine node shape based on surrounding tokens
    fn determine_node_shape(&self, _id: &str, open_bracket_kind: &TokenKind, close_bracket_kind: &TokenKind) -> NodeShape {
        match (open_bracket_kind, close_bracket_kind) {
            (TokenKind::OpenBracket, TokenKind::CloseBracket) => {
                // Peek to see if it's [[ ]] or [ ]
                if self.peek_prev_kind() == Some(TokenKind::OpenBracket) && self.peek_next_kind() == Some(TokenKind::CloseBracket) {
                    NodeShape::Subroutine // [[Label]]
                } else {
                    NodeShape::Rectangle // [Label]
                }
            },
            (TokenKind::OpenParen, TokenKind::CloseParen) => {
                // Peek to see if it's (( )) or ( )
                if self.peek_prev_kind() == Some(TokenKind::OpenParen) && self.peek_next_kind() == Some(TokenKind::CloseParen) {
                    NodeShape::Circle // ((Label))
                } else if self.peek_prev_kind() == Some(TokenKind::OpenBracket) && self.peek_next_kind() == Some(TokenKind::CloseParen) {
                    NodeShape::Cylinder // [(Label)]
                } else {
                    NodeShape::RoundedRectangle // (Label)
                }
            },
            (TokenKind::OpenBrace, TokenKind::CloseBrace) => {
                // Peek to see if it's {{ }} or { }
                if self.peek_prev_kind() == Some(TokenKind::OpenBrace) && self.peek_next_kind() == Some(TokenKind::CloseBrace) {
                    NodeShape::Hexagon // {{Label}} or Database
                } else {
                    NodeShape::Rhombus // {Label}
                }
            },
            (TokenKind::OpenAngle, TokenKind::CloseBracket) => NodeShape::Asymmetric, // >Label]
            // Add other complex shapes like parallelogram, trapezoid, etc.
            // This requires more sophisticated lookahead or token patterns.
            // For now, these are basic.
            _ => NodeShape::Unknown, // Fallback for unhandled combinations
        }
    }

    // --- Utility Methods ---

    fn advance(&mut self) {
        self.previous_token = self.current_token.clone();
        self.current_token = self.lexer.next_token();
        // Skip `Newline` tokens unless they are explicitly part of the grammar we're parsing
        // For general item parsing, newlines often just separate statements.
        // We'll keep them in the lexer for accurate span reporting, but skip them here
        // if they're not syntactically meaningful at the current point.
        // This loop ensures we always land on a non-newline token.
        while self.current_token.kind == TokenKind::Newline {
            self.previous_token = self.current_token.clone();
            self.current_token = self.lexer.next_token();
        }
    }

    fn is_at_end(&self) -> bool {
        self.current_token.kind == TokenKind::Eof
    }

    fn expect_token_kind_and_advance(&mut self, expected_kinds: &[TokenKind], message: &str) -> Option<Token> {
        if expected_kinds.contains(&self.current_token.kind) {
            let token = self.current_token.clone();
            self.advance();
            Some(token)
        } else {
            self.emit_error(
                format!("{} Found '{}'.", message, self.current_token.lexeme),
                self.current_token.span.start,
                self.current_token.span.end,
                "E0105",
                &format!("Ensure the correct syntax is used. Expected one of: {}.", expected_kinds.iter().map(|k| format!("{:?}", k)).collect::<Vec<_>>().join(", ")),
            );
            self.errors_encountered += 1;
            self.synchronize(); // Attempt error recovery
            None
        }
    }

    fn emit_error(&mut self, message: String, start: usize, end: usize, code: &str, help: &str) {
        self.emitter.error(message, start, end, code, help);
        self.errors_encountered += 1;
    }

    /// Error recovery mechanism.
    /// Skips tokens until it finds a token that likely starts a new statement or graph item,
    /// or a token that might help synchronize the parser.
    fn synchronize(&mut self) {
        // Skip until we find a token that typically starts a new statement or block.
        // This prevents cascading errors and allows parsing to continue.
        while !self.is_at_end() {
            match self.current_token.kind {
                TokenKind::Graph | TokenKind::Flowchart | TokenKind::Sequence | TokenKind::Class |
                TokenKind::Subgraph | TokenKind::End | // Subgraph boundaries are good sync points
                TokenKind::Identifier | // Node IDs often start new statements
                TokenKind::Newline => { // Newlines can also indicate a new statement
                    return;
                }
                _ => {
                    self.advance();
                }
            }
        }
    }

    // Helper to get the kind of the token that was just processed before current_token.
    // Useful for look-behind in shape determination.
    fn peek_prev_kind(&self) -> Option<TokenKind> {
        // This is a bit tricky with `advance` skipping newlines.
        // For simple cases like `((Label))`, it would be `previous_token`'s kind.
        // For more robust look-behind, the lexer might need to expose more.
        // For now, let's assume `previous_token` is sufficient for immediate look-behind.
        Some(self.previous_token.kind.clone())
    }

    // Helper to get the kind of the token *after* the current_token.
    // This requires the lexer to support peeking beyond the current token.
    fn peek_next_kind(&self) -> Option<TokenKind> {
        self.lexer.peek_kind() // Assuming lexer has this method
    }

    fn token_kind_to_char(&self, kind: TokenKind) -> char {
        match kind {
            TokenKind::OpenParen => '(',
            TokenKind::CloseParen => ')',
            TokenKind::OpenBracket => '[',
            TokenKind::CloseBracket => ']',
            TokenKind::OpenBrace => '{',
            TokenKind::CloseBrace => '}',
            TokenKind::OpenAngle => '<',
            TokenKind::CloseAngle => '>',
            TokenKind::Pipe => '|',
            _ => '?', // Fallback
        }
    }
}

// Add a method to Lexer to peek the kind of the *next* token without consuming.
// This is essential for the parser's lookahead logic.
impl<'a> Lexer<'a> {
    pub fn peek_kind(&mut self) -> Option<TokenKind> {
        let current_state = (self.start, self.current, self.line, self.col, self.chars.clone());
        let next_token = self.next_token(); // This advances the lexer state
        let kind = next_token.kind.clone();
        // Restore lexer state
        (self.start, self.current, self.line, self.col, self.chars) = current_state;
        if kind == TokenKind::Eof { None } else { Some(kind) }
    }
}

// Extend TokenKind to easily check if it's an arrow or line
impl TokenKind {
    pub fn is_arrow_or_line(&self) -> bool {
        matches!(self,
            TokenKind::ArrowNormal | TokenKind::ArrowOpen | TokenKind::ArrowCross |
            TokenKind::ArrowBidirectionalNormal | TokenKind::ArrowBidirectionalOpen |
            TokenKind::ArrowBidirectionalCross | TokenKind::ArrowThick |
            TokenKind::ArrowDot | TokenKind::ArrowDotThick |
            TokenKind::LineSolid | TokenKind::LineDashed | TokenKind::LineThick
        )
    }
}