pub mod cursor;
pub mod scan_error;
mod scan_error_ext;
pub mod token;

pub use scan_error::{ScanError, ScanErrorKind};
pub use token::{Token, TokenKind, Tokens};

use crate::{Location, Source};
use cursor::Cursor;
use scan_error_ext::ScanErrorExt;
use std::{collections::HashMap, sync::LazyLock};

type TokenMap = HashMap<&'static str, TokenKind>;

/// Initialize a token map using 'key => value' notation.
macro_rules! token_map {
    ($($key:expr => $value:expr),+ $(,)?) => {
        {
            let mut map: TokenMap = HashMap::new();
            $(map.insert($key, $value);)+
            map
        }
    }
}

/// Nexus lexical scanner.
pub struct Scanner<'a> {
    comment_nesting_level: u32, // Tracks nested (multi-line) comment level.
    source: Source<'a>,
    cursor: Cursor<'a>,
    tokens: Tokens,
}

impl<'a> Scanner<'a> {
    /// Create a new scanner from a source string.
    ///
    /// # Example
    ///
    /// ```
    /// use nexus_lang::{Scanner, Source};
    ///
    /// let source = Source::new("fn main() {}");
    /// let scanner = Scanner::new(source);
    /// ```
    ///
    #[must_use]
    pub fn new(source: Source<'a>) -> Self {
        Self {
            comment_nesting_level: 0,
            source,
            cursor: Cursor::new(source.code()),
            tokens: Vec::new(),
        }
    }

    /// Scan the source string and generate a list of tokens.
    ///
    /// # Errors
    ///
    /// This function will return a `ScanError` if the source string can not be parsed properly.
    ///
    /// Note: The current implementation will bail out at the first error it encounters.
    ///
    /// # Example
    ///
    /// ```
    /// use nexus_lang::{Scanner, Source};
    ///
    /// let source = Source::new("fn main() {}");
    /// let scanner = Scanner::new(source);
    /// let tokens = scanner.scan().expect("Scanning failed.");
    /// ```
    ///
    pub fn scan(mut self) -> Result<Tokens, ScanError> {
        while let Some(c) = self.cursor.value() {
            if self.comment_nesting_level > 0 {
                match c {
                    '/' if self.cursor.peek() == Some('*') => {
                        // Nested comment start.
                        self.comment_nesting_level += 1;
                        self.cursor.advance(); // Skip the '*'.
                    }
                    '*' if self.cursor.peek() == Some('/') => {
                        // Comment end.
                        self.comment_nesting_level -= 1;
                        self.cursor.advance(); // Skip the '/'.
                    }
                    _ => {}
                }

                self.cursor.advance();
                continue;
            }

            match c {
                '(' => self.push_token(TokenKind::LParen, "("),
                ')' => self.push_token(TokenKind::RParen, ")"),
                '{' => self.push_token(TokenKind::LBrace, "{"),
                '}' => self.push_token(TokenKind::RBrace, "}"),
                '[' => self.push_token(TokenKind::LBracket, "["),
                ']' => self.push_token(TokenKind::RBracket, "]"),
                ':' => self.push_token(TokenKind::Colon, ":"),
                ';' => self.push_token(TokenKind::Semicolon, ";"),
                '&' => self.peek_match(('&', TokenKind::And, "&&"), (TokenKind::Amp, "&")),
                '+' => self.peek_match(('=', TokenKind::PlusIs, "+="), (TokenKind::Plus, "+")),
                '-' => match self.cursor.peek() {
                    Some('>') => {
                        self.cursor.advance();
                        self.push_token(TokenKind::Arrow, "->");
                    }
                    Some('=') => {
                        self.cursor.advance();
                        self.push_token(TokenKind::MinusIs, "-=");
                    }
                    _ => self.push_token(TokenKind::Minus, "-"),
                },
                '*' => self.peek_match(('=', TokenKind::StarIs, "*="), (TokenKind::Star, "*")),
                '/' => {
                    match self.cursor.peek() {
                        Some('/') => {
                            self.cursor.advance_while(|c| c != '\n');
                            continue;
                        }
                        Some('*') => {
                            self.comment_nesting_level = 1;
                            self.cursor.advance();
                            // Skip the normal cursor advancement to avoid double-advancing.
                            continue;
                        }
                        Some('=') => {
                            self.cursor.advance();
                            self.push_token(TokenKind::SlashIs, "/=");
                        }
                        _ => self.push_token(TokenKind::Slash, "/"),
                    }
                }
                '\\' => self.push_token(TokenKind::Backslash, "\\"),
                '%' => {
                    self.peek_match(('=', TokenKind::PercentIs, "%="), (TokenKind::Percent, "%"));
                }
                ',' => self.push_token(TokenKind::Comma, ","),
                '.' => match self.cursor.peek() {
                    Some('.') => {
                        self.cursor.advance();
                        if self.cursor.peek() == Some('=') {
                            self.cursor.advance();
                            self.push_token(TokenKind::RangeInclusive, "..=");
                        } else {
                            self.push_token(TokenKind::Range, "..");
                        }
                    }
                    _ => self.push_token(TokenKind::Dot, "."),
                },
                '_' => self.push_token(TokenKind::Underscore, "_"),
                '=' => self.peek_match(('=', TokenKind::Eq, "=="), (TokenKind::Is, "=")),
                '>' => self.peek_match(('=', TokenKind::GtEq, ">="), (TokenKind::Gt, ">")),
                '<' => self.peek_match(('=', TokenKind::LtEq, "<="), (TokenKind::Lt, "<")),
                '!' => self.peek_match(('=', TokenKind::NotEq, "!="), (TokenKind::Bang, "!")),
                '|' => match self.cursor.peek() {
                    Some('|') => {
                        self.cursor.advance();
                        self.push_token(TokenKind::DoublePipe, "||");
                    }
                    _ => self.push_token(TokenKind::Pipe, "|"),
                },
                '"' => {
                    self.parse_string()?;
                    // Skip the normal cursor advancement since parse_string already consumes the closing quote.
                    continue;
                }
                c if c.is_ascii_digit() => self.parse_number()?,
                c if c.is_alphabetic() => self.parse_alphabetic(),
                c if c.is_whitespace() => {
                    // Whitespace is ignored, cursor will advance below
                }
                _ => return Err(self.error(ScanErrorKind::UnexpectedCharacter(c))),
            }

            self.cursor.advance();
        }

        Ok(self.tokens)
    }

    /// Parse string literal.
    fn parse_string(&mut self) -> Result<(), ScanError> {
        let mut string_value = String::new();
        let start_index = self.cursor.index();
        let start_loc = self.cursor.location();

        // Skip opening quote.
        self.cursor.consume();

        while let Some(c) = self.cursor.consume() {
            match c {
                '"' => {
                    // Found closing quote - capture the original lexeme including quotes.
                    let end_index = self.cursor.index();
                    let lexeme = &self.source.code()[start_index..end_index];
                    self.push_token_at(TokenKind::String(string_value), lexeme, start_loc);
                    return Ok(());
                }
                '\\' => {
                    // Handle escape sequences.
                    if let Some(escape_char) = self.cursor.consume() {
                        match escape_char {
                            '"' => string_value.push('"'),
                            '\\' => string_value.push('\\'),
                            'n' => string_value.push('\n'),
                            'r' => string_value.push('\r'),
                            't' => string_value.push('\t'),
                            'b' => string_value.push('\u{0008}'), // backspace
                            'f' => string_value.push('\u{000C}'), // form feed
                            '0' => string_value.push('\0'),
                            'x' => {
                                // Parse \x followed by 2 hex digits.
                                let hex_value = self.parse_hex_escape(2)?;
                                if let Some(ch) = char::from_u32(hex_value) {
                                    string_value.push(ch);
                                } else {
                                    return Err(
                                        self.error(ScanErrorKind::InvalidEscapeSequence('x'))
                                    );
                                }
                            }
                            'u' => {
                                // Parse \u followed by 4 hex digits.
                                let hex_value = self.parse_hex_escape(4)?;
                                if let Some(ch) = char::from_u32(hex_value) {
                                    string_value.push(ch);
                                } else {
                                    return Err(
                                        self.error(ScanErrorKind::InvalidEscapeSequence('u'))
                                    );
                                }
                            }
                            'U' => {
                                // Parse \U followed by 8 hex digits.
                                let hex_value = self.parse_hex_escape(8)?;
                                if let Some(ch) = char::from_u32(hex_value) {
                                    string_value.push(ch);
                                } else {
                                    return Err(
                                        self.error(ScanErrorKind::InvalidEscapeSequence('U'))
                                    );
                                }
                            }
                            _ => {
                                return Err(
                                    self.error(ScanErrorKind::InvalidEscapeSequence(escape_char))
                                );
                            }
                        }
                    } else {
                        return Err(self.error(ScanErrorKind::IncompleteEscapeSequence('\\')));
                    }
                }
                '\n' | '\r' => {
                    return Err(self.error(ScanErrorKind::UnterminatedStringLiteral(string_value)));
                }
                _ => {
                    // Regular character.
                    string_value.push(c);
                }
            }
        }

        // If we reach here, the string was not terminated.
        Err(self.error(ScanErrorKind::UnterminatedStringLiteral(string_value)))
    }

    /// Parse hex escape sequence with specified number of digits.
    fn parse_hex_escape(&mut self, num_digits: usize) -> Result<u32, ScanError> {
        let mut hex_value = 0u32;

        for _ in 0..num_digits {
            if let Some(hex_char) = self.cursor.consume() {
                if let Some(digit) = hex_char.to_digit(16) {
                    hex_value = hex_value * 16 + digit;
                } else {
                    return Err(self.error(ScanErrorKind::InvalidHexDigit(hex_char)));
                }
            } else {
                let escape_type = match num_digits {
                    2 => 'x',
                    4 => 'u',
                    8 => 'U',
                    _ => '?',
                };
                return Err(self.error(ScanErrorKind::IncompleteEscapeSequence(escape_type)));
            }
        }

        Ok(hex_value)
    }

    /// Parse number literal.
    fn parse_number(&mut self) -> Result<(), ScanError> {
        // Save the starting location before consuming the token.
        let start_loc = self.cursor.location();

        let first_char = self.cursor.value().unwrap(); // NOTE: The cursor was peeked before.
        let mut s = first_char.to_string(); // Stripped (no underscores); used for f64 parsing.
        let mut raw = first_char.to_string(); // Raw literal including underscores; used for errors.

        let mut has_dot = false;
        let mut prev_was_underscore = false;
        let mut fractional_digit_seen = false;

        while let Some(p) = self.cursor.peek() {
            if p.is_ascii_digit() {
                self.cursor.advance();
                let c = self.cursor.value().unwrap(); // NOTE: The cursor was peeked before.
                s.push(c);
                raw.push(c);
                prev_was_underscore = false;
                if has_dot {
                    fractional_digit_seen = true;
                }
            } else if p == '_' {
                self.cursor.advance();
                raw.push('_');
                if prev_was_underscore || (has_dot && !fractional_digit_seen) {
                    // Adjacent underscores (e.g. `1__0`) or underscore as first fractional char (e.g. `1._0`).
                    return Err(self.error(ScanErrorKind::InvalidNumberUnderscore(raw)));
                }
                prev_was_underscore = true;
            } else if p == '.' && !has_dot {
                if prev_was_underscore {
                    // Trailing underscore before decimal point (e.g. `1_.5`).
                    return Err(self.error(ScanErrorKind::InvalidNumberUnderscore(raw)));
                }
                has_dot = true;
                self.cursor.advance();
                let c = self.cursor.value().unwrap(); // NOTE: The cursor was peeked before.
                s.push(c);
                raw.push(c);
            } else {
                break;
            }
        }

        if prev_was_underscore {
            // Trailing underscore at end of literal (e.g. `1_`).
            return Err(self.error(ScanErrorKind::InvalidNumberUnderscore(raw)));
        }

        if s.ends_with('.') {
            match self.cursor.peek() {
                Some('.') => {
                    // Range token ('..') or inclusive range ('..='); remove trailing dot and advance cursor.
                    s.pop();
                    self.cursor.advance();

                    // NOTE: We know we're parsing a digit, so unwrapping is fine.
                    self.push_token_at(TokenKind::Number(s.parse().unwrap()), s, start_loc);

                    if self.cursor.peek() == Some('=') {
                        self.cursor.advance();
                        self.push_token(TokenKind::RangeInclusive, "..=");
                    } else {
                        self.push_token(TokenKind::Range, "..");
                    }
                }
                _ => {
                    return Err(self.error(ScanErrorKind::MalformedNumberLiteral(raw)));
                }
            }
        } else {
            // NOTE: We know we're parsing a digit, so unwrapping is fine.
            self.push_token_at(TokenKind::Number(s.parse().unwrap()), s, start_loc);
        }

        Ok(())
    }

    /// Parse keyword/identifier.
    fn parse_alphabetic(&mut self) {
        static KEYWORDS: LazyLock<TokenMap> = LazyLock::new(|| {
            token_map! {
                "true"     => TokenKind::True,
                "false"    => TokenKind::False,
                "assert"   => TokenKind::Assert,
                "const"    => TokenKind::Const,
                "let"      => TokenKind::Let,
                "mut"      => TokenKind::Mut,
                "fn"       => TokenKind::Function,
                "if"       => TokenKind::If,
                "else"     => TokenKind::Else,
                "for"      => TokenKind::For,
                "in"       => TokenKind::In,
                "loop"     => TokenKind::Loop,
                "while"    => TokenKind::While,
                "break"    => TokenKind::Break,
                "continue" => TokenKind::Continue,
                "return"   => TokenKind::Return,
                "use"      => TokenKind::Use,
                "bool"     => TokenKind::BoolId,
                "Node"     => TokenKind::NodeId,
                "Group"    => TokenKind::GroupId,
                "Number"   => TokenKind::NumberId,
                "PortRef"  => TokenKind::PortRefId,
                "String"   => TokenKind::StringId,
                "port"     => TokenKind::Port,
                "print"    => TokenKind::Print,
                "node"     => TokenKind::Node,
                "group"    => TokenKind::Group,
                "type"     => TokenKind::Type,
                "out"      => TokenKind::Out,
                "inout"    => TokenKind::Inout,
                "self"     => TokenKind::Self_,
            }
        });

        // Save the starting location before consuming the token.
        let start_loc = self.cursor.location();

        // NOTE: The cursor value was peeked before we arrive here.
        let mut s = self.cursor.value().unwrap().to_string();

        // Eagerly consume all consecutive alphanumeric characters to form an identifier.
        while let Some(p) = self.cursor.peek() {
            if p.is_alphanumeric() || p == '_' {
                self.cursor.advance();
                // NOTE: The cursor value was peeked before advancing.
                s.push(self.cursor.value().unwrap());
            } else {
                break;
            }
        }

        if let Some(t) = KEYWORDS.get(s.as_str()) {
            self.push_token_at(t.clone(), s, start_loc);
        } else {
            self.push_token_at(TokenKind::Identifier(s.clone()), s, start_loc);
        }
    }

    /// Append a token to the internal token stream.
    fn push_token(&mut self, kind: TokenKind, lexeme: impl Into<String>) {
        self.tokens.push(Token {
            kind,
            lexeme: lexeme.into(),
            loc: self.cursor.location(),
        });
    }

    /// Append a token to the internal token stream with a specific location.
    fn push_token_at(&mut self, kind: TokenKind, lexeme: impl Into<String>, loc: Location) {
        self.tokens.push(Token {
            kind,
            lexeme: lexeme.into(),
            loc,
        });
    }

    /// Helper for matching on peek with a default token fallback.
    ///
    /// Handles common pattern of checking next character and either:
    ///
    /// - Advancing cursor and returning a compound token if it matches,
    /// - Or just returning a simple token if it doesn't match.
    ///
    fn peek_match(&mut self, compound: (char, TokenKind, &str), default: (TokenKind, &str)) {
        let (peek_char, compound_kind, compound_lexeme) = compound;
        let (default_kind, default_lexeme) = default;

        match self.cursor.peek() {
            Some(c) if c == peek_char => {
                self.push_token(compound_kind, compound_lexeme);
                self.cursor.advance();
            }
            _ => self.push_token(default_kind, default_lexeme),
        }
    }
}

impl ScanErrorExt for Scanner<'_> {
    fn error(&self, kind: ScanErrorKind) -> ScanError {
        ScanError::new(self.source.into(), &self.cursor.location(), kind)
    }
}
