/// Iris Lexer - Tokenizer
/// 
/// Converts raw source code into a stream of tokens.

#[derive(Debug, Clone, PartialEq)]
pub enum TokenType {
    // Literals
    Null,
    Bool(bool),
    Number(f64),
    String(String),
    Identifier(String),

    // Keywords
    Let,
    Const,
    Fn,
    If,
    Else,
    While,
    For,
    In,
    Return,
    Break,
    Continue,
    True,
    False,
    And,
    Or,
    Not,
    Import,

    // Operators
    Plus,      // +
    Minus,     // -
    Star,      // *
    Slash,     // /
    Percent,   // %
    Eq,        // ==
    Neq,       // !=
    Lt,        // <
    Gt,        // >
    Lte,       // <=
    Gte,       // >=
    Assign,    // =
    PlusAssign,  // +=
    MinusAssign, // -=
    StarAssign,  // *=
    SlashAssign, // /=

    // Delimiters
    LParen,    // (
    RParen,    // )
    LBrace,    // {
    RBrace,    // }
    LBracket,  // [
    RBracket,  // ]
    Semicolon, // ;
    Colon,     // :
    Comma,     // ,
    Dot,       // .
    Arrow,     // =>

    // Special
    Comment(String),
    Newline,
    Eof,
}

#[derive(Debug, Clone)]
pub struct Token {
    pub token_type: TokenType,
    pub line: usize,
    pub column: usize,
}

pub fn tokenize(source: &str) -> Result<Vec<Token>, String> {
    let mut tokens = Vec::new();
    let mut chars = source.chars().peekable();
    let mut line = 1;
    let mut column = 1;

    while let Some(&ch) = chars.peek() {
        let start_col = column;

        match ch {
            // Whitespace
            ' ' | '\t' | '\r' => {
                chars.next();
                column += 1;
            }
            '\n' => {
                chars.next();
                tokens.push(Token { token_type: TokenType::Newline, line, column });
                line += 1;
                column = 1;
            }

            // Comments
            '/' if chars.clone().nth(1) == Some('/') => {
                chars.next(); chars.next();
                column += 2;
                let mut comment = String::new();
                while let Some(&c) = chars.peek() {
                    if c == '\n' { break; }
                    comment.push(c);
                    chars.next();
                    column += 1;
                }
                tokens.push(Token { token_type: TokenType::Comment(comment), line, column: start_col });
            }
            '/' if chars.clone().nth(1) == Some('*') => {
                chars.next(); chars.next();
                column += 2;
                let mut comment = String::new();
                let mut depth = 1;
                while depth > 0 {
                    match chars.next() {
                        Some('*') if chars.peek() == Some(&'/') => {
                            chars.next();
                            depth -= 1;
                        }
                        Some('/') if chars.peek() == Some(&'*') => {
                            chars.next();
                            depth += 1;
                        }
                        Some('\n') => { line += 1; column = 1; }
                        Some(c) => comment.push(c),
                        None => return Err(format!("Unterminated block comment at line {}", line)),
                    }
                }
                tokens.push(Token { token_type: TokenType::Comment(comment), line, column: start_col });
            }

            // Strings
            '"' => {
                chars.next();
                column += 1;
                let mut string = String::new();
                while let Some(&c) = chars.peek() {
                    if c == '"' { break; }
                    if c == '\\' {
                        chars.next();
                        column += 1;
                        match chars.peek() {
                            Some(&'n') => string.push('\n'),
                            Some(&'t') => string.push('\t'),
                            Some(&'r') => string.push('\r'),
                            Some(&'\\') => string.push('\\'),
                            Some(&'"') => string.push('"'),
                            Some(&c) => string.push(c),
                            None => return Err(format!("Unterminated string escape at line {}", line)),
                        }
                    } else {
                        string.push(c);
                    }
                    chars.next();
                    column += 1;
                }
                if chars.peek() != Some(&'"') {
                    return Err(format!("Unterminated string at line {}", line));
                }
                chars.next();
                column += 1;
                tokens.push(Token { token_type: TokenType::String(string), line, column: start_col });
            }

            // Numbers
            '0'..='9' => {
                let mut num_str = String::new();
                while let Some(&c) = chars.peek() {
                    if c.is_ascii_digit() || c == '.' {
                        num_str.push(c);
                        chars.next();
                        column += 1;
                    } else {
                        break;
                    }
                }
                match num_str.parse::<f64>() {
                    Ok(n) => tokens.push(Token { token_type: TokenType::Number(n), line, column: start_col }),
                    Err(_) => return Err(format!("Invalid number '{}' at line {}", num_str, line)),
                }
            }

            // Identifiers and Keywords
            'a'..='z' | 'A'..='Z' | '_' => {
                let mut ident = String::new();
                while let Some(&c) = chars.peek() {
                    if c.is_alphanumeric() || c == '_' {
                        ident.push(c);
                        chars.next();
                        column += 1;
                    } else {
                        break;
                    }
                }
                let token_type = match ident.as_str() {
                    "null" => TokenType::Null,
                    "true" => TokenType::Bool(true),
                    "false" => TokenType::Bool(false),
                    "let" => TokenType::Let,
                    "const" => TokenType::Const,
                    "fn" => TokenType::Fn,
                    "if" => TokenType::If,
                    "else" => TokenType::Else,
                    "while" => TokenType::While,
                    "for" => TokenType::For,
                    "in" => TokenType::In,
                    "return" => TokenType::Return,
                    "break" => TokenType::Break,
                    "continue" => TokenType::Continue,
                    "and" => TokenType::And,
                    "or" => TokenType::Or,
                    "not" => TokenType::Not,
                    "import" => TokenType::Import,
                    _ => TokenType::Identifier(ident),
                };
                tokens.push(Token { token_type, line, column: start_col });
            }

            // Multi-char operators
            '+' if chars.clone().nth(1) == Some('=') => {
                chars.next(); chars.next();
                column += 2;
                tokens.push(Token { token_type: TokenType::PlusAssign, line, column: start_col });
            }
            '-' if chars.clone().nth(1) == Some('=') => {
                chars.next(); chars.next();
                column += 2;
                tokens.push(Token { token_type: TokenType::MinusAssign, line, column: start_col });
            }
            '*' if chars.clone().nth(1) == Some('=') => {
                chars.next(); chars.next();
                column += 2;
                tokens.push(Token { token_type: TokenType::StarAssign, line, column: start_col });
            }
            '/' if chars.clone().nth(1) == Some('=') => {
                chars.next(); chars.next();
                column += 2;
                tokens.push(Token { token_type: TokenType::SlashAssign, line, column: start_col });
            }
            '=' if chars.clone().nth(1) == Some('=') => {
                chars.next(); chars.next();
                column += 2;
                tokens.push(Token { token_type: TokenType::Eq, line, column: start_col });
            }
            '!' if chars.clone().nth(1) == Some('=') => {
                chars.next(); chars.next();
                column += 2;
                tokens.push(Token { token_type: TokenType::Neq, line, column: start_col });
            }
            '<' if chars.clone().nth(1) == Some('=') => {
                chars.next(); chars.next();
                column += 2;
                tokens.push(Token { token_type: TokenType::Lte, line, column: start_col });
            }
            '>' if chars.clone().nth(1) == Some('=') => {
                chars.next(); chars.next();
                column += 2;
                tokens.push(Token { token_type: TokenType::Gte, line, column: start_col });
            }
            '=' if chars.clone().nth(1) == Some('>') => {
                chars.next(); chars.next();
                column += 2;
                tokens.push(Token { token_type: TokenType::Arrow, line, column: start_col });
            }

            // Single-char operators
            '+' => { chars.next(); column += 1; tokens.push(Token { token_type: TokenType::Plus, line, column: start_col }); }
            '-' => { chars.next(); column += 1; tokens.push(Token { token_type: TokenType::Minus, line, column: start_col }); }
            '*' => { chars.next(); column += 1; tokens.push(Token { token_type: TokenType::Star, line, column: start_col }); }
            '/' => { chars.next(); column += 1; tokens.push(Token { token_type: TokenType::Slash, line, column: start_col }); }
            '%' => { chars.next(); column += 1; tokens.push(Token { token_type: TokenType::Percent, line, column: start_col }); }
            '=' => { chars.next(); column += 1; tokens.push(Token { token_type: TokenType::Assign, line, column: start_col }); }
            '<' => { chars.next(); column += 1; tokens.push(Token { token_type: TokenType::Lt, line, column: start_col }); }
            '>' => { chars.next(); column += 1; tokens.push(Token { token_type: TokenType::Gt, line, column: start_col }); }

            // Delimiters
            '(' => { chars.next(); column += 1; tokens.push(Token { token_type: TokenType::LParen, line, column: start_col }); }
            ')' => { chars.next(); column += 1; tokens.push(Token { token_type: TokenType::RParen, line, column: start_col }); }
            '{' => { chars.next(); column += 1; tokens.push(Token { token_type: TokenType::LBrace, line, column: start_col }); }
            '}' => { chars.next(); column += 1; tokens.push(Token { token_type: TokenType::RBrace, line, column: start_col }); }
            '[' => { chars.next(); column += 1; tokens.push(Token { token_type: TokenType::LBracket, line, column: start_col }); }
            ']' => { chars.next(); column += 1; tokens.push(Token { token_type: TokenType::RBracket, line, column: start_col }); }
            ';' => { chars.next(); column += 1; tokens.push(Token { token_type: TokenType::Semicolon, line, column: start_col }); }
            ':' => { chars.next(); column += 1; tokens.push(Token { token_type: TokenType::Colon, line, column: start_col }); }
            ',' => { chars.next(); column += 1; tokens.push(Token { token_type: TokenType::Comma, line, column: start_col }); }
            '.' => { chars.next(); column += 1; tokens.push(Token { token_type: TokenType::Dot, line, column: start_col }); }

            _ => return Err(format!("Unexpected character '{}' at line {}, column {}", ch, line, column)),
        }
    }

    tokens.push(Token { token_type: TokenType::Eof, line, column });
    Ok(tokens)
}

