//! Lexer/Tokenizer for Mini-F# source code.
//!
//! This module implements the lexical analyzer that converts source text into
//! a stream of tokens. The lexer supports:
//!
//! - Literals: integers, floats, booleans, strings
//! - Keywords: let, rec, and, in, if, then, else, fun, true, false, async, return, yield
//! - Computation expression keywords: let!, do!, return!, yield!
//! - Identifiers: alphanumeric names starting with letter or underscore
//! - Operators: arithmetic, comparison, logical, cons (::)
//! - Punctuation: parentheses, arrows, commas, brackets, semicolons
//! - Array syntax: [|, |], <-, .
//! - Anonymous record syntax: {|, |}
//! - Comments: single-line (//) and multi-line ((* *)) with nesting support
//! - Position tracking for error reporting
//!
//! # Example
//!
//! ```rust
//! use fusabi_frontend::lexer::{Lexer, Token};
//!
//! let mut lexer = Lexer::new("let x = 42");
//! let tokens = lexer.tokenize().unwrap();
//! assert_eq!(tokens[0].token, Token::Let);
//! assert_eq!(tokens[1].token, Token::Ident("x".to_string()));
//! ```

pub use crate::span::{Position, Span};
use std::fmt;

/// Token types in Mini-F# source code.
#[derive(Debug, Clone, PartialEq)]
pub enum Token {
    // Literals
    /// Integer literal (e.g., 42, -10)
    Int(i64),
    /// Floating-point literal (e.g., 3.15, -0.5)
    Float(f64),
    /// Boolean literal (true or false)
    Bool(bool),
    /// String literal (e.g., "hello")
    String(String),

    // Identifiers
    /// Identifier (variable or function name)
    Ident(String),

    // Keywords
    /// let keyword
    Let,
    /// rec keyword (for recursive bindings)
    Rec,
    /// and keyword (for mutual recursion)
    AndKeyword,
    /// in keyword
    In,
    /// if keyword
    If,
    /// then keyword
    Then,
    /// else keyword
    Else,
    /// fun keyword
    Fun,
    /// match keyword
    Match,
    /// type keyword
    Type,
    /// with keyword (for record updates)
    With,
    /// of keyword (for discriminated unions)
    Of,
    /// open keyword (for importing modules)
    Open,
    /// module keyword (for module definitions)
    Module,
    /// do keyword (for side-effect expressions)
    Do,
    /// while keyword (for while loops)
    While,
    /// break keyword (for breaking out of loops)
    Break,
    /// continue keyword (for continuing to next iteration)
    Continue,
    /// async keyword (for async computation expressions)
    Async,
    /// return keyword (for returning from functions)
    Return,
    /// yield keyword (for generator functions)
    Yield,
    /// let! keyword (for computation expression binding)
    LetBang,
    /// do! keyword (for computation expression side-effect)
    DoBang,
    /// return! keyword (for computation expression return)
    ReturnBang,
    /// yield! keyword (for computation expression yield)
    YieldBang,

    // Operators
    /// + operator
    Plus,
    /// - operator
    Minus,
    /// * operator
    Star,
    /// / operator
    Slash,
    /// = operator
    Eq,
    /// == operator (equality comparison)
    EqEq,
    /// <> operator (inequality)
    Neq,
    /// < operator
    Lt,
    /// <= operator
    Lte,
    /// > operator
    Gt,
    /// >= operator
    Gte,
    /// && operator
    And,
    /// || operator
    Or,
    /// :: operator (cons)
    ColonColon,
    /// <- operator (array assignment)
    LArrow,
    /// |> operator (pipeline)
    PipeRight,
    /// ++ operator (string concatenation)
    PlusPlus,

    // Punctuation
    /// ( left parenthesis
    LParen,
    /// ) right parenthesis
    RParen,
    /// [ left bracket
    LBracket,
    /// ] right bracket
    RBracket,
    /// [| left array bracket
    LBracketPipe,
    /// |] right array bracket
    PipeRBracket,
    /// -> arrow
    Arrow,
    /// , comma
    Comma,
    /// ; semicolon
    Semicolon,
    /// . dot
    Dot,
    /// { left brace
    LBrace,
    /// } right brace
    RBrace,
    /// {| left anonymous record brace
    LBracePipe,
    /// |} right anonymous record brace
    PipeRBrace,
    /// : colon
    Colon,
    /// | pipe
    Pipe,
    /// _ underscore
    Underscore,

    // Directives
    /// #load directive (e.g., #load "path.fsx")
    LoadDirective(String),

    // Special
    /// End of file marker
    Eof,
}

impl fmt::Display for Token {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Token::Int(n) => write!(f, "Int({})", n),
            Token::Float(n) => write!(f, "Float({})", n),
            Token::Bool(b) => write!(f, "Bool({})", b),
            Token::String(s) => write!(f, "String(\"{}\")", s),
            Token::Ident(s) => write!(f, "Ident({})", s),
            Token::Let => write!(f, "let"),
            Token::In => write!(f, "in"),
            Token::Rec => write!(f, "rec"),
            Token::AndKeyword => write!(f, "and"),
            Token::If => write!(f, "if"),
            Token::Then => write!(f, "then"),
            Token::Else => write!(f, "else"),
            Token::Fun => write!(f, "fun"),
            Token::Match => write!(f, "match"),
            Token::Type => write!(f, "type"),
            Token::With => write!(f, "with"),
            Token::Plus => write!(f, "+"),
            Token::Minus => write!(f, "-"),
            Token::Star => write!(f, "*"),
            Token::Slash => write!(f, "/"),
            Token::Eq => write!(f, "="),
            Token::EqEq => write!(f, "=="),
            Token::Neq => write!(f, "<>"),
            Token::Lt => write!(f, "<"),
            Token::Lte => write!(f, "<="),
            Token::Gt => write!(f, ">"),
            Token::Gte => write!(f, ">="),
            Token::And => write!(f, "&&"),
            Token::Or => write!(f, "||"),
            Token::ColonColon => write!(f, "::"),
            Token::LArrow => write!(f, "<-"),
            Token::PipeRight => write!(f, "|>"),
            Token::PlusPlus => write!(f, "++"),
            Token::LParen => write!(f, "("),
            Token::RParen => write!(f, ")"),
            Token::LBracket => write!(f, "["),
            Token::RBracket => write!(f, "]"),
            Token::LBracketPipe => write!(f, "[|"),
            Token::PipeRBracket => write!(f, "|]"),
            Token::Arrow => write!(f, "->"),
            Token::Comma => write!(f, ","),
            Token::Semicolon => write!(f, ";"),
            Token::Dot => write!(f, "."),
            Token::LBrace => write!(f, "{{"),
            Token::RBrace => write!(f, "}}"),
            Token::LBracePipe => write!(f, "{{|"),
            Token::PipeRBrace => write!(f, "|}}"),
            Token::Colon => write!(f, ":"),
            Token::Pipe => write!(f, "|"),
            Token::Underscore => write!(f, "_"),
            Token::Of => write!(f, "of"),
            Token::Open => write!(f, "open"),
            Token::Module => write!(f, "module"),
            Token::Do => write!(f, "do"),
            Token::While => write!(f, "while"),
            Token::Break => write!(f, "break"),
            Token::Continue => write!(f, "continue"),
            Token::Async => write!(f, "async"),
            Token::Return => write!(f, "return"),
            Token::Yield => write!(f, "yield"),
            Token::LetBang => write!(f, "let!"),
            Token::DoBang => write!(f, "do!"),
            Token::ReturnBang => write!(f, "return!"),
            Token::YieldBang => write!(f, "yield!"),
            Token::LoadDirective(path) => write!(f, "#load \"{}\"", path),
            Token::Eof => write!(f, "EOF"),
        }
    }
}

/// A token with its position in the source code.
#[derive(Debug, Clone, PartialEq)]
pub struct TokenWithPos {
    /// The token
    pub token: Token,
    /// Position in source
    pub pos: Position,
}

impl TokenWithPos {
    /// Create a new token with position.
    pub fn new(token: Token, pos: Position) -> Self {
        TokenWithPos { token, pos }
    }
}

/// A token with its span in the source code.
#[derive(Debug, Clone, PartialEq)]
pub struct TokenWithSpan {
    /// The token
    pub token: Token,
    /// Span in source
    pub span: Span,
}

impl TokenWithSpan {
    /// Create a new token with span.
    pub fn new(token: Token, span: Span) -> Self {
        TokenWithSpan { token, span }
    }
}

/// Lexical analysis errors.
#[derive(Debug, Clone, PartialEq)]
pub enum LexError {
    /// Unexpected character encountered
    UnexpectedChar(char, Position),
    /// Unterminated string literal
    UnterminatedString(Position),
    /// Invalid number format
    InvalidNumber(String, Position),
    /// Unterminated multi-line comment
    UnterminatedComment(Position),
    /// Unknown directive
    UnknownDirective(String, Position),
}

impl fmt::Display for LexError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            LexError::UnexpectedChar(ch, pos) => {
                write!(f, "Unexpected character '{}' at {}", ch, pos)
            }
            LexError::UnterminatedString(pos) => {
                write!(f, "Unterminated string literal at {}", pos)
            }
            LexError::InvalidNumber(s, pos) => {
                write!(f, "Invalid number '{}' at {}", s, pos)
            }
            LexError::UnterminatedComment(pos) => {
                write!(f, "Unterminated multi-line comment at {}", pos)
            }
            LexError::UnknownDirective(name, pos) => {
                write!(f, "Unknown directive '{}' at {}", name, pos)
            }
        }
    }
}

impl std::error::Error for LexError {}

/// Lexer for Mini-F# source code.
///
/// The lexer converts source text into a stream of tokens with position information.
pub struct Lexer {
    /// Input characters
    input: Vec<char>,
    /// Current position in input
    pos: usize,
    /// Current line number (1-indexed)
    line: usize,
    /// Current column number (1-indexed)
    column: usize,
}

impl Lexer {
    /// Create a new lexer for the given input string.
    pub fn new(input: &str) -> Self {
        Lexer {
            input: input.chars().collect(),
            pos: 0,
            line: 1,
            column: 1,
        }
    }

    /// Tokenize the entire input, returning a vector of tokens with positions.
    pub fn tokenize(&mut self) -> Result<Vec<TokenWithPos>, LexError> {
        let mut tokens = Vec::new();

        while !self.is_at_end() {
            self.skip_whitespace_and_comments()?;
            if self.is_at_end() {
                break;
            }

            let start_pos = self.current_position();
            let token = self.next_token()?;
            tokens.push(TokenWithPos::new(token, start_pos));
        }

        tokens.push(TokenWithPos::new(Token::Eof, self.current_position()));

        Ok(tokens)
    }

    /// Tokenize with span information.
    pub fn tokenize_with_spans(&mut self) -> Result<Vec<TokenWithSpan>, LexError> {
        let mut tokens = Vec::new();

        while !self.is_at_end() {
            self.skip_whitespace_and_comments()?;
            if self.is_at_end() {
                break;
            }

            let start_pos = self.current_position();
            let token = self.next_token()?;
            let end_pos = self.current_position();
            let span = Span::new(start_pos, end_pos);
            tokens.push(TokenWithSpan::new(token, span));
        }

        let eof_pos = self.current_position();
        tokens.push(TokenWithSpan::new(Token::Eof, Span::point(eof_pos)));

        Ok(tokens)
    }

    /// Lex the next token from the input.
    fn next_token(&mut self) -> Result<Token, LexError> {
        let ch = self.current_char();

        match ch {
            '0'..='9' => self.lex_number(),
            'a'..='z' | 'A'..='Z' => self.lex_identifier_or_keyword(),
            '_' => {
                // Peek ahead to see if this is an identifier like _foo or just _
                if !self.is_at_end_or(1)
                    && (self.peek_char().is_alphanumeric() || self.peek_char() == '_')
                {
                    self.lex_identifier_or_keyword()
                } else {
                    self.advance();
                    Ok(Token::Underscore)
                }
            }
            '"' => self.lex_string(),
            '+' => self.lex_plus_or_plusplus(),
            '-' => self.lex_minus_or_arrow(),
            '*' => {
                self.advance();
                Ok(Token::Star)
            }
            '/' => {
                self.advance();
                Ok(Token::Slash)
            }
            '=' => self.lex_eq_or_eqeq(),
            '<' => self.lex_lt_or_lte_or_neq_or_larrow(),
            '>' => self.lex_gt_or_gte(),
            '&' => self.lex_and(),
            '|' => {
                self.advance();
                if !self.is_at_end() && self.current_char() == '|' {
                    self.advance();
                    Ok(Token::Or)
                } else if !self.is_at_end() && self.current_char() == '>' {
                    self.advance();
                    Ok(Token::PipeRight)
                } else if !self.is_at_end() && self.current_char() == ']' {
                    self.advance();
                    Ok(Token::PipeRBracket)
                } else if !self.is_at_end() && self.current_char() == '}' {
                    self.advance();
                    Ok(Token::PipeRBrace)
                } else {
                    Ok(Token::Pipe)
                }
            }
            ':' => self.lex_colon_or_coloncolon(),
            '(' => {
                self.advance();
                Ok(Token::LParen)
            }
            ')' => {
                self.advance();
                Ok(Token::RParen)
            }
            '[' => self.lex_lbracket_or_lbracket_pipe(),
            ']' => {
                self.advance();
                Ok(Token::RBracket)
            }
            ',' => {
                self.advance();
                Ok(Token::Comma)
            }
            ';' => {
                self.advance();
                Ok(Token::Semicolon)
            }
            '{' => self.lex_lbrace_or_lbrace_pipe(),
            '}' => {
                self.advance();
                Ok(Token::RBrace)
            }
            '.' => {
                self.advance();
                Ok(Token::Dot)
            }
            '#' => self.lex_directive(),
            _ => Err(LexError::UnexpectedChar(ch, self.current_position())),
        }
    }

    /// Lex a number (integer or float).
    fn lex_number(&mut self) -> Result<Token, LexError> {
        let start = self.pos;
        let start_pos = self.current_position();

        while !self.is_at_end() && self.current_char().is_ascii_digit() {
            self.advance();
        }

        // Check for float
        if !self.is_at_end()
            && self.current_char() == '.'
            && !self.is_at_end_or(1)
            && self.peek_char().is_ascii_digit()
        {
            self.advance(); // consume '.'
            while !self.is_at_end() && self.current_char().is_ascii_digit() {
                self.advance();
            }
            let s: String = self.input[start..self.pos].iter().collect();
            s.parse::<f64>()
                .map(Token::Float)
                .map_err(|_| LexError::InvalidNumber(s, start_pos))
        } else {
            let s: String = self.input[start..self.pos].iter().collect();
            s.parse::<i64>()
                .map(Token::Int)
                .map_err(|_| LexError::InvalidNumber(s, start_pos))
        }
    }

    /// Lex an identifier or keyword.
    fn lex_identifier_or_keyword(&mut self) -> Result<Token, LexError> {
        let start = self.pos;

        while !self.is_at_end()
            && (self.current_char().is_alphanumeric() || self.current_char() == '_')
        {
            self.advance();
        }

        let s: String = self.input[start..self.pos].iter().collect();

        // Check for bang-suffix tokens (let!, do!, return!, yield!)
        if !self.is_at_end() && self.current_char() == '!' {
            self.advance(); // consume the '!'
            let token = match s.as_str() {
                "let" => Token::LetBang,
                "do" => Token::DoBang,
                "return" => Token::ReturnBang,
                "yield" => Token::YieldBang,
                _ => {
                    // Not a recognized bang keyword, treat as identifier with bang
                    Token::Ident(format!("{}!", s))
                }
            };
            return Ok(token);
        }

        let token = match s.as_str() {
            "let" => Token::Let,
            "in" => Token::In,
            "rec" => Token::Rec,
            "match" => Token::Match,
            "and" => Token::AndKeyword,
            "if" => Token::If,
            "then" => Token::Then,
            "else" => Token::Else,
            "fun" => Token::Fun,
            "type" => Token::Type,
            "with" => Token::With,
            "of" => Token::Of,
            "open" => Token::Open,
            "module" => Token::Module,
            "do" => Token::Do,
            "while" => Token::While,
            "break" => Token::Break,
            "continue" => Token::Continue,
            "async" => Token::Async,
            "return" => Token::Return,
            "yield" => Token::Yield,
            "true" => Token::Bool(true),
            "false" => Token::Bool(false),
            _ => Token::Ident(s),
        };

        Ok(token)
    }

    /// Lex a string literal.
    fn lex_string(&mut self) -> Result<Token, LexError> {
        let start_pos = self.current_position();
        self.advance(); // consume opening "

        let mut s = String::new();

        while !self.is_at_end() && self.current_char() != '"' {
            let ch = self.current_char();
            if ch == '\\' {
                self.advance();
                if self.is_at_end() {
                    return Err(LexError::UnterminatedString(start_pos));
                }
                let escaped = match self.current_char() {
                    'n' => '\n',
                    't' => '\t',
                    'r' => '\r',
                    '\\' => '\\',
                    '"' => '"',
                    c => c, // Unknown escape, just include the character
                };
                s.push(escaped);
                self.advance();
            } else {
                s.push(ch);
                self.advance();
            }
        }

        if self.is_at_end() {
            return Err(LexError::UnterminatedString(start_pos));
        }

        self.advance(); // consume closing "
        Ok(Token::String(s))
    }

    /// Lex + or ++.
    fn lex_plus_or_plusplus(&mut self) -> Result<Token, LexError> {
        self.advance();
        if !self.is_at_end() && self.current_char() == '+' {
            self.advance();
            Ok(Token::PlusPlus)
        } else {
            Ok(Token::Plus)
        }
    }

    /// Lex - or ->.
    fn lex_minus_or_arrow(&mut self) -> Result<Token, LexError> {
        self.advance();
        if !self.is_at_end() && self.current_char() == '>' {
            self.advance();
            Ok(Token::Arrow)
        } else {
            Ok(Token::Minus)
        }
    }

    /// Lex = or ==.
    fn lex_eq_or_eqeq(&mut self) -> Result<Token, LexError> {
        self.advance();
        if !self.is_at_end() && self.current_char() == '=' {
            self.advance();
            Ok(Token::EqEq)
        } else {
            Ok(Token::Eq)
        }
    }

    /// Lex <, <=, <>, or <-.
    fn lex_lt_or_lte_or_neq_or_larrow(&mut self) -> Result<Token, LexError> {
        self.advance();
        if !self.is_at_end() {
            match self.current_char() {
                '=' => {
                    self.advance();
                    Ok(Token::Lte)
                }
                '>' => {
                    self.advance();
                    Ok(Token::Neq)
                }
                '-' => {
                    self.advance();
                    Ok(Token::LArrow)
                }
                _ => Ok(Token::Lt),
            }
        } else {
            Ok(Token::Lt)
        }
    }

    /// Lex > or >=.
    fn lex_gt_or_gte(&mut self) -> Result<Token, LexError> {
        self.advance();
        if !self.is_at_end() && self.current_char() == '=' {
            self.advance();
            Ok(Token::Gte)
        } else {
            Ok(Token::Gt)
        }
    }

    /// Lex &&.
    fn lex_and(&mut self) -> Result<Token, LexError> {
        let pos = self.current_position();
        self.advance();
        if !self.is_at_end() && self.current_char() == '&' {
            self.advance();
            Ok(Token::And)
        } else {
            Err(LexError::UnexpectedChar('&', pos))
        }
    }

    /// Lex : or ::.
    fn lex_colon_or_coloncolon(&mut self) -> Result<Token, LexError> {
        let _pos = self.current_position();
        self.advance();
        if !self.is_at_end() && self.current_char() == ':' {
            self.advance();
            Ok(Token::ColonColon)
        } else {
            Ok(Token::Colon)
        }
    }

    /// Lex [ or [|.
    fn lex_lbracket_or_lbracket_pipe(&mut self) -> Result<Token, LexError> {
        self.advance();
        if !self.is_at_end() && self.current_char() == '|' {
            self.advance();
            Ok(Token::LBracketPipe)
        } else {
            Ok(Token::LBracket)
        }
    }

    /// Lex { or {|.
    fn lex_lbrace_or_lbrace_pipe(&mut self) -> Result<Token, LexError> {
        self.advance();
        if !self.is_at_end() && self.current_char() == '|' {
            self.advance();
            Ok(Token::LBracePipe)
        } else {
            Ok(Token::LBrace)
        }
    }

    /// Lex a directive (#load "path").
    fn lex_directive(&mut self) -> Result<Token, LexError> {
        let start_pos = self.current_position();
        self.advance(); // consume '#'

        // Read the directive name
        let directive_start = self.pos;
        while !self.is_at_end()
            && (self.current_char().is_alphanumeric() || self.current_char() == '_')
        {
            self.advance();
        }

        let directive: String = self.input[directive_start..self.pos].iter().collect();

        // Skip whitespace between directive and path
        while !self.is_at_end() && matches!(self.current_char(), ' ' | '\t') {
            self.advance();
        }

        match directive.as_str() {
            "load" => {
                // Expect a string literal for the path
                if self.is_at_end() || self.current_char() != '"' {
                    return Err(LexError::UnknownDirective(
                        format!("#load requires a string path"),
                        start_pos,
                    ));
                }
                let path_token = self.lex_string()?;
                if let Token::String(path) = path_token {
                    Ok(Token::LoadDirective(path))
                } else {
                    Err(LexError::UnknownDirective(
                        format!("#load requires a string path"),
                        start_pos,
                    ))
                }
            }
            _ => Err(LexError::UnknownDirective(directive, start_pos)),
        }
    }

    /// Skip whitespace and comments.
    fn skip_whitespace_and_comments(&mut self) -> Result<(), LexError> {
        loop {
            if self.is_at_end() {
                break;
            }

            match self.current_char() {
                ' ' | '\t' | '\r' => {
                    self.advance();
                }
                '\n' => {
                    self.line += 1;
                    self.column = 0; // Will be incremented to 1 by advance()
                    self.advance();
                }
                '/' if !self.is_at_end_or(1) && self.peek_char() == '/' => {
                    // Single-line comment
                    self.skip_single_line_comment();
                }
                '(' if !self.is_at_end_or(1) && self.peek_char() == '*' => {
                    // Multi-line comment
                    self.skip_multiline_comment()?;
                }
                _ => break,
            }
        }
        Ok(())
    }

    /// Skip a single-line comment.
    fn skip_single_line_comment(&mut self) {
        while !self.is_at_end() && self.current_char() != '\n' {
            self.advance();
        }
    }

    /// Skip a multi-line comment with support for nested comments.
    /// F#-style comments: (* ... *) with nesting support.
    fn skip_multiline_comment(&mut self) -> Result<(), LexError> {
        let start_pos = self.current_position();
        let mut depth = 1;

        self.advance(); // skip '('
        self.advance(); // skip '*'

        while depth > 0 && !self.is_at_end() {
            let ch = self.current_char();

            if ch == '(' && !self.is_at_end_or(1) && self.peek_char() == '*' {
                // Nested comment start
                depth += 1;
                self.advance();
                self.advance();
            } else if ch == '*' && !self.is_at_end_or(1) && self.peek_char() == ')' {
                // Comment end
                depth -= 1;
                self.advance();
                self.advance();
            } else if ch == '\n' {
                // Track line numbers in comments
                self.line += 1;
                self.column = 0; // Will be incremented to 1 by advance()
                self.advance();
            } else {
                self.advance();
            }
        }

        if depth > 0 {
            return Err(LexError::UnterminatedComment(start_pos));
        }

        Ok(())
    }

    /// Get the current character without consuming it.
    fn current_char(&self) -> char {
        if self.is_at_end() {
            '\0'
        } else {
            self.input[self.pos]
        }
    }

    /// Peek at the next character without consuming it.
    fn peek_char(&self) -> char {
        if self.pos + 1 >= self.input.len() {
            '\0'
        } else {
            self.input[self.pos + 1]
        }
    }

    /// Advance to the next character.
    fn advance(&mut self) {
        if !self.is_at_end() {
            self.pos += 1;
            self.column += 1;
        }
    }

    /// Check if we're at the end of input.
    fn is_at_end(&self) -> bool {
        self.pos >= self.input.len()
    }

    /// Check if we're at the end or at position + offset is beyond end.
    fn is_at_end_or(&self, offset: usize) -> bool {
        self.pos + offset >= self.input.len()
    }

    /// Get the current position.
    fn current_position(&self) -> Position {
        Position::new(self.line, self.column, self.pos)
    }
}

// Tests are preserved from original file - they test TokenWithPos
#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_lex_integer() {
        let mut lexer = Lexer::new("42");
        let tokens = lexer.tokenize().unwrap();
        assert_eq!(tokens.len(), 2);
        assert_eq!(tokens[0].token, Token::Int(42));
        assert_eq!(tokens[1].token, Token::Eof);
    }

    #[test]
    fn test_tokenize_with_spans() {
        let mut lexer = Lexer::new("let x = 42");
        let tokens = lexer.tokenize_with_spans().unwrap();
        assert_eq!(tokens.len(), 5); // let, x, =, 42, EOF
        assert_eq!(tokens[0].token, Token::Let);
        assert!(tokens[0].span.is_single_line());
    }

    #[test]
    fn test_lex_anonymous_record_tokens() {
        let mut lexer = Lexer::new("{| x = 1 |}");
        let tokens = lexer.tokenize().unwrap();
        assert_eq!(tokens[0].token, Token::LBracePipe);
        assert_eq!(tokens[4].token, Token::PipeRBrace);
    }

    #[test]
    fn test_lex_pipe_disambiguation() {
        // Test that |} is recognized as PipeRBrace, not Pipe + RBrace
        let mut lexer = Lexer::new("|}");
        let tokens = lexer.tokenize().unwrap();
        assert_eq!(tokens.len(), 2); // |} and EOF
        assert_eq!(tokens[0].token, Token::PipeRBrace);
    }

    // Multi-line comment tests
    #[test]
    fn test_simple_multiline_comment() {
        let mut lexer = Lexer::new("(* comment *) let x = 42");
        let tokens = lexer.tokenize().unwrap();
        assert_eq!(tokens.len(), 5); // let, x, =, 42, EOF
        assert_eq!(tokens[0].token, Token::Let);
    }

    #[test]
    fn test_multiline_comment_multiline() {
        let source = r#"(* This is a
        multi-line comment
        spanning three lines *)
        let x = 10"#;
        let mut lexer = Lexer::new(source);
        let tokens = lexer.tokenize().unwrap();
        assert_eq!(tokens[0].token, Token::Let);
        assert_eq!(tokens[1].token, Token::Ident("x".to_string()));
    }

    #[test]
    fn test_nested_multiline_comments() {
        let source = "(* outer (* inner *) still outer *) let x = 5";
        let mut lexer = Lexer::new(source);
        let tokens = lexer.tokenize().unwrap();
        assert_eq!(tokens[0].token, Token::Let);
    }

    #[test]
    fn test_deeply_nested_comments() {
        let source = "(* level 1 (* level 2 (* level 3 *) level 2 *) level 1 *) let x = 1";
        let mut lexer = Lexer::new(source);
        let tokens = lexer.tokenize().unwrap();
        assert_eq!(tokens[0].token, Token::Let);
    }

    #[test]
    fn test_inline_multiline_comment() {
        let source = "let x = (* inline comment *) 42";
        let mut lexer = Lexer::new(source);
        let tokens = lexer.tokenize().unwrap();
        assert_eq!(tokens[0].token, Token::Let);
        assert_eq!(tokens[1].token, Token::Ident("x".to_string()));
        assert_eq!(tokens[2].token, Token::Eq);
        assert_eq!(tokens[3].token, Token::Int(42));
    }

    #[test]
    fn test_multiline_comment_with_special_chars() {
        let source = r#"(* Comment with special chars: !@#$%^&*()[]{}:;"'<>,.?/\|`~ *) let x = 1"#;
        let mut lexer = Lexer::new(source);
        let tokens = lexer.tokenize().unwrap();
        assert_eq!(tokens[0].token, Token::Let);
    }

    #[test]
    fn test_multiple_multiline_comments() {
        let source = "(* first *) let (* second *) x (* third *) = 42";
        let mut lexer = Lexer::new(source);
        let tokens = lexer.tokenize().unwrap();
        assert_eq!(tokens[0].token, Token::Let);
        assert_eq!(tokens[1].token, Token::Ident("x".to_string()));
        assert_eq!(tokens[2].token, Token::Eq);
        assert_eq!(tokens[3].token, Token::Int(42));
    }

    #[test]
    fn test_unterminated_multiline_comment() {
        let source = "(* This comment is not terminated let x = 42";
        let mut lexer = Lexer::new(source);
        let result = lexer.tokenize();
        assert!(result.is_err());
        match result.unwrap_err() {
            LexError::UnterminatedComment(_) => (),
            _ => panic!("Expected UnterminatedComment error"),
        }
    }

    #[test]
    fn test_unterminated_nested_comment() {
        let source = "(* outer (* inner *) still outer";
        let mut lexer = Lexer::new(source);
        let result = lexer.tokenize();
        assert!(result.is_err());
        match result.unwrap_err() {
            LexError::UnterminatedComment(_) => (),
            _ => panic!("Expected UnterminatedComment error"),
        }
    }

    #[test]
    fn test_mixed_single_and_multiline_comments() {
        let source = r#"
        // Single-line comment
        let x = 42 (* inline multi-line *)
        // Another single-line
        (* Multi-line
           spanning lines *)
        "#;
        let mut lexer = Lexer::new(source);
        let tokens = lexer.tokenize().unwrap();
        assert_eq!(tokens[0].token, Token::Let);
        assert_eq!(tokens[1].token, Token::Ident("x".to_string()));
        assert_eq!(tokens[2].token, Token::Eq);
        assert_eq!(tokens[3].token, Token::Int(42));
    }

    #[test]
    fn test_comment_does_not_affect_string() {
        // Comments should not be recognized inside strings
        let source = r#"let s = "(* not a comment *)""#;
        let mut lexer = Lexer::new(source);
        let tokens = lexer.tokenize().unwrap();
        assert_eq!(tokens[0].token, Token::Let);
        assert_eq!(tokens[1].token, Token::Ident("s".to_string()));
        assert_eq!(tokens[2].token, Token::Eq);
        assert_eq!(
            tokens[3].token,
            Token::String("(* not a comment *)".to_string())
        );
    }

    #[test]
    fn test_empty_multiline_comment() {
        let source = "(**) let x = 1";
        let mut lexer = Lexer::new(source);
        let tokens = lexer.tokenize().unwrap();
        assert_eq!(tokens[0].token, Token::Let);
    }

    #[test]
    fn test_multiline_comment_with_asterisks() {
        let source = "(* ** *** **** *) let x = 1";
        let mut lexer = Lexer::new(source);
        let tokens = lexer.tokenize().unwrap();
        assert_eq!(tokens[0].token, Token::Let);
    }

    #[test]
    fn test_lex_plusplus() {
        let mut lexer = Lexer::new("\"hello\" ++ \"world\"");
        let tokens = lexer.tokenize().unwrap();
        assert_eq!(tokens.len(), 4); // "hello", ++, "world", EOF
        assert_eq!(tokens[0].token, Token::String("hello".to_string()));
        assert_eq!(tokens[1].token, Token::PlusPlus);
        assert_eq!(tokens[2].token, Token::String("world".to_string()));
        assert_eq!(tokens[3].token, Token::Eof);
    }

    #[test]
    fn test_lex_plusplus_vs_plus() {
        let mut lexer = Lexer::new("1 + 2 ++ 3");
        let tokens = lexer.tokenize().unwrap();
        assert_eq!(tokens[0].token, Token::Int(1));
        assert_eq!(tokens[1].token, Token::Plus);
        assert_eq!(tokens[2].token, Token::Int(2));
        assert_eq!(tokens[3].token, Token::PlusPlus);
        assert_eq!(tokens[4].token, Token::Int(3));
    }

    #[test]
    fn test_lex_plusplus_no_space() {
        let mut lexer = Lexer::new("a++b");
        let tokens = lexer.tokenize().unwrap();
        assert_eq!(tokens[0].token, Token::Ident("a".to_string()));
        assert_eq!(tokens[1].token, Token::PlusPlus);
        assert_eq!(tokens[2].token, Token::Ident("b".to_string()));
    }

    #[test]
    fn test_lex_async_keyword() {
        let mut lexer = Lexer::new("async");
        let tokens = lexer.tokenize().unwrap();
        assert_eq!(tokens.len(), 2);
        assert_eq!(tokens[0].token, Token::Async);
    }

    #[test]
    fn test_lex_return_keyword() {
        let mut lexer = Lexer::new("return");
        let tokens = lexer.tokenize().unwrap();
        assert_eq!(tokens.len(), 2);
        assert_eq!(tokens[0].token, Token::Return);
    }

    #[test]
    fn test_lex_yield_keyword() {
        let mut lexer = Lexer::new("yield");
        let tokens = lexer.tokenize().unwrap();
        assert_eq!(tokens.len(), 2);
        assert_eq!(tokens[0].token, Token::Yield);
    }

    #[test]
    fn test_lex_let_bang() {
        let mut lexer = Lexer::new("let!");
        let tokens = lexer.tokenize().unwrap();
        assert_eq!(tokens.len(), 2);
        assert_eq!(tokens[0].token, Token::LetBang);
    }

    #[test]
    fn test_lex_do_bang() {
        let mut lexer = Lexer::new("do!");
        let tokens = lexer.tokenize().unwrap();
        assert_eq!(tokens.len(), 2);
        assert_eq!(tokens[0].token, Token::DoBang);
    }

    #[test]
    fn test_lex_return_bang() {
        let mut lexer = Lexer::new("return!");
        let tokens = lexer.tokenize().unwrap();
        assert_eq!(tokens.len(), 2);
        assert_eq!(tokens[0].token, Token::ReturnBang);
    }

    #[test]
    fn test_lex_yield_bang() {
        let mut lexer = Lexer::new("yield!");
        let tokens = lexer.tokenize().unwrap();
        assert_eq!(tokens.len(), 2);
        assert_eq!(tokens[0].token, Token::YieldBang);
    }

    #[test]
    fn test_lex_let_vs_let_bang() {
        let mut lexer = Lexer::new("let x = let! y = 42");
        let tokens = lexer.tokenize().unwrap();
        assert_eq!(tokens[0].token, Token::Let);
        assert_eq!(tokens[1].token, Token::Ident("x".to_string()));
        assert_eq!(tokens[2].token, Token::Eq);
        assert_eq!(tokens[3].token, Token::LetBang);
        assert_eq!(tokens[4].token, Token::Ident("y".to_string()));
        assert_eq!(tokens[5].token, Token::Eq);
        assert_eq!(tokens[6].token, Token::Int(42));
    }

    #[test]
    fn test_lex_do_vs_do_bang() {
        let mut lexer = Lexer::new("do x do! y");
        let tokens = lexer.tokenize().unwrap();
        assert_eq!(tokens[0].token, Token::Do);
        assert_eq!(tokens[1].token, Token::Ident("x".to_string()));
        assert_eq!(tokens[2].token, Token::DoBang);
        assert_eq!(tokens[3].token, Token::Ident("y".to_string()));
    }

    #[test]
    fn test_lex_computation_expression() {
        let source = "async { let! x = getAsync() do! printAsync x return! computeAsync x }";
        let mut lexer = Lexer::new(source);
        let tokens = lexer.tokenize().unwrap();
        let async_pos = tokens.iter().position(|t| t.token == Token::Async).unwrap();
        let letbang_pos = tokens
            .iter()
            .position(|t| t.token == Token::LetBang)
            .unwrap();
        let dobang_pos = tokens
            .iter()
            .position(|t| t.token == Token::DoBang)
            .unwrap();
        let returnbang_pos = tokens
            .iter()
            .position(|t| t.token == Token::ReturnBang)
            .unwrap();
        assert!(async_pos < letbang_pos);
        assert!(letbang_pos < dobang_pos);
        assert!(dobang_pos < returnbang_pos);
    }

    #[test]
    fn test_lex_bang_on_non_keyword() {
        let mut lexer = Lexer::new("foo!");
        let tokens = lexer.tokenize().unwrap();
        assert_eq!(tokens.len(), 2);
        assert_eq!(tokens[0].token, Token::Ident("foo!".to_string()));
    }

    #[test]
    fn test_lex_yield_in_expression() {
        let mut lexer = Lexer::new("yield x + yield! y");
        let tokens = lexer.tokenize().unwrap();
        assert_eq!(tokens[0].token, Token::Yield);
        assert_eq!(tokens[1].token, Token::Ident("x".to_string()));
        assert_eq!(tokens[2].token, Token::Plus);
        assert_eq!(tokens[3].token, Token::YieldBang);
        assert_eq!(tokens[4].token, Token::Ident("y".to_string()));
    }
}
