Chapter 6: Rich Diagnostics: Emitting Compiler-Style Error Messages

# Cargo.toml

[package]
name = "mermaid-analyzer"
version = "0.1.0"
edition = "2021"

[dependencies]
ariadne = "0.4.0" # Add this line
log = "0.4"
env_logger = "0.11"
# ... other dependencies from previous chapters

// src/diagnostics/span.rs

/// A unique identifier for a source file.
/// For a CLI tool processing a single input, this might just be "input".
#[derive(Debug, Clone, PartialEq, Eq, Hash, Copy)]
pub struct SourceId(pub &'static str); // Using &'static str for simplicity for now

impl From<&'static str> for SourceId {
    fn from(s: &'static str) -> Self {
        SourceId(s)
    }
}

/// Represents a contiguous region in the source code.
/// Stores byte offsets for precise highlighting and line/column for human readability.
#[derive(Debug, Clone, PartialEq, Eq, Copy)]
pub struct Span {
    pub source_id: SourceId,
    pub start: usize, // Start byte offset
    pub end: usize,   // End byte offset (exclusive)
    pub start_line: usize,
    pub end_line: usize,
    pub start_column: usize,
    pub end_column: usize,
}

impl Span {
    /// Creates a new span.
    pub fn new(
        source_id: SourceId,
        start: usize,
        end: usize,
        start_line: usize,
        end_line: usize,
        start_column: usize,
        end_column: usize,
    ) -> Self {
        Self {
            source_id,
            start,
            end,
            start_line,
            end_line,
            start_column,
            end_column,
        }
    }

    /// Creates a dummy span for cases where a real span isn't available (e.g., internal errors).
    pub fn dummy() -> Self {
        Self {
            source_id: SourceId("dummy"),
            start: 0,
            end: 0,
            start_line: 1,
            end_line: 1,
            start_column: 1,
            end_column: 1,
        }
    }

    /// Returns the length of the span in bytes.
    pub fn len(&self) -> usize {
        self.end - self.start
    }

    /// Checks if the span is empty.
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Combines two spans into a single span that covers both.
    /// Assumes both spans are from the same source.
    pub fn merge(&self, other: &Self) -> Self {
        assert_eq!(self.source_id, other.source_id, "Cannot merge spans from different sources");

        let start = self.start.min(other.start);
        let end = self.end.max(other.end);

        let start_line = self.start_line.min(other.start_line);
        let end_line = self.end_line.max(other.end_line);

        let start_column = if self.start_line == start_line {
            self.start_column
        } else {
            other.start_column
        };

        let end_column = if self.end_line == end_line {
            self.end_column
        } else {
            other.end_column
        };


        Self {
            source_id: self.source_id,
            start,
            end,
            start_line,
            end_line,
            start_column,
            end_column,
        }
    }
}

// Implement ariadne's Span trait for our Span struct
impl ariadne::Span for Span {
    type SourceId = SourceId;

    fn source(&self) -> &Self::SourceId {
        &self.source_id
    }

    fn start(&self) -> usize {
        self.start
    }

    fn end(&self) -> usize {
        self.end
    }
}

// src/diagnostics/error_codes.rs

/// Represents the severity level of a diagnostic.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Severity {
    Error,
    Warning,
    Note,
    Help,
}

/// Unique error codes for different types of issues.
/// M = Mermaid Analyzer
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ErrorCode {
    // Lexer errors (M001-M099)
    M001, // Unexpected character
    M002, // Unterminated string literal

    // Parser errors (M100-M199)
    M100, // Expected token, found something else
    M101, // Unmatched parenthesis/bracket
    M102, // Missing diagram type declaration (e.g., 'graph TD')
    M103, // Invalid node ID
    M104, // Malformed edge definition
    M105, // Invalid diagram type

    // Semantic / Validation errors (M200-M299) - will be used in future chapters
    M200, // Duplicate node ID
    M201, // Undefined node in edge
    M202, // Invalid direction for diagram type
    M203, // Cyclic dependency detected
    // ... more codes as we expand
}

impl ErrorCode {
    /// Returns a short, descriptive message for the error code.
    pub fn message(&self) -> &'static str {
        match self {
            ErrorCode::M001 => "Unexpected character",
            ErrorCode::M002 => "Unterminated string literal",
            ErrorCode::M100 => "Unexpected token",
            ErrorCode::M101 => "Unmatched parenthesis or bracket",
            ErrorCode::M102 => "Missing diagram type declaration",
            ErrorCode::M103 => "Invalid node ID format",
            ErrorCode::M104 => "Malformed edge definition",
            ErrorCode::M105 => "Invalid or unsupported diagram type",
            ErrorCode::M200 => "Duplicate node identifier",
            ErrorCode::M201 => "Edge refers to an undefined node",
            ErrorCode::M202 => "Invalid direction for this diagram type",
            ErrorCode::M203 => "Cyclic dependency detected in graph",
        }
    }

    /// Returns a longer, more detailed help message for the error code.
    pub fn help(&self) -> Option<&'static str> {
        match self {
            ErrorCode::M001 => Some("This character is not allowed in Mermaid syntax at this position. Please check for typos or invalid symbols."),
            ErrorCode::M002 => Some("String literals must be closed with a matching quote. Add a '\"' or apostrophe to terminate the string."),
            ErrorCode::M100 => Some("The parser expected a different kind of token here. This usually indicates a syntax error. Check the Mermaid syntax documentation for this construct."),
            ErrorCode::M101 => Some("Ensure all parentheses, brackets, and braces are correctly matched and balanced. Each opening symbol must have a corresponding closing symbol."),
            ErrorCode::M102 => Some("All Mermaid diagrams must start with a declaration like 'graph TD', 'sequenceDiagram', or 'classDiagram'. Add a diagram type declaration at the beginning of your code."),
            ErrorCode::M103 => Some("Node IDs must follow Mermaid naming conventions. They typically consist of alphanumeric characters and underscores, or be quoted if they contain special characters."),
            ErrorCode::M104 => Some("An edge definition must specify two nodes and a valid arrow type (e.g., 'A --> B'). Check for missing nodes, invalid arrow syntax, or extra characters."),
            ErrorCode::M105 => Some("The specified diagram type is either misspelled or not supported by Mermaid, or not yet implemented by this analyzer. Refer to Mermaid documentation for supported types."),
            ErrorCode::M200 => Some("Node IDs must be unique within a diagram. Rename one of the nodes to resolve this conflict."),
            ErrorCode::M201 => Some("All nodes referenced in an edge must be defined in the diagram. Ensure node '{node_id}' is declared before it's used in an edge."),
            ErrorCode::M202 => Some("The chosen direction (e.g., 'TD', 'LR') is not valid for this specific diagram type. Consult the Mermaid documentation for valid directions."),
            ErrorCode::M203 => Some("A cycle was detected in your graph, which might indicate a logical error or infinite loop in certain contexts. Review the relationships between nodes."),
            _ => None,
        }
    }
}

// src/diagnostics/mod.rs

pub mod error_codes;
pub mod span;
pub mod emitter;

use error_codes::{ErrorCode, Severity};
use span::{Span, SourceId};

/// Represents a label associated with a diagnostic.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DiagnosticLabel {
    pub span: Span,
    pub message: String,
    pub is_primary: bool, // True for the main point of interest
}

/// A comprehensive diagnostic message, similar to a compiler error.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Diagnostic {
    pub severity: Severity,
    pub code: ErrorCode,
    pub message: String,
    pub labels: Vec<DiagnosticLabel>,
    pub notes: Vec<String>,
    pub help: Option<String>,
}

impl Diagnostic {
    /// Creates a new `DiagnosticBuilder` for constructing a diagnostic.
    pub fn new(severity: Severity, code: ErrorCode) -> DiagnosticBuilder {
        DiagnosticBuilder::new(severity, code)
    }

    /// Convenience method for creating an error diagnostic.
    pub fn error(code: ErrorCode) -> DiagnosticBuilder {
        DiagnosticBuilder::new(Severity::Error, code)
    }

    /// Convenience method for creating a warning diagnostic.
    pub fn warning(code: ErrorCode) -> DiagnosticBuilder {
        DiagnosticBuilder::new(Severity::Warning, code)
    }

    /// Convenience method for creating a note diagnostic.
    pub fn note(code: ErrorCode) -> DiagnosticBuilder {
        DiagnosticBuilder::new(Severity::Note, code)
    }
}

/// A builder for constructing `Diagnostic` instances more ergonomically.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DiagnosticBuilder {
    severity: Severity,
    code: ErrorCode,
    message: Option<String>,
    labels: Vec<DiagnosticLabel>,
    notes: Vec<String>,
    help: Option<String>,
}

impl DiagnosticBuilder {
    pub fn new(severity: Severity, code: ErrorCode) -> Self {
        Self {
            severity,
            code,
            message: None,
            labels: Vec::new(),
            notes: Vec::new(),
            help: None,
        }
    }

    /// Sets the primary message for the diagnostic. If not set, a default from `ErrorCode` is used.
    pub fn with_message(mut self, message: impl Into<String>) -> Self {
        self.message = Some(message.into());
        self
    }

    /// Adds a primary label, which points to the main location of the issue.
    pub fn with_primary_label(mut self, span: Span, message: impl Into<String>) -> Self {
        self.labels.push(DiagnosticLabel {
            span,
            message: message.into(),
            is_primary: true,
        });
        self
    }

    /// Adds a secondary label, for additional context or related locations.
    pub fn with_secondary_label(mut self, span: Span, message: impl Into<String>) -> Self {
        self.labels.push(DiagnosticLabel {
            span,
            message: message.into(),
            is_primary: false,
        });
        self
    }

    /// Adds a note, which is additional textual information.
    pub fn with_note(mut self, note: impl Into<String>) -> Self {
        self.notes.push(note.into());
        self
    }

    /// Sets the help message. If not set, a default from `ErrorCode` is used.
    pub fn with_help(mut self, help: impl Into<String>) -> Self {
        self.help = Some(help.into());
        self
    }

    /// Builds the `Diagnostic` instance.
    pub fn build(self) -> Diagnostic {
        let code_message = self.code.message();
        let code_help = self.code.help();

        Diagnostic {
            severity: self.severity,
            code: self.code,
            message: self.message.unwrap_or_else(|| code_message.to_string()),
            labels: self.labels,
            notes: self.notes,
            help: self.help.or_else(|| code_help.map(|s| s.to_string())),
        }
    }
}

// src/lexer/token.rs

use crate::diagnostics::span::Span; // Import Span

/// Represents the different types of tokens in Mermaid syntax.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum TokenType {
    // Keywords
    Graph,
    Flowchart,
    SequenceDiagram,
    ClassDiagram,
    StateDiagram,
    Gantt,
    Pie,
    GitGraph,
    Ermaid,
    Journey,
    C4c,
    Mindmap,
    Timeline,
    Block,
    SanKey,
    RequirementDiagram,
    // Directions (for graph/flowchart)
    TD, // Top-Down
    BT, // Bottom-Top
    LR, // Left-Right
    RL, // Right-Left
    TB, // Top-Bottom (alias for TD)

    // Structural elements
    OpenParen,    // (
    CloseParen,   // )
    OpenBracket,  // [
    CloseBracket, // ]
    OpenBrace,    // {
    CloseBrace,   // }
    DoubleOpenBracket, // [[
    DoubleCloseBracket, // ]]
    Colon,        // :
    SemiColon,    // ;
    Comma,        // ,
    Dot,          // .
    Equals,       // =
    Plus,         // +
    Minus,        // -
    Star,         // *
    Hash,         // #
    Percent,      // %
    Pipe,         // |
    Backslash,    // \

    // Arrows and connectors
    ArrowRight,     // -->
    ArrowLeft,      // <--
    DoubleArrow,    // <-->
    ThickArrowRight, // ==>
    ThickArrowLeft,  // <==
    ThickDoubleArrow, // <==>
    DottedArrowRight, // -.->
    DottedArrowLeft,  // <-.
    DottedDoubleArrow, // <-->
    CrossArrowRight, // --x
    OpenCircleArrowRight, // --o
    ArrowRightWithText, // --- text -->
    ArrowLeftWithText, // <-- text ---
    DoubleArrowWithText, // <--- text --->
    ThickArrowRightWithText, // === text ==>
    DottedArrowRightWithText, // -.- text -.->

    // Literals and identifiers
    Identifier(String),
    StringLiteral(String), // Quoted strings
    Number(String),        // Numeric literals

    // Comments
    LineComment(String),   // %% comment
    BlockComment(String),  // /* comment */

    // Special
    NewLine,
    Whitespace, // Ignored by parser, but useful for span tracking in lexer
    EOF,
}

/// Represents a token found by the lexer, including its type and source span.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Token {
    pub token_type: TokenType,
    pub span: Span, // Added Span here
}

impl Token {
    pub fn new(token_type: TokenType, span: Span) -> Self {
        Token { token_type, span }
    }
}

// src/lexer/lexer.rs

use crate::diagnostics::span::{Span, SourceId};
use crate::diagnostics::{Diagnostic, DiagnosticBuilder, error_codes::{ErrorCode, Severity}}; // Import Diagnostic components
use crate::lexer::token::{Token, TokenType};
use log::debug;

pub struct Lexer<'a> {
    source: &'a str,
    source_id: SourceId,
    chars: std::iter::Peekable<std::str::Chars<'a>>,
    current_offset: usize, // Current byte offset
    current_line: usize,
    current_column: usize,
    diagnostics: Vec<Diagnostic>, // Collect diagnostics
}

impl<'a> Lexer<'a> {
    pub fn new(source: &'a str, source_id: SourceId) -> Self {
        Lexer {
            source,
            source_id,
            chars: source.chars().peekable(),
            current_offset: 0,
            current_line: 1,
            current_column: 1,
            diagnostics: Vec::new(),
        }
    }

    /// Lexes the entire source and returns a vector of tokens and any diagnostics.
    pub fn lex(mut self) -> (Vec<Token>, Vec<Diagnostic>) {
        let mut tokens = Vec::new();

        while self.peek().is_some() {
            let start_offset = self.current_offset;
            let start_line = self.current_line;
            let start_column = self.current_column;

            if let Some(token_type) = self.next_token() {
                let end_offset = self.current_offset;
                let end_line = self.current_line;
                let end_column = self.current_column;

                let span = Span::new(
                    self.source_id,
                    start_offset,
                    end_offset,
                    start_line,
                    end_line,
                    start_column,
                    end_column,
                );
                tokens.push(Token::new(token_type, span));
            } else {
                // If next_token returns None, it means an error occurred and a diagnostic was emitted.
                // We try to recover by advancing past the invalid character.
                self.advance();
            }
        }

        let eof_span = Span::new(
            self.source_id,
            self.current_offset,
            self.current_offset,
            self.current_line,
            self.current_line,
            self.current_column,
            self.current_column,
        );
        tokens.push(Token::new(TokenType::EOF, eof_span));

        (tokens, self.diagnostics)
    }

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

    fn advance(&mut self) -> Option<char> {
        if let Some(c) = self.chars.next() {
            let char_len = c.len_utf8();
            self.current_offset += char_len;
            if c == '\n' {
                self.current_line += 1;
                self.current_column = 1;
            } else {
                self.current_column += 1;
            }
            Some(c)
        } else {
            None
        }
    }

    fn advance_if<F>(&mut self, predicate: F) -> Option<char>
    where
        F: FnOnce(char) -> bool,
    {
        if let Some(c) = self.peek() {
            if predicate(c) {
                return self.advance();
            }
        }
        None
    }

    fn advance_while<F>(&mut self, predicate: F) -> String
    where
        F: Fn(char) -> bool,
    {
        let mut s = String::new();
        while let Some(c) = self.peek() {
            if predicate(c) {
                s.push(self.advance().unwrap());
            } else {
                break;
            }
        }
        s
    }

    fn next_token(&mut self) -> Option<TokenType> {
        self.skip_whitespace();

        let start_offset = self.current_offset;
        let start_line = self.current_line;
        let start_column = self.current_column;

        let Some(c) = self.advance() else {
            return None; // EOF handled by lex function
        };

        let token_type = match c {
            '(' => TokenType::OpenParen,
            ')' => TokenType::CloseParen,
            '[' => {
                if self.advance_if(|c| c == '[').is_some() {
                    TokenType::DoubleOpenBracket
                } else {
                    TokenType::OpenBracket
                }
            }
            ']' => {
                if self.advance_if(|c| c == ']').is_some() {
                    TokenType::DoubleCloseBracket
                } else {
                    TokenType::CloseBracket
                }
            }
            '{' => TokenType::OpenBrace,
            '}' => TokenType::CloseBrace,
            ':' => TokenType::Colon,
            ';' => TokenType::SemiColon,
            ',' => TokenType::Comma,
            '.' => TokenType::Dot,
            '=' => {
                if self.advance_if(|c| c == '=').is_some() {
                    if self.advance_if(|c| c == '>').is_some() {
                        TokenType::ThickArrowRight
                    } else if self.advance_if(|c| c == '<').is_some() {
                        TokenType::ThickDoubleArrow // <==>
                    } else {
                        // Error: Malformed '==' or '==>'
                        // For now, treat as Equals, let parser handle
                        TokenType::Equals // Treat as single '=', let parser report if invalid
                    }
                } else {
                    TokenType::Equals
                }
            }
            '+' => TokenType::Plus,
            '-' => {
                // Handle arrows starting with '-'
                if self.peek() == Some('.') {
                    self.advance(); // consume '.'
                    if self.advance_if(|c| c == '>').is_some() {
                        TokenType::DottedArrowRight
                    } else {
                        // Error: Malformed '-.'
                        // Treat as minus, let parser handle
                        debug!("Malformed dotted arrow: '-.'");
                        TokenType::Minus
                    }
                } else if self.peek() == Some('-') {
                    self.advance(); // consume second '-'
                    if self.advance_if(|c| c == '>').is_some() {
                        TokenType::ArrowRight
                    } else if self.advance_if(|c| c == 'x').is_some() {
                        TokenType::CrossArrowRight
                    } else if self.advance_if(|c| c == 'o').is_some() {
                        TokenType::OpenCircleArrowRight
                    } else {
                        // Could be '---' for text
                        let _ = self.advance_while(|ch| ch == '-'); // Consume remaining dashes
                        // This needs more context for ArrowRightWithText, handle in parser for now
                        TokenType::ArrowRight // Simplified for now, parser will differentiate
                    }
                } else {
                    TokenType::Minus
                }
            }
            '*' => TokenType::Star,
            '#' => TokenType::Hash,
            '%' => {
                if self.peek() == Some('%') {
                    self.advance(); // consume second '%'
                    let comment_text = self.advance_while(|ch| ch != '\n');
                    TokenType::LineComment(comment_text.trim_end().to_string())
                } else {
                    TokenType::Percent
                }
            }
            '|' => TokenType::Pipe,
            '\\' => TokenType::Backslash,
            '<' => {
                if self.advance_if(|c| c == '=').is_some() {
                    if self.advance_if(|c| c == '=').is_some() {
                        if self.advance_if(|c| c == '>').is_some() {
                             TokenType::ThickDoubleArrow // <==>
                        } else {
                            TokenType::ThickArrowLeft // <==
                        }
                    } else {
                        // Error: Malformed '<='
                        debug!("Malformed arrow: '<='");
                        // For now, treat as single '<', let parser handle
                        TokenType::DoubleArrow // Simplified, parser will refine
                    }
                } else if self.advance_if(|c| c == '-').is_some() {
                    if self.peek() == Some('.') {
                        self.advance(); // consume '.'
                        TokenType::DottedArrowLeft
                    } else if self.advance_if(|c| c == '-').is_some() {
                        if self.advance_if(|c| c == '-').is_some() {
                            // Could be <--- text --->, let parser parse the text
                            TokenType::DoubleArrowWithText // Simplified
                        } else {
                            TokenType::ArrowLeft // <--
                        }
                    } else {
                        // Error: Malformed '<-'
                        debug!("Malformed arrow: '<-'");
                        TokenType::ArrowLeft // Simplified
                    }
                } else {
                    // This is problematic. '<' alone is not a valid token in Mermaid
                    // For now, treat as invalid character.
                    self.report_lexer_error(
                        ErrorCode::M001,
                        Span::new(
                            self.source_id,
                            start_offset,
                            self.current_offset,
                            start_line,
                            self.current_line,
                            start_column,
                            self.current_column,
                        ),
                        format!("Unexpected character '{}'. Mermaid arrows usually start with '-' or '<-'.", c),
                    );
                    return None; // Indicate an error, lexer will skip this char
                }
            }
            '"' | '\'' => {
                let quote_char = c;
                let mut content = String::new();
                let mut terminated = false;
                while let Some(next_char) = self.peek() {
                    if next_char == quote_char {
                        self.advance(); // consume closing quote
                        terminated = true;
                        break;
                    }
                    if next_char == '\n' {
                        // String literals cannot span multiple lines without explicit escaping
                        break;
                    }
                    content.push(self.advance().unwrap());
                }

                if !terminated {
                    self.report_lexer_error(
                        ErrorCode::M002,
                        Span::new(
                            self.source_id,
                            start_offset,
                            self.current_offset,
                            start_line,
                            self.current_line,
                            start_column,
                            self.current_column,
                        ),
                        format!("Unterminated string literal. Expected '{}'.", quote_char),
                    );
                    return None; // Indicate error
                }
                TokenType::StringLiteral(content)
            }
            c if c.is_ascii_digit() => {
                let num_str = c.to_string() + &self.advance_while(|ch| ch.is_ascii_digit() || ch == '.');
                // Basic check, full number validation could be a parser/validator task
                TokenType::Number(num_str)
            }
            c if c.is_alphabetic() => {
                let ident_str = c.to_string() + &self.advance_while(|ch| ch.is_alphanumeric() || ch == '_');
                match ident_str.as_str() {
                    // Keywords
                    "graph" => TokenType::Graph,
                    "flowchart" => TokenType::Flowchart,
                    "sequenceDiagram" => TokenType::SequenceDiagram,
                    "classDiagram" => TokenType::ClassDiagram,
                    "stateDiagram" => TokenType::StateDiagram,
                    "gantt" => TokenType::Gantt,
                    "pie" => TokenType::Pie,
                    "gitGraph" => TokenType::GitGraph,
                    "erDiagram" => TokenType::Ermaid, // Typo: ermaid -> erDiagram
                    "journey" => TokenType::Journey,
                    "C4Context" => TokenType::C4c,
                    "mindmap" => TokenType::Mindmap,
                    "timeline" => TokenType::Timeline,
                    "block" => TokenType::Block,
                    "sankey" => TokenType::SanKey,
                    "requirementDiagram" => TokenType::RequirementDiagram,

                    // Directions
                    "TD" | "TB" => TokenType::TD,
                    "BT" => TokenType::BT,
                    "LR" => TokenType::LR,
                    "RL" => TokenType::RL,
                    _ => TokenType::Identifier(ident_str),
                }
            }
            _ => {
                let span = Span::new(
                    self.source_id,
                    start_offset,
                    self.current_offset,
                    start_line,
                    self.current_line,
                    start_column,
                    self.current_column,
                );
                self.report_lexer_error(
                    ErrorCode::M001,
                    span,
                    format!("Unexpected character '{}'", c),
                );
                return None; // Indicate error, lexer will skip this char
            }
        };
        Some(token_type)
    }

    fn skip_whitespace(&mut self) {
        let _ = self.advance_while(|c| c.is_whitespace() && c != '\n');
    }

    // Helper to report a lexer-specific diagnostic
    fn report_lexer_error(&mut self, code: ErrorCode, span: Span, message: String) {
        let diag = Diagnostic::error(code)
            .with_message(message)
            .with_primary_label(span, "unexpected character")
            .build();
        self.diagnostics.push(diag);
    }
}

// src/parser/parser.rs (partial, imports only)

use crate::lexer::token::{Token, TokenType};
use crate::parser::ast::*;
use crate::diagnostics::span::{Span, SourceId};
use crate::diagnostics::{Diagnostic, DiagnosticBuilder, error_codes::{ErrorCode, Severity}}; // New imports
use log::{debug, error, warn};

// ... rest of the file

// src/parser/parser.rs (partial, modify Parser struct and parse method)

pub struct Parser<'a> {
    tokens: &'a [Token],
    current: usize,
    source_id: SourceId,
    diagnostics: Vec<Diagnostic>, // Collect diagnostics
}

impl<'a> Parser<'a> {
    pub fn new(tokens: &'a [Token], source_id: SourceId) -> Self {
        Parser {
            tokens,
            current: 0,
            source_id,
            diagnostics: Vec::new(),
        }
    }

    /// Parses the entire stream of tokens into a Diagram AST,
    /// collecting all encountered diagnostics.
    pub fn parse(mut self) -> (Option<Diagram>, Vec<Diagnostic>) {
        debug!("Starting parsing...");

        // Ensure we have tokens to process
        if self.peek().token_type == TokenType::EOF && self.tokens.len() == 1 {
            let diag = Diagnostic::error(ErrorCode::M102)
                .with_message("Empty input: missing diagram type declaration")
                .with_primary_label(self.peek().span, "expected 'graph TD', 'sequenceDiagram', etc.")
                .build();
            self.diagnostics.push(diag);
            return (None, self.diagnostics);
        }

        let diagram = self.parse_diagram_declaration();
        let mut ast = match diagram {
            Ok(d) => d,
            Err(_) => {
                // Error already reported by parse_diagram_declaration
                // Attempt to synchronize by skipping until a new line or EOF
                self.synchronize();
                return (None, self.diagnostics);
            }
        };

        match &mut ast {
            Diagram::Flowchart(flowchart) => {
                while self.peek().token_type != TokenType::EOF {
                    match self.parse_flowchart_statement() {
                        Ok(Some(stmt)) => flowchart.statements.push(stmt),
                        Ok(None) => {
                            // Recovered, but no statement parsed (e.g., skip whitespace)
                            self.synchronize(); // Try to skip to next statement
                        },
                        Err(_) => {
                            // Error reported by parse_flowchart_statement, attempt to synchronize
                            self.synchronize();
                        }
                    }
                }
            }
            // ... handle other diagram types similarly if they have statements
            _ => {
                // For now, other diagram types might not have further statements parsed
                // We'll expand this as we implement more diagram types.
                debug!("Diagram type {:?} does not yet support further statement parsing.", ast);
            }
        }


        if !self.diagnostics.is_empty() && self.diagnostics.iter().any(|d| d.severity == Severity::Error) {
            (None, self.diagnostics) // If any errors, return None AST
        } else {
            (Some(ast), self.diagnostics) // Otherwise, return AST
        }
    }

    // Helper to report a parser-specific diagnostic
    fn report_parser_error(&mut self, code: ErrorCode, span: Span, message: String, primary_label_msg: Option<String>) {
        let mut builder = Diagnostic::error(code)
            .with_message(message);
        if let Some(label_msg) = primary_label_msg {
            builder = builder.with_primary_label(span, label_msg);
        } else {
            builder = builder.with_primary_label(span, self.peek().token_type.to_string());
        }
        if let Some(help_msg) = code.help() {
            builder = builder.with_help(help_msg);
        }
        self.diagnostics.push(builder.build());
    }

    // Helper to consume a token of expected type, reporting an error if not found.
    fn consume(&mut self, expected: TokenType, error_code: ErrorCode, error_msg: &str) -> Result<Token, Diagnostic> {
        let token = self.peek();
        if token.token_type == expected {
            Ok(self.advance())
        } else {
            let span = token.span;
            let current_token_type = token.token_type.to_string();
            let diag = Diagnostic::error(error_code)
                .with_message(format!("{}. Expected '{}', but found '{}'.", error_msg, expected.to_string(), current_token_type))
                .with_primary_label(span, format!("expected '{}'", expected.to_string()))
                .with_note(format!("The parser expected a '{}' here to continue parsing this construct.", expected.to_string()))
                .with_help(error_code.help().unwrap_or("Review Mermaid syntax for this section.").to_string())
                .build();
            self.diagnostics.push(diag.clone());
            Err(diag) // Return an Err to indicate failure in this specific parsing step
        }
    }

    // New method for error recovery: skip tokens until a likely synchronization point.
    fn synchronize(&mut self) {
        debug!("Attempting to synchronize parser after error.");
        self.advance(); // Consume the erroneous token

        while self.peek().token_type != TokenType::EOF {
            match self.previous().token_type {
                TokenType::SemiColon | TokenType::NewLine => return, // End of statement
                TokenType::OpenBrace | TokenType::OpenParen | TokenType::OpenBracket => return, // Start of a new block
                _ => {}
            }
            // Skip until we find a likely start of a new statement or block
            match self.peek().token_type {
                TokenType::Graph | TokenType::Flowchart | TokenType::SequenceDiagram | TokenType::ClassDiagram => return,
                TokenType::NewLine => {
                    self.advance(); // Consume newline and return
                    return;
                },
                _ => {
                    self.advance();
                }
            }
        }
    }

    // --- Modify existing parsing methods to use diagnostics ---

    fn parse_diagram_declaration(&mut self) -> Result<Diagram, Diagnostic> {
        let start_token = self.peek();
        let diagram_type = match start_token.token_type {
            TokenType::Graph | TokenType::Flowchart => {
                self.advance(); // consume 'graph' or 'flowchart'
                let direction_token = self.peek();
                let direction = match direction_token.token_type {
                    TokenType::TD | TokenType::TB => { self.advance(); FlowchartDirection::TD },
                    TokenType::BT => { self.advance(); FlowchartDirection::BT },
                    TokenType::LR => { self.advance(); FlowchartDirection::LR },
                    TokenType::RL => { self.advance(); FlowchartDirection::RL },
                    _ => {
                        let diag = Diagnostic::error(ErrorCode::M102)
                            .with_message(format!("Missing or invalid direction for '{}' diagram.", start_token.token_type.to_string()))
                            .with_primary_label(direction_token.span, "expected 'TD', 'BT', 'LR', or 'RL'")
                            .with_help(ErrorCode::M102.help().unwrap().to_string())
                            .build();
                        self.diagnostics.push(diag.clone());
                        // Attempt to recover by assuming TD and continuing
                        FlowchartDirection::TD
                    }
                };
                return Ok(Diagram::Flowchart(FlowchartDiagram { direction, statements: Vec::new() }));
            }
            TokenType::SequenceDiagram => {
                self.advance();
                return Ok(Diagram::Sequence(SequenceDiagram { statements: Vec::new() }));
            }
            TokenType::ClassDiagram => {
                self.advance();
                return Ok(Diagram::Class(ClassDiagram { statements: Vec::new() }));
            }
            _ => {
                let diag = Diagnostic::error(ErrorCode::M102)
                    .with_message(format!("Missing or invalid diagram type declaration. Found '{}'.", start_token.token_type.to_string()))
                    .with_primary_label(start_token.span, "expected 'graph', 'flowchart', 'sequenceDiagram', etc.")
                    .with_help(ErrorCode::M102.help().unwrap().to_string())
                    .build();
                self.diagnostics.push(diag.clone());
                return Err(diag);
            }
        };
    }

    // Example of a statement parser that might return multiple diagnostics or an error
    fn parse_flowchart_statement(&mut self) -> Result<Option<FlowchartStatement>, Diagnostic> {
        self.skip_newlines_and_whitespace(); // Helper to advance past newlines/whitespace

        if self.peek().token_type == TokenType::EOF {
            return Ok(None);
        }

        let start_token = self.peek();
        if start_token.token_type == TokenType::Identifier || start_token.token_type == TokenType::StringLiteral {
            let node_id_token = self.advance();
            let node_id = match &node_id_token.token_type {
                TokenType::Identifier(s) => s.clone(),
                TokenType::StringLiteral(s) => s.clone(),
                _ => unreachable!(), // Handled by if condition
            };

            // Node definition or edge
            if self.peek().token_type == TokenType::OpenBracket {
                // Node definition: A[Label]
                self.advance(); // consume '['
                let label_token = self.peek();
                let label = match &label_token.token_type {
                    TokenType::Identifier(s) => { self.advance(); s.clone() },
                    TokenType::StringLiteral(s) => { self.advance(); s.clone() },
                    _ => {
                        let diag = Diagnostic::error(ErrorCode::M103)
                            .with_message("Expected a node label inside brackets.")
                            .with_primary_label(label_token.span, "expected label (identifier or string)")
                            .with_help(ErrorCode::M103.help().unwrap().to_string())
                            .build();
                        self.diagnostics.push(diag);
                        // Attempt to recover by using a dummy label
                        "MISSING_LABEL".to_string()
                    }
                };
                match self.consume(TokenType::CloseBracket, ErrorCode::M101, "Expected ']' to close node label.") {
                    Ok(_) => { /* OK */ },
                    Err(_) => { /* Error already reported */ }
                }
                return Ok(Some(FlowchartStatement::Node(Node { id: node_id, label: Some(label) })));
            } else if self.is_arrow_token(&self.peek().token_type) {
                // Edge definition: A --> B
                let arrow_token = self.advance(); // consume arrow
                let end_node_token = self.peek();
                let end_node_id = match &end_node_token.token_type {
                    TokenType::Identifier(s) => { self.advance(); s.clone() },
                    TokenType::StringLiteral(s) => { self.advance(); s.clone() },
                    _ => {
                        let diag = Diagnostic::error(ErrorCode::M104)
                            .with_message("Expected an end node identifier after arrow.")
                            .with_primary_label(end_node_token.span, "expected node ID")
                            .with_help(ErrorCode::M104.help().unwrap().to_string())
                            .build();
                        self.diagnostics.push(diag);
                        // Attempt to recover by using a dummy node
                        "UNDEFINED_NODE".to_string()
                    }
                };
                return Ok(Some(FlowchartStatement::Edge(Edge {
                    from: node_id,
                    to: end_node_id,
                    arrow_type: self.map_token_to_arrow_type(&arrow_token.token_type),
                    label: None, // For now, no labels on edges
                })));
            } else {
                // Just a node definition without a label or an edge
                return Ok(Some(FlowchartStatement::Node(Node { id: node_id, label: None })));
            }
        }

        // If we reach here, it's an unexpected token at the start of a statement
        let diag = Diagnostic::error(ErrorCode::M100)
            .with_message(format!("Unexpected token '{}' at the start of a statement.", start_token.token_type.to_string()))
            .with_primary_label(start_token.span, "unexpected token")
            .with_help("Expected a node ID, an edge, or a new diagram declaration. This might indicate a syntax error or a missing semicolon.")
            .build();
        self.diagnostics.push(diag);
        Err(diag) // Indicate error, parser will synchronize
    }

    fn skip_newlines_and_whitespace(&mut self) {
        while self.peek().token_type == TokenType::NewLine || self.peek().token_type == TokenType::Whitespace {
            self.advance();
        }
    }

    // ... other helper methods like is_arrow_token, map_token_to_arrow_type
    fn is_arrow_token(&self, token_type: &TokenType) -> bool {
        matches!(
            token_type,
            TokenType::ArrowRight
                | TokenType::ArrowLeft
                | TokenType::DoubleArrow
                | TokenType::ThickArrowRight
                | TokenType::ThickArrowLeft
                | TokenType::ThickDoubleArrow
                | TokenType::DottedArrowRight
                | TokenType::DottedArrowLeft
                | TokenType::DottedDoubleArrow
                | TokenType::CrossArrowRight
                | TokenType::OpenCircleArrowRight
                | TokenType::ArrowRightWithText
                | TokenType::ArrowLeftWithText
                | TokenType::DoubleArrowWithText
                | TokenType::ThickArrowRightWithText
                | TokenType::DottedArrowRightWithText
        )
    }

    fn map_token_to_arrow_type(&self, token_type: &TokenType) -> ArrowType {
        match token_type {
            TokenType::ArrowRight => ArrowType::Solid,
            TokenType::ArrowLeft => ArrowType::SolidLeft,
            TokenType::DoubleArrow => ArrowType::SolidDouble,
            TokenType::ThickArrowRight => ArrowType::Thick,
            TokenType::ThickArrowLeft => ArrowType::ThickLeft,
            TokenType::ThickDoubleArrow => ArrowType::ThickDouble,
            TokenType::DottedArrowRight => ArrowType::Dotted,
            TokenType::DottedArrowLeft => ArrowType::DottedLeft,
            TokenType::DottedDoubleArrow => ArrowType::DottedDouble,
            TokenType::CrossArrowRight => ArrowType::Cross,
            TokenType::OpenCircleArrowRight => ArrowType::OpenCircle,
            // For now, these are simplified, actual text parsing will happen here later
            TokenType::ArrowRightWithText => ArrowType::Solid,
            TokenType::ArrowLeftWithText => ArrowType::SolidLeft,
            TokenType::DoubleArrowWithText => ArrowType::SolidDouble,
            TokenType::ThickArrowRightWithText => ArrowType::Thick,
            TokenType::DottedArrowRightWithText => ArrowType::Dotted,
            _ => {
                warn!("Unknown arrow token type encountered: {:?}", token_type);
                ArrowType::Solid // Default to solid arrow for now
            }
        }
    }

    fn peek(&self) -> &Token {
        self.tokens.get(self.current).unwrap_or_else(|| {
            // Should not happen if EOF is always the last token
            &self.tokens[self.tokens.len() - 1]
        })
    }

    fn previous(&self) -> &Token {
        self.tokens.get(self.current - 1).unwrap_or_else(|| {
            // Should not happen if current is > 0
            &self.tokens[0]
        })
    }

    fn advance(&mut self) -> Token {
        if self.current < self.tokens.len() {
            self.current += 1;
        }
        self.previous().clone()
    }
}