Chapter 3: Crafting the Parser: Building the Abstract Syntax Tree (AST)

// src/parser/mod.rs
pub mod ast;
pub mod parser;

pub use self::parser::Parser;
pub use self::ast::{
    Diagram, DiagramType, Statement, Node, Edge, Subgraph, ArrowType, LineType, NodeShape,
    EdgeDirection,
};

// src/lib.rs
pub mod lexer;
pub mod diagnostics;
pub mod parser; // Add this line

// Re-export key types for easier access
pub use lexer::{Lexer, Token, TokenType, Location};
pub use diagnostics::{Diagnostic, DiagnosticLevel, DiagnosticCode, DiagnosticEmitter};
pub use parser::{
    Parser, Diagram, DiagramType, Statement, Node, Edge, Subgraph, ArrowType, LineType, NodeShape,
    EdgeDirection,
};

// src/parser/ast.rs
use crate::{Location, TokenType}; // Import Location and TokenType from root crate

/// Represents the top-level structure of a Mermaid diagram.
#[derive(Debug, PartialEq, Eq, Clone)]
pub struct Diagram {
    pub diagram_type: DiagramType,
    pub statements: Vec<Statement>,
    pub location: Location, // Location of the 'graph' or 'flowchart' keyword
}

/// Enumerates the supported types of Mermaid diagrams.
#[derive(Debug, PartialEq, Eq, Clone)]
pub enum DiagramType {
    Flowchart(Option<FlowchartDirection>), // e.g., graph TD, flowchart LR
    Sequence,
    Class,
    // Add other diagram types as needed (e.g., Gantt, State, Git, etc.)
}

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

/// Represents a single statement within a Mermaid diagram.
#[derive(Debug, PartialEq, Eq, Clone)]
pub enum Statement {
    Node(Node),
    Edge(Edge),
    Subgraph(Subgraph),
    Directive(Directive), // e.g., %%{init: ...}
    Comment(String, Location), // e.g., %% This is a comment
    EmptyLine(Location), // For tracking formatting, useful for formatter
    // Potentially add other statement types like Class members, Sequence messages, etc.
}

/// Represents a node in the diagram.
#[derive(Debug, PartialEq, Eq, Clone)]
pub struct Node {
    pub id: String,
    pub label: Option<String>, // Optional label for nodes like A[Label]
    pub shape: NodeShape,
    pub location: Location, // Location of the node ID
    pub label_location: Option<Location>, // Location of the label, if present
}

/// Enumerates the possible shapes a node can have.
#[derive(Debug, PartialEq, Eq, Clone)]
pub enum NodeShape {
    Rectangle,          // A[Text]
    RoundedRectangle,   // A(Text)
    Stadium,            // A([Text])
    Subroutine,         // A[[Text]]
    Cylinder,           // A[(Text)]
    Circle,             // A((Text))
    Asymmetric,         // A>Text]
    Rhombus,            // A{Text}
    Hexagon,            // A{{Text}}
    Parallelogram,      // A[/Text/]
    ParallelogramAlt,   // A[\Text\]
    Trapezoid,          // A[/Text\]
    TrapezoidAlt,       // A[\Text/]
    DoubleCircle,       // A(((Text)))
    // Add other shapes as per Mermaid specification
    Default, // Implicit rectangle if no shape specified (e.g., A)
}

/// Represents an edge (connection) between two nodes.
#[derive(Debug, PartialEq, Eq, Clone)]
pub struct Edge {
    pub source: String,
    pub target: String,
    pub label: Option<String>,
    pub arrow_type: ArrowType,
    pub line_type: LineType,
    pub location: Location, // Location of the first token of the edge (source ID)
    pub arrow_location: Location, // Location of the arrow/line tokens
    pub label_location: Option<Location>, // Location of the label, if present
}

/// Enumerates the types of arrows that can be used in an edge.
#[derive(Debug, PartialEq, Eq, Clone)]
pub enum ArrowType {
    None,          // ---
    ArrowHead,     // -->
    Cross,         // --x
    Open,          // --o
    Bidirectional, // <-->
    // Add other arrow types (e.g., ===>, --+>, etc.)
}

/// Enumerates the types of lines that can be used in an edge.
#[derive(Debug, PartialEq, Eq, Clone)]
pub enum LineType {
    Solid,       // ---
    Dotted,      // -.->
    Thick,       // ===
    Invisible,   // ~~~
    // Add other line types
}

/// Represents a subgraph in the diagram.
#[derive(Debug, PartialEq, Eq, Clone)]
pub struct Subgraph {
    pub id: String,
    pub label: Option<String>, // Subgraph label (e.g., subgraph "My Subgraph")
    pub statements: Vec<Statement>,
    pub location: Location, // Location of the 'subgraph' keyword
    pub label_location: Option<Location>, // Location of the label, if present
}

/// Represents a directive, like `%%{init: ...}`.
#[derive(Debug, PartialEq, Eq, Clone)]
pub struct Directive {
    pub content: String, // The full content of the directive
    pub location: Location,
}

// Helper to determine node shape from token type
impl NodeShape {
    pub fn from_token_type(token_type: &TokenType) -> Option<Self> {
        match token_type {
            TokenType::NodeOpenBracket => Some(NodeShape::Rectangle),
            TokenType::NodeOpenParen => Some(NodeShape::RoundedRectangle),
            TokenType::NodeOpenSquareBracketParen => Some(NodeShape::Stadium),
            TokenType::NodeOpenDoubleSquareBracket => Some(NodeShape::Subroutine),
            TokenType::NodeOpenParenSquareBracket => Some(NodeShape::Cylinder),
            TokenType::NodeOpenDoubleParen => Some(NodeShape::Circle),
            TokenType::NodeOpenGt => Some(NodeShape::Asymmetric),
            TokenType::NodeOpenBrace => Some(NodeShape::Rhombus),
            TokenType::NodeOpenDoubleBrace => Some(NodeShape::Hexagon),
            TokenType::NodeOpenSlash => Some(NodeShape::Parallelogram),
            TokenType::NodeOpenBackslash => Some(NodeShape::ParallelogramAlt),
            TokenType::NodeOpenSlashBackslash => Some(NodeShape::Trapezoid),
            TokenType::NodeOpenBackslashSlash => Some(NodeShape::TrapezoidAlt),
            TokenType::NodeOpenTripleParen => Some(NodeShape::DoubleCircle),
            _ => None,
        }
    }

    pub fn get_closing_token_type(&self) -> TokenType {
        match self {
            NodeShape::Rectangle => TokenType::NodeCloseBracket,
            NodeShape::RoundedRectangle => TokenType::NodeCloseParen,
            NodeShape::Stadium => TokenType::NodeCloseSquareBracketParen,
            NodeShape::Subroutine => TokenType::NodeCloseDoubleSquareBracket,
            NodeShape::Cylinder => TokenType::NodeCloseParenSquareBracket,
            NodeShape::Circle => TokenType::NodeCloseDoubleParen,
            NodeShape::Asymmetric => TokenType::NodeCloseBracket, // Asymmetric closes with ]
            NodeShape::Rhombus => TokenType::NodeCloseBrace,
            NodeShape::Hexagon => TokenType::NodeCloseDoubleBrace,
            NodeShape::Parallelogram => TokenType::NodeCloseSlash,
            NodeShape::ParallelogramAlt => TokenType::NodeCloseBackslash,
            NodeShape::Trapezoid => TokenType::NodeCloseBackslash, // Trapezoid closes with \]
            NodeShape::TrapezoidAlt => TokenType::NodeCloseSlash, // TrapezoidAlt closes with /]
            NodeShape::DoubleCircle => TokenType::NodeCloseTripleParen,
            NodeShape::Default => TokenType::Identifier, // No explicit closing for default
        }
    }
}

// Helper to determine arrow and line types from token type
impl EdgeDirection {
    // This enum might be useful for representing the entire arrow token
    // For now, we'll parse ArrowType and LineType separately.
    // This can be expanded if we need more complex arrow structures.
}

// src/parser/parser.rs
use crate::{
    diagnostics::{Diagnostic, DiagnosticCode, DiagnosticEmitter, DiagnosticLevel},
    lexer::{Lexer, Token, TokenType},
    parser::ast::*, // Import all AST types
    Location,
};
use std::iter::Peekable;
use std::vec::IntoIter;

/// The Parser transforms a stream of tokens into an Abstract Syntax Tree (AST).
pub struct Parser<'a> {
    lexer: Peekable<IntoIter<Token>>,
    current_token: Token,
    peek_token: Token,
    diagnostics: DiagnosticEmitter<'a>,
    source_code: &'a str,
}

impl<'a> Parser<'a> {
    /// Creates a new Parser instance.
    /// It consumes the lexer's token stream and initializes the current and peek tokens.
    pub fn new(lexer: Lexer, source_code: &'a str, diagnostics: DiagnosticEmitter<'a>) -> Self {
        let mut token_iter = lexer.into_iter().peekable();

        // Prime the current and peek tokens
        let current_token = token_iter.next().unwrap_or_else(|| Token {
            token_type: TokenType::Eof,
            literal: String::new(),
            location: Location::new(0, 0, 0, 0),
        });
        let peek_token = token_iter.next().unwrap_or_else(|| Token {
            token_type: TokenType::Eof,
            literal: String::new(),
            location: Location::new(0, 0, 0, 0),
        });

        Parser {
            lexer: token_iter,
            current_token,
            peek_token,
            diagnostics,
            source_code,
        }
    }

    /// Advances the parser's token stream by one.
    fn advance_token(&mut self) {
        self.current_token = self.peek_token.clone();
        self.peek_token = self.lexer.next().unwrap_or_else(|| Token {
            token_type: TokenType::Eof,
            literal: String::new(),
            location: Location::new(self.source_code.len(), self.source_code.len(), 0, 0),
        });
    }

    /// Checks if the current token matches the expected type.
    fn current_token_is(&self, token_type: TokenType) -> bool {
        self.current_token.token_type == token_type
    }

    /// Checks if the peek token matches the expected type.
    fn peek_token_is(&self, token_type: TokenType) -> bool {
        self.peek_token.token_type == token_type
    }

    /// Consumes the current token if it matches the expected type, otherwise reports an error.
    /// Returns the consumed token or an error diagnostic.
    fn expect_token(&mut self, expected: TokenType) -> Result<Token, Diagnostic> {
        if self.current_token_is(expected) {
            let token = self.current_token.clone();
            self.advance_token();
            Ok(token)
        } else {
            let diagnostic = Diagnostic::new(
                DiagnosticCode::UnexpectedToken,
                format!(
                    "Expected token {:?}, but found {:?}",
                    expected, self.current_token.token_type
                ),
                self.current_token.location,
                DiagnosticLevel::Error,
            )
            .with_help(format!(
                "Ensure that '{}' is used correctly here.",
                expected.to_string()
            ));
            self.diagnostics.emit(diagnostic.clone());
            Err(diagnostic)
        }
    }

    /// Consumes the current token if it matches any of the expected types, otherwise reports an error.
    fn expect_any_token(&mut self, expected_types: &[TokenType]) -> Result<Token, Diagnostic> {
        if expected_types.contains(&self.current_token.token_type) {
            let token = self.current_token.clone();
            self.advance_token();
            Ok(token)
        } else {
            let expected_str = expected_types
                .iter()
                .map(|t| format!("{:?}", t))
                .collect::<Vec<_>>()
                .join(" or ");
            let diagnostic = Diagnostic::new(
                DiagnosticCode::UnexpectedToken,
                format!(
                    "Expected one of {}, but found {:?}",
                    expected_str, self.current_token.token_type
                ),
                self.current_token.location,
                DiagnosticLevel::Error,
            )
            .with_help(format!(
                "Ensure that one of the expected tokens is used here."
            ));
            self.diagnostics.emit(diagnostic.clone());
            Err(diagnostic)
        }
    }

    /// Parses the entire Mermaid diagram. This is the entry point for parsing.
    pub fn parse(&mut self) -> Result<Diagram, Vec<Diagnostic>> {
        let mut collected_diagnostics = Vec::new();

        // First, parse the diagram type (e.g., 'graph', 'flowchart', 'sequence')
        let diagram_type_token = match self.expect_any_token(&[
            TokenType::KeywordGraph,
            TokenType::KeywordFlowchart,
            TokenType::KeywordSequence,
            TokenType::KeywordClass,
            // Add other diagram type keywords here
        ]) {
            Ok(t) => t,
            Err(diag) => {
                collected_diagnostics.push(diag);
                // Attempt to recover by assuming a default diagram type or skipping until a known keyword
                // For now, we'll just return the error. Robust error recovery comes later.
                return Err(collected_diagnostics);
            }
        };

        let diagram_type_location = diagram_type_token.location;
        let diagram_type = match self.parse_diagram_type(&diagram_type_token) {
            Ok(dt) => dt,
            Err(diag) => {
                collected_diagnostics.push(diag);
                // Recovery: assume Flowchart(None)
                DiagramType::Flowchart(None)
            }
        };

        let mut statements = Vec::new();
        while !self.current_token_is(TokenType::Eof) {
            match self.parse_statement() {
                Ok(Some(stmt)) => statements.push(stmt),
                Ok(None) => { /* No statement parsed, possibly a skipped invalid token */ }
                Err(diag) => {
                    collected_diagnostics.push(diag);
                    // Error recovery: skip tokens until a likely statement start or EOF
                    self.synchronize_parser();
                }
            }
        }

        if collected_diagnostics.is_empty() {
            Ok(Diagram {
                diagram_type,
                statements,
                location: diagram_type_location,
            })
        } else {
            Err(collected_diagnostics)
        }
    }

    /// Parses the specific diagram type and its associated properties (e.g., direction for flowchart).
    fn parse_diagram_type(&mut self, token: &Token) -> Result<DiagramType, Diagnostic> {
        match token.token_type {
            TokenType::KeywordGraph | TokenType::KeywordFlowchart => {
                // Flowcharts can have an optional direction (TD, LR, etc.)
                if self.peek_token_is(TokenType::KeywordTD) {
                    self.advance_token(); // Consume TD
                    self.advance_token(); // Consume the keyword itself
                    Ok(DiagramType::Flowchart(Some(FlowchartDirection::TopToBottom)))
                } else if self.peek_token_is(TokenType::KeywordLR) {
                    self.advance_token(); // Consume LR
                    self.advance_token(); // Consume the keyword itself
                    Ok(DiagramType::Flowchart(Some(FlowchartDirection::LeftToRight)))
                } else if self.peek_token_is(TokenType::KeywordBT) {
                    self.advance_token(); // Consume BT
                    self.advance_token(); // Consume the keyword itself
                    Ok(DiagramType::Flowchart(Some(FlowchartDirection::BottomToTop)))
                } else if self.peek_token_is(TokenType::KeywordRL) {
                    self.advance_token(); // Consume RL
                    self.advance_token(); // Consume the keyword itself
                    Ok(DiagramType::Flowchart(Some(FlowchartDirection::RightToLeft)))
                } else {
                    // No explicit direction, default for flowchart/graph is TD
                    Ok(DiagramType::Flowchart(Some(FlowchartDirection::TopToBottom)))
                }
            }
            TokenType::KeywordSequence => Ok(DiagramType::Sequence),
            TokenType::KeywordClass => Ok(DiagramType::Class),
            _ => {
                // This should ideally not be reached if `expect_any_token` handles it correctly.
                // But as a fallback, report an error.
                Err(Diagnostic::new(
                    DiagnosticCode::InvalidDiagramType,
                    format!("Unknown diagram type: {:?}", token.literal),
                    token.location,
                    DiagnosticLevel::Error,
                ))
            }
        }
    }

    /// Parses a single statement (Node, Edge, Subgraph, Directive, Comment, or EmptyLine).
    /// Returns `Ok(Some(Statement))` if a statement was successfully parsed.
    /// Returns `Ok(None)` if the current token doesn't start a valid statement but isn't an error
    /// (e.g., whitespace tokens that should be ignored at this level).
    /// Returns `Err(Diagnostic)` if a syntax error prevents parsing a statement.
    fn parse_statement(&mut self) -> Result<Option<Statement>, Diagnostic> {
        // Skip whitespace and comments at the statement level for now.
        // The lexer already handles most whitespace, but empty lines or comments
        // between statements might be passed through.
        while self.current_token_is(TokenType::Whitespace) || self.current_token_is(TokenType::Newline) {
            self.advance_token();
        }

        if self.current_token_is(TokenType::Eof) {
            return Ok(None);
        }

        if self.current_token_is(TokenType::DirectiveStart) {
            return self.parse_directive().map(Some);
        }

        if self.current_token_is(TokenType::Comment) {
            let comment_token = self.current_token.clone();
            self.advance_token();
            return Ok(Some(Statement::Comment(comment_token.literal, comment_token.location)));
        }

        // A statement typically starts with an identifier (for a node or source of an edge)
        // or a keyword like 'subgraph'.
        if self.current_token_is(TokenType::Identifier) {
            let id_token = self.current_token.clone();
            if self.peek_token_is(TokenType::Arrow)
                || self.peek_token_is(TokenType::DottedArrow)
                || self.peek_token_is(TokenType::ThickArrow)
                || self.peek_token_is(TokenType::InvisibleArrow)
                || self.peek_token_is(TokenType::BidirectionalArrow)
                || self.peek_token_is(TokenType::EdgeOpenParen) // A---B style
                || self.peek_token_is(TokenType::EdgeOpenBracket) // A--label--B style
            {
                // This looks like an edge
                return self.parse_edge().map(Some);
            } else {
                // This looks like a node
                return self.parse_node().map(Some);
            }
        } else if self.current_token_is(TokenType::KeywordSubgraph) {
            return self.parse_subgraph().map(Some);
        }

        // If we reach here, it's an unexpected token at the start of a statement.
        let diagnostic = Diagnostic::new(
            DiagnosticCode::UnexpectedToken,
            format!(
                "Unexpected token '{:?}' at the start of a statement.",
                self.current_token.token_type
            ),
            self.current_token.location,
            DiagnosticLevel::Error,
        )
        .with_help("Expected a node ID, 'subgraph' keyword, a directive, or a comment.");
        self.diagnostics.emit(diagnostic.clone());

        // Error recovery: advance token to try and find a valid statement start
        self.advance_token();
        Err(diagnostic)
    }

    /// Parses a node definition (e.g., `A[Start]`, `B(Process)`, `C`).
    fn parse_node(&mut self) -> Result<Node, Diagnostic> {
        let id_token = self.expect_token(TokenType::Identifier)?;
        let id = id_token.literal.clone();
        let mut label: Option<String> = None;
        let mut node_shape = NodeShape::Default;
        let mut label_location: Option<Location> = None;

        // Check for node shape and label
        if let Some(shape) = NodeShape::from_token_type(&self.current_token.token_type) {
            let open_bracket_token = self.current_token.clone();
            self.advance_token(); // Consume opening bracket/paren

            // Node label can be a quoted string or unquoted identifier-like text
            if self.current_token_is(TokenType::StringLiteral) {
                label = Some(self.current_token.literal.clone());
                label_location = Some(self.current_token.location);
                self.advance_token(); // Consume string literal
            } else if self.current_token_is(TokenType::Identifier) {
                // Allow identifiers as labels if not quoted (Mermaid often does this)
                label = Some(self.current_token.literal.clone());
                label_location = Some(self.current_token.location);
                self.advance_token(); // Consume identifier as label
            } else {
                // Empty label is allowed (e.g., A[])
                label = Some("".to_string());
                label_location = Some(self.current_token.location); // Location of the empty string
            }

            // Expect closing bracket/paren
            let expected_closing_type = shape.get_closing_token_type();
            self.expect_token(expected_closing_type)?;
            node_shape = shape;
        }

        Ok(Node {
            id,
            label,
            shape: node_shape,
            location: id_token.location,
            label_location,
        })
    }

    /// Parses an edge definition (e.g., `A --> B`, `B -- "label" --> C`).
    fn parse_edge(&mut self) -> Result<Edge, Diagnostic> {
        let source_token = self.expect_token(TokenType::Identifier)?;
        let source_id = source_token.literal.clone();
        let edge_start_location = source_token.location; // Start location of the edge

        let mut label: Option<String> = None;
        let mut arrow_type = ArrowType::None; // Default to no arrow
        let mut line_type = LineType::Solid;   // Default to solid line
        let mut arrow_location = self.current_token.location; // Default, will update

        // Handle optional label between source and arrow (e.g., A -- label -- B)
        if self.current_token_is(TokenType::EdgeOpenParen) || self.current_token_is(TokenType::EdgeOpenBracket) {
            let open_token = self.current_token.clone();
            self.advance_token(); // Consume '(' or '['

            if self.current_token_is(TokenType::StringLiteral) {
                label = Some(self.current_token.literal.clone());
                self.advance_token(); // Consume label
            } else if self.current_token_is(TokenType::Identifier) {
                label = Some(self.current_token.literal.clone());
                self.advance_token(); // Consume label
            } else {
                // Empty label
            }

            // Expect closing token
            let expected_close = match open_token.token_type {
                TokenType::EdgeOpenParen => TokenType::EdgeCloseParen,
                TokenType::EdgeOpenBracket => TokenType::EdgeCloseBracket,
                _ => unreachable!(), // Should be caught by the if condition
            };
            self.expect_token(expected_close)?;
        }

        // Parse arrow/line type
        let arrow_line_token = self.expect_any_token(&[
            TokenType::Arrow,
            TokenType::DottedArrow,
            TokenType::ThickArrow,
            TokenType::InvisibleArrow,
            TokenType::BidirectionalArrow,
            // Add more specific arrow types as they are tokenized
        ])?;
        arrow_location = arrow_line_token.location;

        match arrow_line_token.token_type {
            TokenType::Arrow => { // e.g., -->, ---
                if arrow_line_token.literal.ends_with('>') {
                    arrow_type = ArrowType::ArrowHead;
                } else if arrow_line_token.literal.ends_with('x') { // --x
                    arrow_type = ArrowType::Cross;
                } else if arrow_line_token.literal.ends_with('o') { // --o
                    arrow_type = ArrowType::Open;
                }
                line_type = LineType::Solid; // Default for '---' and '-->'
            }
            TokenType::DottedArrow => { // e.g., -.->
                arrow_type = ArrowType::ArrowHead;
                line_type = LineType::Dotted;
            }
            TokenType::ThickArrow => { // e.g., ===>, ===
                arrow_type = ArrowType::ArrowHead;
                line_type = LineType::Thick;
            }
            TokenType::InvisibleArrow => { // e.g., ~~~
                arrow_type = ArrowType::None;
                line_type = LineType::Invisible;
            }
            TokenType::BidirectionalArrow => { // e.g., <-->
                arrow_type = ArrowType::Bidirectional;
                line_type = LineType::Solid;
            }
            _ => { /* Already handled by expect_any_token */ }
        }

        // Handle optional label after arrow (e.g., A --> "label" B)
        if self.current_token_is(TokenType::StringLiteral) {
            label = Some(self.current_token.literal.clone());
            self.advance_token(); // Consume label
        } else if self.current_token_is(TokenType::Identifier) {
            label = Some(self.current_token.literal.clone());
            self.advance_token(); // Consume label
        }

        let target_token = self.expect_token(TokenType::Identifier)?;
        let target_id = target_token.literal.clone();

        Ok(Edge {
            source: source_id,
            target: target_id,
            label,
            arrow_type,
            line_type,
            location: edge_start_location,
            arrow_location,
            label_location: None, // TODO: Accurately track label location
        })
    }

    /// Parses a subgraph definition.
    fn parse_subgraph(&mut self) -> Result<Subgraph, Diagnostic> {
        let subgraph_keyword_token = self.expect_token(TokenType::KeywordSubgraph)?;
        let subgraph_location = subgraph_keyword_token.location;

        let id_token = self.expect_token(TokenType::Identifier)?;
        let id = id_token.literal.clone();

        let mut label: Option<String> = None;
        let mut label_location: Option<Location> = None;

        // Optional subgraph label (e.g., subgraph MySubgraph ["My Label"])
        if self.current_token_is(TokenType::NodeOpenBracket) {
            self.advance_token(); // Consume '['
            if self.current_token_is(TokenType::StringLiteral) {
                label = Some(self.current_token.literal.clone());
                label_location = Some(self.current_token.location);
                self.advance_token(); // Consume string literal
            } else if self.current_token_is(TokenType::Identifier) {
                // Allow unquoted identifiers as labels
                label = Some(self.current_token.literal.clone());
                label_location = Some(self.current_token.location);
                self.advance_token(); // Consume identifier
            }
            self.expect_token(TokenType::NodeCloseBracket)?; // Consume ']'
        }

        let mut statements = Vec::new();
        // Subgraph body can be enclosed in { } or just indented lines
        // For now, let's assume it's followed by statements until 'end' or EOF
        // A more robust parser would check for indentation or explicit `{}`
        while !self.current_token_is(TokenType::KeywordEnd) && !self.current_token_is(TokenType::Eof) {
            match self.parse_statement() {
                Ok(Some(stmt)) => statements.push(stmt),
                Ok(None) => { /* Continue */ },
                Err(diag) => {
                    self.diagnostics.emit(diag);
                    self.synchronize_parser();
                }
            }
        }

        self.expect_token(TokenType::KeywordEnd)?; // Consume 'end' keyword

        Ok(Subgraph {
            id,
            label,
            statements,
            location: subgraph_location,
            label_location,
        })
    }

    /// Parses a directive (e.g., `%%{init: { ... } }%%`).
    fn parse_directive(&mut self) -> Result<Directive, Diagnostic> {
        let directive_start_token = self.expect_token(TokenType::DirectiveStart)?;
        let content = directive_start_token.literal.clone(); // The lexer should have captured full content

        // The lexer is designed to capture the full directive content as part of the DirectiveStart token.
        // So, we just need to consume the token and extract its literal.
        // No explicit '%%}' token is expected here if the lexer captures the whole thing.

        Ok(Directive {
            content,
            location: directive_start_token.location,
        })
    }

    /// Implements basic error recovery by advancing tokens until a likely synchronization point.
    fn synchronize_parser(&mut self) {
        loop {
            if self.current_token_is(TokenType::Eof) {
                return;
            }

            match self.current_token.token_type {
                TokenType::KeywordGraph
                | TokenType::KeywordFlowchart
                | TokenType::KeywordSequence
                | TokenType::KeywordClass
                | TokenType::KeywordSubgraph
                | TokenType::Identifier => {
                    // These tokens usually start a new statement or diagram
                    return;
                }
                _ => {
                    // Skip over current token
                    self.advance_token();
                }
            }
        }
    }
}

// src/parser/parser_tests.rs
#[cfg(test)]
mod tests {
    use super::*;
    use crate::lexer::Lexer;
    use crate::diagnostics::DiagnosticEmitter;
    use crate::Location;

    // Helper function to create a parser from source code
    fn setup_parser<'a>(source: &'a str, emitter: DiagnosticEmitter<'a>) -> Parser<'a> {
        let lexer = Lexer::new(source);
        Parser::new(lexer, source, emitter)
    }

    #[test]
    fn test_parse_simple_flowchart() {
        let source = "graph TD\n    A[Start] --> B(Process)";
        let emitter = DiagnosticEmitter::new(source);
        let mut parser = setup_parser(source, emitter.clone());

        let result = parser.parse();
        assert!(result.is_ok(), "Parsing failed with errors: {:?}", emitter.get_diagnostics());

        let diagram = result.unwrap();
        assert_eq!(diagram.diagram_type, DiagramType::Flowchart(Some(FlowchartDirection::TopToBottom)));
        assert_eq!(diagram.statements.len(), 2); // One node, one edge

        // Verify Node A
        if let Statement::Node(node_a) = &diagram.statements[0] {
            assert_eq!(node_a.id, "A");
            assert_eq!(node_a.label, Some("Start".to_string()));
            assert_eq!(node_a.shape, NodeShape::Rectangle);
        } else {
            panic!("Expected Node statement, found {:?}", diagram.statements[0]);
        }

        // Verify Edge A --> B
        if let Statement::Edge(edge_ab) = &diagram.statements[1] {
            assert_eq!(edge_ab.source, "A");
            assert_eq!(edge_ab.target, "B");
            assert_eq!(edge_ab.label, None);
            assert_eq!(edge_ab.arrow_type, ArrowType::ArrowHead);
            assert_eq!(edge_ab.line_type, LineType::Solid);
        } else {
            panic!("Expected Edge statement, found {:?}", diagram.statements[1]);
        }

        assert!(emitter.get_diagnostics().is_empty(), "Unexpected diagnostics: {:?}", emitter.get_diagnostics());
    }

    #[test]
    fn test_parse_flowchart_with_direction_lr() {
        let source = "flowchart LR\n    Node1 --> Node2";
        let emitter = DiagnosticEmitter::new(source);
        let mut parser = setup_parser(source, emitter.clone());

        let result = parser.parse();
        assert!(result.is_ok(), "Parsing failed with errors: {:?}", emitter.get_diagnostics());
        let diagram = result.unwrap();
        assert_eq!(diagram.diagram_type, DiagramType::Flowchart(Some(FlowchartDirection::LeftToRight)));
        assert!(emitter.get_diagnostics().is_empty());
    }

    #[test]
    fn test_parse_node_shapes_and_labels() {
        let source = "graph TD\n    A[Rect]\n    B(Rounded)\n    C>Asym]\n    D{Rhombus}";
        let emitter = DiagnosticEmitter::new(source);
        let mut parser = setup_parser(source, emitter.clone());
        let result = parser.parse().unwrap();

        assert_eq!(result.statements.len(), 4);
        if let Statement::Node(a) = &result.statements[0] { assert_eq!(a.shape, NodeShape::Rectangle); assert_eq!(a.label, Some("Rect".to_string())); }
        if let Statement::Node(b) = &result.statements[1] { assert_eq!(b.shape, NodeShape::RoundedRectangle); assert_eq!(b.label, Some("Rounded".to_string())); }
        if let Statement::Node(c) = &result.statements[2] { assert_eq!(c.shape, NodeShape::Asymmetric); assert_eq!(c.label, Some("Asym".to_string())); }
        if let Statement::Node(d) = &result.statements[3] { assert_eq!(d.shape, NodeShape::Rhombus); assert_eq!(d.label, Some("Rhombus".to_string())); }
        assert!(emitter.get_diagnostics().is_empty());
    }

    #[test]
    fn test_parse_edge_labels_and_types() {
        let source = "graph TD\n    A -- Label A --> B\n    B --o C\n    C ===> D\n    D -.-> E\n    E ~~~ F";
        let emitter = DiagnosticEmitter::new(source);
        let mut parser = setup_parser(source, emitter.clone());
        let result = parser.parse().unwrap();

        assert_eq!(result.statements.len(), 5);
        if let Statement::Edge(e) = &result.statements[0] { assert_eq!(e.label, Some("Label A".to_string())); assert_eq!(e.arrow_type, ArrowType::ArrowHead); }
        if let Statement::Edge(e) = &result.statements[1] { assert_eq!(e.label, None); assert_eq!(e.arrow_type, ArrowType::Open); }
        if let Statement::Edge(e) = &result.statements[2] { assert_eq!(e.label, None); assert_eq!(e.arrow_type, ArrowType::ArrowHead); assert_eq!(e.line_type, LineType::Thick); }
        if let Statement::Edge(e) = &result.statements[3] { assert_eq!(e.label, None); assert_eq!(e.arrow_type, ArrowType::ArrowHead); assert_eq!(e.line_type, LineType::Dotted); }
        if let Statement::Edge(e) = &result.statements[4] { assert_eq!(e.label, None); assert_eq!(e.arrow_type, ArrowType::None); assert_eq!(e.line_type, LineType::Invisible); }
        assert!(emitter.get_diagnostics().is_empty());
    }

    #[test]
    fn test_parse_subgraph() {
        let source = r#"
graph TD
    subgraph MySubgraph ["A Subgraph"]
        A --> B
    end
    B --> C
"#;
        let emitter = DiagnosticEmitter::new(source);
        let mut parser = setup_parser(source, emitter.clone());
        let result = parser.parse().unwrap();

        assert_eq!(result.statements.len(), 2);
        if let Statement::Subgraph(subgraph) = &result.statements[0] {
            assert_eq!(subgraph.id, "MySubgraph");
            assert_eq!(subgraph.label, Some("A Subgraph".to_string()));
            assert_eq!(subgraph.statements.len(), 1);
            if let Statement::Edge(edge) = &subgraph.statements[0] {
                assert_eq!(edge.source, "A");
                assert_eq!(edge.target, "B");
            } else { panic!("Expected edge in subgraph"); }
        } else { panic!("Expected subgraph"); }

        if let Statement::Edge(edge) = &result.statements[1] {
            assert_eq!(edge.source, "B");
            assert_eq!(edge.target, "C");
        } else { panic!("Expected edge outside subgraph"); }
        assert!(emitter.get_diagnostics().is_empty());
    }

    #[test]
    fn test_parse_directive() {
        let source = "%%{init: { 'theme': 'forest' } }%%\ngraph TD\nA --> B";
        let emitter = DiagnosticEmitter::new(source);
        let mut parser = setup_parser(source, emitter.clone());
        let result = parser.parse().unwrap();

        assert_eq!(result.statements.len(), 2);
        if let Statement::Directive(directive) = &result.statements[0] {
            assert_eq!(directive.content, "{init: { 'theme': 'forest' } }"); // Lexer should capture content without %%{}%%
        } else { panic!("Expected directive"); }
        assert!(emitter.get_diagnostics().is_empty());
    }

    #[test]
    fn test_parse_syntax_error_missing_diagram_type() {
        let source = "A --> B"; // Missing 'graph TD'
        let emitter = DiagnosticEmitter::new(source);
        let mut parser = setup_parser(source, emitter.clone());

        let result = parser.parse();
        assert!(result.is_err());
        let diagnostics = emitter.get_diagnostics();
        assert_eq!(diagnostics.len(), 1);
        assert_eq!(diagnostics[0].code, DiagnosticCode::UnexpectedToken);
        assert!(diagnostics[0].message.contains("Expected one of KeywordGraph or KeywordFlowchart or KeywordSequence or KeywordClass"));
    }

    #[test]
    fn test_parse_syntax_error_unmatched_bracket() {
        let source = "graph TD\n    A[Start --> B"; // Missing closing bracket
        let emitter = DiagnosticEmitter::new(source);
        let mut parser = setup_parser(source, emitter.clone());

        let result = parser.parse();
        assert!(result.is_err()); // Should return an error
        let diagnostics = emitter.get_diagnostics();
        assert_eq!(diagnostics.len(), 1);
        assert_eq!(diagnostics[0].code, DiagnosticCode::UnexpectedToken);
        assert!(diagnostics[0].message.contains("Expected token NodeCloseBracket"));
    }

    #[test]
    fn test_parse_empty_input() {
        let source = "";
        let emitter = DiagnosticEmitter::new(source);
        let mut parser = setup_parser(source, emitter.clone());

        let result = parser.parse();
        assert!(result.is_err()); // Empty input is invalid as it lacks a diagram type
        let diagnostics = emitter.get_diagnostics();
        assert_eq!(diagnostics.len(), 1);
        assert_eq!(diagnostics[0].code, DiagnosticCode::UnexpectedToken);
    }
}