Chapter 2: Designing the Lexer: Tokenization of Mermaid Syntax

// src/lexer/mod.rs

pub mod error;
pub mod lexer;
pub mod token;

// Re-export common types for easier access
pub use error::LexerError;
pub use lexer::Lexer;
pub use token::{Span, Token, TokenKind};

// src/lexer/token.rs

use std::fmt;

/// Represents a span of text in the source code.
/// Stores the start and end byte indices, and the literal string slice.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct Span<'a> {
    pub start: usize,
    pub end: usize,
    pub literal: &'a str,
}

impl<'a> Span<'a> {
    /// Creates a new `Span`.
    pub fn new(start: usize, end: usize, literal: &'a str) -> Self {
        Self { start, end, literal }
    }

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

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

impl<'a> fmt::Display for Span<'a> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "({}, {}: '{}')", self.start, self.end, self.literal)
    }
}

/// Represents the kind of a token.
/// This enum will grow as we support more Mermaid syntax.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum TokenKind {
    // Keywords
    Graph, Flowchart, SequenceDiagram, ClassDiagram, StateDiagram, ErDiagram,
    Gantt, Pie, GitGraph, Journey, C4Context, C4Container, C4Component,
    C4Dynamic, Mindmap, Timeline, QuadrantChart, RequirementDiagram,
    Blockdiag, Zenuml, EntityRelationship, Info, Config, Theme, Direction,
    Subgraph, Link, Click, Call, Activate, Deactivate, Alt, Opt, Loop, Par,
    Break, Critical, Note, Rect, Box, Actor, Participant, As, Title, Section,
    End, Class, Enum, Interface, Abstract, Static, Public, Private, Protected,
    Package, Namespace, Implements, Extends, Aggregation, Composition,
    Association, Dependency, Generalization, Realization, Style, ClassDef,
    LinkStyle, Fill, Stroke, StrokeWidth, FontSize, FontFamily, TextAlign,
    Left, Right, Up, Down, Fork, Join, Choice, State, Entry, Exit, On, Off,
    Hide, Show, LoopKeyword,
    // Note: Some keywords might overlap with identifiers, lexer must prioritize keywords.
    // E.g., `loop` can be a keyword or a node ID. Context matters for parser.

    // Directions
    LR, RL, TB, BT, TD,

    // Operators and Punctuation
    OpenParen,      // (
    CloseParen,     // )
    OpenBracket,    // [
    CloseBracket,   // ]
    OpenBrace,      // {
    CloseBrace,     // }
    LessThan,       // <
    GreaterThan,    // >
    Plus,           // +
    Minus,          // -
    Star,           // *
    Slash,          // /
    Equal,          // =
    Tilde,          // ~
    Colon,          // :
    Comma,          // ,
    Dot,            // .
    Underscore,     // _
    Pipe,           // |
    Hash,           // #
    Bang,           // !
    At,             // @
    Dollar,         // $
    Percent,        // %
    Caret,          // ^
    Ampersand,      // &
    Backtick,       // ` (for quoted strings)

    // Arrows and Connectors
    ArrowNormal,    // -->, ->
    ArrowOpen,      // -->, ->
    ArrowCross,     // --x, -x
    ArrowCircle,    // --o, -o
    ArrowDouble,    // <--, <-
    ArrowDoubleOpen,// <-->
    ArrowDoubleCross,// <-->x
    ArrowDoubleCircle,// <-->o
    ArrowThick,     // ===>, ==>
    ArrowDotted,    // -.->, .->
    ArrowThickDotted, // -=.=>

    // Special
    String,         // "..." or '...' or `...` (raw value)
    Identifier,     // Regular unquoted identifier (e.g., NodeA, my_node)
    Number,         // e.g., 123, 0.5
    Comment,        // %% single line comment
    Directive,      // %%{ ... }%% multiline directive/comment
    Whitespace,     // For internal use, usually ignored by parser
    Newline,        // Significant whitespace for line breaks

    // End of File
    Eof,

    // Error token for lexing failures
    Error,
}

impl TokenKind {
    /// Checks if the given string is a Mermaid keyword.
    pub fn from_keyword(s: &str) -> Option<TokenKind> {
        match s {
            "graph" => Some(TokenKind::Graph),
            "flowchart" => Some(TokenKind::Flowchart),
            "sequenceDiagram" => Some(TokenKind::SequenceDiagram),
            "classDiagram" => Some(TokenKind::ClassDiagram),
            "stateDiagram" => Some(TokenKind::StateDiagram),
            "erDiagram" => Some(TokenKind::ErDiagram),
            "gantt" => Some(TokenKind::Gantt),
            "pie" => Some(TokenKind::Pie),
            "gitGraph" => Some(TokenKind::GitGraph),
            "journey" => Some(TokenKind::Journey),
            "C4Context" => Some(TokenKind::C4Context),
            "C4Container" => Some(TokenKind::C4Container),
            "C4Component" => Some(TokenKind::C4Component),
            "C4Dynamic" => Some(TokenKind::C4Dynamic),
            "mindmap" => Some(TokenKind::Mindmap),
            "timeline" => Some(TokenKind::Timeline),
            "quadrantChart" => Some(TokenKind::QuadrantChart),
            "requirementDiagram" => Some(TokenKind::RequirementDiagram),
            "blockdiag" => Some(TokenKind::Blockdiag),
            "zenuml" => Some(TokenKind::Zenuml),
            "entityRelationship" => Some(TokenKind::EntityRelationship),
            "info" => Some(TokenKind::Info),
            "config" => Some(TokenKind::Config),
            "theme" => Some(TokenKind::Theme),
            "direction" => Some(TokenKind::Direction),
            "subgraph" => Some(TokenKind::Subgraph),
            "link" => Some(TokenKind::Link),
            "click" => Some(TokenKind::Click),
            "call" => Some(TokenKind::Call),
            "activate" => Some(TokenKind::Activate),
            "deactivate" => Some(TokenKind::Deactivate),
            "alt" => Some(TokenKind::Alt),
            "opt" => Some(TokenKind::Opt),
            "loop" => Some(TokenKind::LoopKeyword), // Differentiate from 'loop' as an identifier
            "par" => Some(TokenKind::Par),
            "break" => Some(TokenKind::Break),
            "critical" => Some(TokenKind::Critical),
            "note" => Some(TokenKind::Note),
            "rect" => Some(TokenKind::Rect),
            "box" => Some(TokenKind::Box),
            "actor" => Some(TokenKind::Actor),
            "participant" => Some(TokenKind::Participant),
            "as" => Some(TokenKind::As),
            "title" => Some(TokenKind::Title),
            "section" => Some(TokenKind::Section),
            "end" => Some(TokenKind::End),
            "class" => Some(TokenKind::Class),
            "enum" => Some(TokenKind::Enum),
            "interface" => Some(TokenKind::Interface),
            "abstract" => Some(TokenKind::Abstract),
            "static" => Some(TokenKind::Static),
            "public" => Some(TokenKind::Public),
            "private" => Some(TokenKind::Private),
            "protected" => Some(TokenKind::Protected),
            "package" => Some(TokenKind::Package),
            "namespace" => Some(TokenKind::Namespace),
            "implements" => Some(TokenKind::Implements),
            "extends" => Some(TokenKind::Extends),
            "aggregation" => Some(TokenKind::Aggregation),
            "composition" => Some(TokenKind::Composition),
            "association" => Some(TokenKind::Association),
            "dependency" => Some(TokenKind::Dependency),
            "generalization" => Some(TokenKind::Generalization),
            "realization" => Some(TokenKind::Realization),
            "style" => Some(TokenKind::Style),
            "classDef" => Some(TokenKind::ClassDef),
            "linkStyle" => Some(TokenKind::LinkStyle),
            "fill" => Some(TokenKind::Fill),
            "stroke" => Some(TokenKind::Stroke),
            "stroke-width" => Some(TokenKind::StrokeWidth),
            "font-size" => Some(TokenKind::FontSize),
            "font-family" => Some(TokenKind::FontFamily),
            "text-align" => Some(TokenKind::TextAlign),
            "left" => Some(TokenKind::Left),
            "right" => Some(TokenKind::Right),
            "up" => Some(TokenKind::Up),
            "down" => Some(TokenKind::Down),
            "fork" => Some(TokenKind::Fork),
            "join" => Some(TokenKind::Join),
            "choice" => Some(TokenKind::Choice),
            "state" => Some(TokenKind::State),
            "entry" => Some(TokenKind::Entry),
            "exit" => Some(TokenKind::Exit),
            "on" => Some(TokenKind::On),
            "off" => Some(TokenKind::Off),
            "hide" => Some(TokenKind::Hide),
            "show" => Some(TokenKind::Show),
            // Directions
            "LR" => Some(TokenKind::LR),
            "RL" => Some(TokenKind::RL),
            "TB" => Some(TokenKind::TB),
            "BT" => Some(TokenKind::BT),
            "TD" => Some(TokenKind::TD),
            _ => None,
        }
    }
}

impl fmt::Display for TokenKind {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{:?}", self) // Default to debug output for now
    }
}

/// Represents a single token found by the lexer.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct Token<'a> {
    pub kind: TokenKind,
    pub span: Span<'a>,
}

impl<'a> Token<'a> {
    /// Creates a new `Token`.
    pub fn new(kind: TokenKind, span: Span<'a>) -> Self {
        Self { kind, span }
    }
}

impl<'a> fmt::Display for Token<'a> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "Token {{ kind: {}, span: {} }}", self.kind, self.span)
    }
}

// src/lexer/error.rs

use crate::lexer::Span; // Import Span from our lexer module
use std::fmt;
use std::error::Error;

/// Represents an error encountered during lexing.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum LexerError<'a> {
    /// An unexpected or unrecognized character was found.
    InvalidCharacter {
        span: Span<'a>,
        found: char,
    },
    /// A string literal was started but not properly terminated.
    UnterminatedString {
        span: Span<'a>,
        missing_delimiter: char,
    },
    /// An invalid escape sequence was found within a string.
    InvalidEscapeSequence {
        span: Span<'a>,
        sequence: &'a str,
    },
    /// An unexpected end of file was reached while expecting more characters.
    UnexpectedEof {
        span: Span<'a>,
        expected: &'static str,
    },
    /// A generic lexer error, used for less specific issues.
    GenericError {
        span: Span<'a>,
        message: String,
    },
}

impl<'a> fmt::Display for LexerError<'a> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            LexerError::InvalidCharacter { span, found } => {
                write!(f, "Error: Invalid character '{}' at {}", found, span)
            }
            LexerError::UnterminatedString { span, missing_delimiter } => {
                write!(f, "Error: Unterminated string starting at {}. Expected '{}'", span, missing_delimiter)
            }
            LexerError::InvalidEscapeSequence { span, sequence } => {
                write!(f, "Error: Invalid escape sequence '{}' at {}", sequence, span)
            }
            LexerError::UnexpectedEof { span, expected } => {
                write!(f, "Error: Unexpected end of file at {}. Expected {}", span, expected)
            }
            LexerError::GenericError { span, message } => {
                write!(f, "Error: {} at {}", message, span)
            }
        }
    }
}

impl<'a> Error for LexerError<'a> {}

// src/lexer/lexer.rs

use crate::lexer::error::LexerError;
use crate::lexer::token::{Span, Token, TokenKind};
use std::str::Chars;
use std::iter::Peekable;

/// The Lexer is responsible for converting a Mermaid source string into a stream of tokens.
pub struct Lexer<'a> {
    source: &'a str,
    chars: Peekable<Chars<'a>>,
    current_pos: usize, // Current byte offset in the source string
    start_pos: usize,   // Start byte offset of the current token being scanned
}

impl<'a> Lexer<'a> {
    /// Creates a new `Lexer` instance from the given Mermaid source code.
    pub fn new(source: &'a str) -> Self {
        Self {
            source,
            chars: source.chars().peekable(),
            current_pos: 0,
            start_pos: 0,
        }
    }

    /// Peeks at the next character without consuming it.
    fn peek(&mut self) -> Option<char> {
        self.chars.peek().copied()
    }

    /// Consumes and returns the next character, advancing the current position.
    fn advance(&mut self) -> Option<char> {
        let c = self.chars.next();
        if let Some(ch) = c {
            self.current_pos += ch.len_utf8();
        }
        c
    }

    /// Consumes the next character if it matches the expected character.
    fn consume_if_match(&mut self, expected: char) -> bool {
        if self.peek() == Some(expected) {
            self.advance();
            true
        } else {
            false
        }
    }

    /// Creates a `Span` for the token just scanned.
    fn create_span(&self) -> Span<'a> {
        Span::new(
            self.start_pos,
            self.current_pos,
            &self.source[self.start_pos..self.current_pos],
        )
    }

    /// Skips whitespace characters.
    fn skip_whitespace(&mut self) {
        while self.peek().map_or(false, |c| c.is_ascii_whitespace() && c != '\n') {
            self.advance();
        }
    }

    /// Reads an identifier or keyword.
    fn read_identifier_or_keyword(&mut self) -> TokenKind {
        self.take_while(|c| c.is_alphanumeric() || c == '_' || c == '-')
            .map_or(TokenKind::Error, |literal| {
                if let Some(kind) = TokenKind::from_keyword(literal) {
                    kind
                } else {
                    TokenKind::Identifier
                }
            })
    }

    /// Reads a string literal (single, double, or backtick quoted).
    fn read_string(&mut self, delimiter: char) -> Result<TokenKind, LexerError<'a>> {
        let start_string_pos = self.current_pos;
        self.advance(); // Consume the opening delimiter

        let mut escaped = false;
        while let Some(c) = self.peek() {
            if c == '\\' && !escaped {
                escaped = true;
                self.advance();
                continue;
            }

            if c == delimiter && !escaped {
                self.advance(); // Consume the closing delimiter
                return Ok(TokenKind::String);
            }

            if c == '\n' && !escaped {
                // Unterminated string on newline
                return Err(LexerError::UnterminatedString {
                    span: Span::new(self.start_pos, self.current_pos, &self.source[self.start_pos..self.current_pos]),
                    missing_delimiter: delimiter,
                });
            }

            escaped = false;
            self.advance();
        }

        // Reached EOF without closing delimiter
        Err(LexerError::UnterminatedString {
            span: Span::new(self.start_pos, start_string_pos, &self.source[self.start_pos..start_string_pos]),
            missing_delimiter: delimiter,
        })
    }

    /// Reads a single-line comment (%% ...).
    fn read_comment(&mut self) -> TokenKind {
        self.advance(); // Consume first '%'
        self.advance(); // Consume second '%'
        self.take_while(|c| c != '\n'); // Consume until newline or EOF
        TokenKind::Comment
    }

    /// Reads a multi-line directive (%%{ ... }%%).
    fn read_directive(&mut self) -> Result<TokenKind, LexerError<'a>> {
        self.advance(); // Consume first '%'
        self.advance(); // Consume second '%'
        self.advance(); // Consume '{'

        let mut balance = 1;
        while let Some(c) = self.advance() {
            match c {
                '{' => balance += 1,
                '}' => balance -= 1,
                _ => {}
            }
            if balance == 0 {
                // Check for closing %%
                if self.peek() == Some('%') {
                    self.advance();
                    if self.peek() == Some('%') {
                        self.advance();
                        return Ok(TokenKind::Directive);
                    }
                }
                // If we get here, it means '}' was found but not followed by '%%'
                // We'll treat this as an unterminated directive for now, or let parser handle it
                // For a strict lexer, let's just make sure we consume the whole thing.
                return Ok(TokenKind::Directive); // Or error if not followed by %%
            }
        }
        // Reached EOF without closing '}'
        Err(LexerError::UnterminatedString {
            span: self.create_span(),
            missing_delimiter: '}',
        })
    }

    /// Reads various arrow types. This is a complex part.
    /// Priority matters: longer arrows first.
    fn read_arrow(&mut self) -> TokenKind {
        let first_char = self.advance().unwrap(); // Already checked to be part of an arrow
        match first_char {
            '-' => {
                if self.consume_if_match('.') {
                    if self.consume_if_match('-') {
                        if self.consume_if_match('>') { return TokenKind::ArrowDotted; } // -.->
                    }
                } else if self.consume_if_match('-') {
                    if self.consume_if_match('>') { return TokenKind::ArrowNormal; } // -->
                    if self.consume_if_match('x') { return TokenKind::ArrowCross; } // --x
                    if self.consume_if_match('o') { return TokenKind::ArrowCircle; } // --o
                    return TokenKind::Minus; // Fallback for '--' or just '-'
                } else if self.consume_if_match('>') {
                    return TokenKind::ArrowOpen; // ->
                }
                TokenKind::Minus // Single '-'
            },
            '=' => {
                if self.consume_if_match('=') {
                    if self.consume_if_match('>') { return TokenKind::ArrowThick; } // ==>
                    if self.consume_if_match('.') {
                        if self.consume_if_match('=') {
                            if self.consume_if_match('>') { return TokenKind::ArrowThickDotted; } // -=.=>
                        }
                    }
                }
                TokenKind::Equal // Single '='
            },
            '<' => {
                if self.consume_if_match('-') {
                    if self.consume_if_match('-') {
                        if self.consume_if_match('>') { return TokenKind::ArrowDoubleOpen; } // <-->
                        if self.consume_if_match('x') { return TokenKind::ArrowDoubleCross; } // <--x
                        if self.consume_if_match('o') { return TokenKind::ArrowDoubleCircle; } // <--o
                        return TokenKind::ArrowDouble; // <---
                    } else if self.consume_if_match('>') {
                        return TokenKind::ArrowDouble; // <->
                    }
                    return TokenKind::ArrowDouble; // <-
                }
                TokenKind::LessThan // Single '<'
            },
            _ => unreachable!("read_arrow called on non-arrow char"),
        }
    }

    /// Reads a number literal. (Simple integer for now, can be extended for floats).
    fn read_number(&mut self) -> TokenKind {
        self.take_while(|c| c.is_ascii_digit());
        // TODO: Extend for floats (e.g., 123.45)
        TokenKind::Number
    }

    /// Consumes characters while the predicate `p` is true, returning the consumed slice.
    fn take_while<P>(&mut self, p: P) -> Option<&'a str>
    where
        P: Fn(char) -> bool,
    {
        let start = self.current_pos;
        while let Some(c) = self.peek() {
            if p(c) {
                self.advance();
            } else {
                break;
            }
        }
        if self.current_pos > start {
            Some(&self.source[start..self.current_pos])
        } else {
            None
        }
    }

    /// The main lexing function: determines the next token.
    fn next_token(&mut self) -> Result<Token<'a>, LexerError<'a>> {
        self.start_pos = self.current_pos; // Mark the start of the new token

        // Skip leading whitespace (but not newlines, as they can be significant)
        self.skip_whitespace();
        if self.current_pos > self.start_pos {
            // If we skipped whitespace, create a whitespace token
            let span = Span::new(self.start_pos, self.current_pos, &self.source[self.start_pos..self.current_pos]);
            // Reset start_pos for the *next* actual token
            self.start_pos = self.current_pos;
            return Ok(Token::new(TokenKind::Whitespace, span));
        }

        let current_char = match self.advance() {
            Some(c) => c,
            None => {
                // If we're at EOF, check if it's already the start_pos, if so, return Eof
                // Otherwise, it means we advanced past the last character and hit EOF.
                // In either case, current_pos is accurate for span, start_pos might need adjustment
                let span = Span::new(self.current_pos, self.current_pos, "");
                return Ok(Token::new(TokenKind::Eof, span));
            }
        };

        let kind = match current_char {
            '\n' => TokenKind::Newline,
            '(' => TokenKind::OpenParen,
            ')' => TokenKind::CloseParen,
            '[' => TokenKind::OpenBracket,
            ']' => TokenKind::CloseBracket,
            '{' => {
                // Potentially a directive if preceded by %%
                // This specific check might be better handled by backtracking or a state machine
                // For simplicity, we'll assume '{' is always OpenBrace unless explicitly handled
                // by read_directive *before* this match.
                // This means '%%{' must be checked first.
                TokenKind::OpenBrace
            },
            '}' => TokenKind::CloseBrace,
            '<' => self.read_arrow(), // Can be part of <--, <-->
            '>' => TokenKind::GreaterThan,
            '+' => TokenKind::Plus,
            '*' => TokenKind::Star,
            '/' => TokenKind::Slash,
            ':' => TokenKind::Colon,
            ',' => TokenKind::Comma,
            '.' => {
                if self.peek() == Some('-') {
                    // Potentially part of a dotted arrow like .->
                    // Backtrack or handle within read_arrow if called
                    // For now, treat as Dot, and let read_arrow handle the full sequence
                    // This implies the lexer has to be smart about multi-character lookaheads.
                    // For simplicity, we'll let read_arrow handle leading chars like '-' or '='
                    // and 'dot' will be a standalone operator here.
                }
                TokenKind::Dot
            },
            '=' => self.read_arrow(), // Can be part of ===>, ==>
            '~' => TokenKind::Tilde,
            '|' => TokenKind::Pipe,
            '#' => TokenKind::Hash,
            '!' => TokenKind::Bang,
            '@' => TokenKind::At,
            '$' => TokenKind::Dollar,
            '%' => {
                if self.peek() == Some('%') {
                    // It's either a comment (%%) or a directive (%%{)
                    self.advance(); // Consume the second '%'
                    if self.peek() == Some('{') {
                        // This is a directive
                        self.start_pos = self.current_pos - 2; // Adjust start_pos to include both '%%'
                        return self.read_directive().map(|k| Token::new(k, self.create_span()));
                    } else {
                        // This is a single-line comment
                        self.start_pos = self.current_pos - 2; // Adjust start_pos to include both '%%'
                        self.read_comment();
                        TokenKind::Comment
                    }
                } else {
                    TokenKind::Percent // Single '%'
                }
            },
            '^' => TokenKind::Caret,
            '&' => TokenKind::Ampersand,
            '`' => return self.read_string('`').map(|k| Token::new(k, self.create_span())),
            '\'' => return self.read_string('\'').map(|k| Token::new(k, self.create_span())),
            '"' => return self.read_string('"').map(|k| Token::new(k, self.create_span())),
            '-' => self.read_arrow(), // Can be part of -->, ---, -.->
            c if c.is_ascii_digit() => {
                // If it's a digit, read the whole number
                self.current_pos -= c.len_utf8(); // Backtrack, read_number will advance again
                self.read_number()
            },
            c if c.is_alphabetic() || c == '_' => {
                // If it's an alphabet or underscore, read identifier/keyword
                self.current_pos -= c.len_utf8(); // Backtrack, read_identifier_or_keyword will advance again
                self.read_identifier_or_keyword()
            },
            _ => {
                // Unrecognized character
                return Err(LexerError::InvalidCharacter {
                    span: self.create_span(),
                    found: current_char,
                });
            }
        };

        Ok(Token::new(kind, self.create_span()))
    }
}

// Implement the Iterator trait for Lexer
impl<'a> Iterator for Lexer<'a> {
    type Item = Result<Token<'a>, LexerError<'a>>;

    fn next(&mut self) -> Option<Self::Item> {
        // Repeatedly call next_token until a non-whitespace token or EOF is found,
        // or an error occurs.
        loop {
            match self.next_token() {
                Ok(token) => {
                    if token.kind == TokenKind::Eof {
                        // Only return Eof once
                        if token.span.is_empty() && self.current_pos == self.source.len() {
                            return None; // Already yielded Eof, or source is empty
                        } else {
                            // First time hitting EOF, return it.
                            return Some(Ok(token));
                        }
                    } else if token.kind == TokenKind::Whitespace {
                        // Skip whitespace and try again
                        self.start_pos = self.current_pos; // Ensure next token starts after whitespace
                        continue;
                    } else {
                        return Some(Ok(token));
                    }
                }
                Err(e) => return Some(Err(e)),
            }
        }
    }
}

// src/main.rs

mod lexer;

use lexer::{Lexer, Token, LexerError};
use std::fs;
use std::path::PathBuf;
use clap::Parser; // For CLI arguments, from Chapter 1

#[derive(Parser, Debug)]
#[command(author, version, about, long_about = None)]
struct Args {
    /// Path to the Mermaid file to process
    #[arg(short, long, value_name = "FILE")]
    file: Option<PathBuf>,

    /// Mermaid code directly as a string
    #[arg(short, long, value_name = "CODE")]
    code: Option<String>,

    /// Output mode: "lint", "fix", "strict"
    #[arg(short, long, default_value = "lint")]
    mode: String,
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let args = Args::parse();

    let mermaid_source = if let Some(file_path) = args.file {
        fs::read_to_string(&file_path)?
    } else if let Some(code_str) = args.code {
        code_str
    } else {
        eprintln!("Error: No input file or Mermaid code provided. Use --file or --code.");
        std::process::exit(1);
    };

    println!("--- Lexing Input ---");
    let lexer = Lexer::new(&mermaid_source);
    let mut tokens: Vec<Token> = Vec::new();
    let mut errors: Vec<LexerError> = Vec::new();

    for item in lexer {
        match item {
            Ok(token) => {
                println!("{}", token);
                tokens.push(token);
            }
            Err(e) => {
                eprintln!("{}", e);
                errors.push(e);
            }
        }
    }

    println!("\n--- Lexing Summary ---");
    println!("Total tokens: {}", tokens.len());
    if !errors.is_empty() {
        eprintln!("Lexing completed with {} errors.", errors.len());
        std::process::exit(1);
    } else {
        println!("Lexing completed successfully with no errors.");
    }

    Ok(())
}

// src/lexer/lexer.rs (continued)

#[cfg(test)]
mod tests {
    use super::*;

    /// Helper function to lex and collect tokens, ignoring whitespace.
    fn lex_and_collect(source: &str) -> Result<Vec<Token>, LexerError> {
        let lexer = Lexer::new(source);
        let mut tokens = Vec::new();
        for item in lexer {
            match item {
                Ok(token) => tokens.push(token),
                Err(e) => return Err(e),
            }
        }
        Ok(tokens)
    }

    /// Helper function to lex and collect token kinds.
    fn lex_and_collect_kinds(source: &str) -> Result<Vec<TokenKind>, LexerError> {
        let lexer = Lexer::new(source);
        let mut kinds = Vec::new();
        for item in lexer {
            match item {
                Ok(token) => kinds.push(token.kind),
                Err(e) => return Err(e),
            }
        }
        Ok(kinds)
    }

    #[test]
    fn test_empty_input() {
        let tokens = lex_and_collect_kinds("").unwrap();
        assert_eq!(tokens, vec![TokenKind::Eof]);
    }

    #[test]
    fn test_basic_flowchart() {
        let source = "graph TD\n    A[Node A] --> B(Node B)";
        let tokens = lex_and_collect_kinds(source).unwrap();
        assert_eq!(
            tokens,
            vec![
                TokenKind::Graph,
                TokenKind::TD,
                TokenKind::Newline,
                TokenKind::Identifier,
                TokenKind::OpenBracket,
                TokenKind::Identifier, // "Node A" (will be parsed as a string literal later)
                TokenKind::CloseBracket,
                TokenKind::ArrowNormal,
                TokenKind::Identifier,
                TokenKind::OpenParen,
                TokenKind::Identifier, // "Node B"
                TokenKind::CloseParen,
                TokenKind::Eof
            ]
        );
    }

    #[test]
    fn test_keywords_and_identifiers() {
        let source = "flowchart LR\nsubgraph MySubgraph\n    Node1 --> Node2\nend";
        let tokens = lex_and_collect_kinds(source).unwrap();
        assert_eq!(
            tokens,
            vec![
                TokenKind::Flowchart,
                TokenKind::LR,
                TokenKind::Newline,
                TokenKind::Subgraph,
                TokenKind::Identifier, // MySubgraph
                TokenKind::Newline,
                TokenKind::Identifier, // Node1
                TokenKind::ArrowNormal,
                TokenKind::Identifier, // Node2
                TokenKind::Newline,
                TokenKind::End,
                TokenKind::Eof
            ]
        );
    }

    #[test]
    fn test_all_arrow_types() {
        let source = "A --> B\nC --- D\nE ==> F\nG --x H\nI --o J\nK -.-> L\nM <--> N\nO ===> P\nQ -.=> R"; // Missing -=.=>
        let tokens = lex_and_collect_kinds(source).unwrap();
        assert!(tokens.contains(&TokenKind::ArrowNormal));
        assert!(tokens.contains(&TokenKind::Minus)); // --- will be Minus, Minus, Minus for now
        assert!(tokens.contains(&TokenKind::ArrowThick)); // ==>
        assert!(tokens.contains(&TokenKind::ArrowCross));
        assert!(tokens.contains(&TokenKind::ArrowCircle));
        assert!(tokens.contains(&TokenKind::ArrowDotted));
        assert!(tokens.contains(&TokenKind::ArrowDoubleOpen));
        assert!(tokens.contains(&TokenKind::ArrowThick)); // ===>
        assert!(tokens.contains(&TokenKind::ArrowThickDotted));
    }

    #[test]
    fn test_quoted_strings() {
        let source = "A[\"Hello World\"]\nB['Another String']\nC[`Backtick String`]";
        let tokens = lex_and_collect_kinds(source).unwrap();
        assert_eq!(
            tokens,
            vec![
                TokenKind::Identifier,
                TokenKind::OpenBracket,
                TokenKind::String,
                TokenKind::CloseBracket,
                TokenKind::Newline,
                TokenKind::Identifier,
                TokenKind::OpenBracket,
                TokenKind::String,
                TokenKind::CloseBracket,
                TokenKind::Newline,
                TokenKind::Identifier,
                TokenKind::OpenBracket,
                TokenKind::String,
                TokenKind::CloseBracket,
                TokenKind::Eof
            ]
        );

        let source_escaped = "A[\"Hello \\\"World\\\"\"]";
        let tokens_escaped = lex_and_collect(source_escaped).unwrap();
        assert_eq!(tokens_escaped[2].span.literal, "\"Hello \\\"World\\\"\"");
    }

    #[test]
    fn test_comments_and_directives() {
        let source = "%% This is a comment\ngraph TD\n%%{ init: { 'theme': 'dark' } }%%\n    A --> B";
        let tokens = lex_and_collect_kinds(source).unwrap();
        assert_eq!(
            tokens,
            vec![
                TokenKind::Comment,
                TokenKind::Newline,
                TokenKind::Graph,
                TokenKind::TD,
                TokenKind::Newline,
                TokenKind::Directive,
                TokenKind::Newline,
                TokenKind::Identifier,
                TokenKind::ArrowNormal,
                TokenKind::Identifier,
                TokenKind::Eof
            ]
        );
    }

    #[test]
    fn test_unterminated_string_error() {
        let source = "A[\"Unterminated string]";
        let error = lex_and_collect(source).unwrap_err();
        assert_eq!(
            error,
            LexerError::UnterminatedString {
                span: Span::new(3, 3, ""), // Span points to the beginning of the string literal
                missing_delimiter: '"',
            }
        );
    }

    #[test]
    fn test_invalid_character_error() {
        let source = "graph TD\n    A[Node A] $ B";
        let error = lex_and_collect(source).unwrap_err();
        assert_eq!(
            error,
            LexerError::InvalidCharacter {
                span: Span::new(19, 20, "$"),
                found: '$',
            }
        );
    }

    #[test]
    fn test_complex_operators() {
        let source = "A + B - C * D / E = F : G | H # I ! J @ K";
        let tokens = lex_and_collect_kinds(source).unwrap();
        assert_eq!(
            tokens,
            vec![
                TokenKind::Identifier, TokenKind::Plus, TokenKind::Identifier,
                TokenKind::Minus, TokenKind::Identifier, TokenKind::Star,
                TokenKind::Identifier, TokenKind::Slash, TokenKind::Identifier,
                TokenKind::Equal, TokenKind::Identifier, TokenKind::Colon,
                TokenKind::Identifier, TokenKind::Pipe, TokenKind::Identifier,
                TokenKind::Hash, TokenKind::Identifier, TokenKind::Bang,
                TokenKind::Identifier, TokenKind::At, TokenKind::Identifier,
                TokenKind::Eof
            ]
        );
    }

    #[test]
    fn test_numbers() {
        let source = "Node1(123) -- 456 --> Node2";
        let tokens = lex_and_collect_kinds(source).unwrap();
        assert_eq!(
            tokens,
            vec![
                TokenKind::Identifier, TokenKind::OpenParen, TokenKind::Number, TokenKind::CloseParen,
                TokenKind::Minus, TokenKind::Minus, // --
                TokenKind::Number,
                TokenKind::ArrowNormal, // -->
                TokenKind::Identifier,
                TokenKind::Eof
            ]
        );
    }
}