//! Lexical analysis and tokenization

use crate::frontend::ast::Span;
use logos::{Lexer, Logos};

/// Process escape sequences in a string literal
fn process_escapes(s: &str) -> String {
    let mut result = String::new();
    let mut chars = s.chars();

    while let Some(ch) = chars.next() {
        if ch == '\\' {
            match chars.next() {
                Some('n') => result.push('\n'),
                Some('t') => result.push('\t'),
                Some('r') => result.push('\r'),
                Some('\\') | None => result.push('\\'), // Backslash or end of string
                Some('"') => result.push('"'),
                Some('\'') => result.push('\''),
                Some('0') => result.push('\0'),
                Some('u') if chars.as_str().starts_with('{') => {
                    // Unicode escape: \u{XXXX}
                    chars.next(); // consume '{'
                    let mut hex = String::new();
                    for hex_char in chars.by_ref() {
                        if hex_char == '}' {
                            break;
                        }
                        hex.push(hex_char);
                    }

                    if let Ok(code_point) = u32::from_str_radix(&hex, 16) {
                        if let Some(unicode_char) = char::from_u32(code_point) {
                            result.push(unicode_char);
                        } else {
                            // Invalid Unicode code point, keep as literal
                            result.push_str("\\u{");
                            result.push_str(&hex);
                            result.push('}');
                        }
                    } else {
                        // Invalid hex, keep as literal
                        result.push_str("\\u{");
                        result.push_str(&hex);
                        result.push('}');
                    }
                }
                Some(other) => {
                    // Unknown escape sequence, keep as literal
                    result.push('\\');
                    result.push(other);
                }
            }
        } else {
            result.push(ch);
        }
    }

    result
}

#[derive(Logos, Debug, PartialEq, Clone)]
#[logos(skip r"[ \t\n\f]+")]
#[logos(skip r"//[^\n]*")]
#[logos(skip r"/\*([^*]|\*[^/])*\*/")]
pub enum Token {
    // Literals
    #[regex(r"[0-9]+(?:i8|i16|i32|i64|i128|isize|u8|u16|u32|u64|u128|usize)?", |lex| {
        let slice = lex.slice();
        // Strip type suffix for parsing the numeric value
        let (num_part, _type_suffix) = if let Some(pos) = slice.find(|c: char| c.is_alphabetic()) {
            (&slice[..pos], Some(&slice[pos..]))
        } else {
            (slice, None)
        };
        num_part.parse::<i64>().ok()
    })]
    Integer(i64),

    #[regex(r"[0-9]+\.[0-9]+([eE][+-]?[0-9]+)?", |lex| lex.slice().parse::<f64>().ok())]
    Float(f64),

    #[regex(r#""([^"\\]|\\.)*""#, |lex| {
        let s = lex.slice();
        let inner = &s[1..s.len()-1];
        Some(process_escapes(inner))
    })]
    String(String),

    #[regex(r#"f"([^"\\]|\\.)*""#, |lex| {
        let s = lex.slice();
        // Remove f" prefix and " suffix
        let inner = &s[2..s.len()-1];
        Some(process_escapes(inner))
    })]
    FString(String),

    #[regex(r#"r"([^"])*""#, |lex| {
        let s = lex.slice();
        // Remove r" prefix and " suffix - no escape processing for raw strings
        Some(s[2..s.len()-1].to_string())
    })]
    RawString(String),

    #[regex(r"'([^'\\]|\\.)'", |lex| {
        let s = lex.slice();
        let inner = &s[1..s.len()-1];
        if inner.len() == 1 {
            inner.chars().next()
        } else if inner.starts_with('\\') && inner.len() == 2 {
            match inner.chars().nth(1) {
                Some('n') => Some('\n'),
                Some('t') => Some('\t'),
                Some('r') => Some('\r'),
                Some('\\') => Some('\\'),
                Some('\'') => Some('\''),
                Some('0') => Some('\0'),
                _ => None,
            }
        } else {
            None
        }
    })]
    Char(char),

    #[token("true", |_| true)]
    #[token("false", |_| false)]
    Bool(bool),

    // Keywords
    #[token("fun")]
    Fun,
    #[token("fn")]
    Fn,
    #[token("let")]
    Let,
    #[token("var")]
    Var,
    #[token("mod")]
    Mod,
    #[token("if")]
    If,
    #[token("else")]
    Else,
    #[token("match")]
    Match,
    #[token("for")]
    For,
    #[token("in")]
    In,
    #[token("while")]
    While,
    #[token("loop")]
    Loop,
    #[token("async")]
    Async,
    #[token("await")]
    Await,
    #[token("throw")]
    Throw,
    #[token("try")]
    Try,
    #[token("catch")]
    Catch,
    #[token("return")]
    Return,
    #[token("command")]
    Command,
    #[token("Ok")]
    Ok,
    #[token("Err")]
    Err,
    #[token("Some")]
    Some,
    #[token("None")]
    None,
    #[token("null")]
    Null,
    #[token("Result")]
    Result,
    #[token("Option")]
    Option,
    #[token("break")]
    Break,
    #[token("continue")]
    Continue,
    #[token("struct")]
    Struct,
    #[token("enum")]
    Enum,
    #[token("impl")]
    Impl,
    #[token("trait")]
    Trait,
    #[token("extend")]
    Extend,
    #[token("actor")]
    Actor,
    #[token("state")]
    State,
    #[token("receive")]
    Receive,
    #[token("send")]
    Send,
    #[token("ask")]
    Ask,
    #[token("type")]
    Type,
    #[token("where")]
    Where,
    #[token("const")]
    Const,
    #[token("static")]
    Static,
    #[token("mut")]
    Mut,
    #[token("pub")]
    Pub,
    #[token("import")]
    Import,
    #[token("use")]
    Use,
    #[token("as")]
    As,
    #[token("module")]
    Module,
    #[token("export")]
    Export,
    #[token("df", priority = 2)]
    DataFrame,

    // Identifiers (lower priority than keywords)
    #[regex(r"[a-zA-Z_][a-zA-Z0-9_]*", |lex| lex.slice().to_string(), priority = 1)]
    Identifier(String),

    // Operators
    #[token("+")]
    Plus,
    #[token("-")]
    Minus,
    #[token("*")]
    Star,
    #[token("/")]
    Slash,
    #[token("%")]
    Percent,
    #[token("**")]
    Power,

    #[token("==")]
    EqualEqual,
    #[token("!=")]
    NotEqual,
    #[token("<?")]
    ActorQuery,
    #[token("<-")]
    LeftArrow,
    #[token("<")]
    Less,
    #[token("<=")]
    LessEqual,
    #[token(">")]
    Greater,
    #[token(">=")]
    GreaterEqual,

    #[token("&&")]
    AndAnd,
    #[token("||")]
    OrOr,
    #[token("!")]
    Bang,

    #[token("&")]
    Ampersand,
    #[token("|")]
    Pipe,
    #[token("^")]
    Caret,
    #[token("~")]
    Tilde,
    #[token("\\")]
    Backslash,
    #[token("<<")]
    LeftShift,

    #[token("=")]
    Equal,
    #[token("+=")]
    PlusEqual,
    #[token("-=")]
    MinusEqual,
    #[token("*=")]
    StarEqual,
    #[token("/=")]
    SlashEqual,
    #[token("%=")]
    PercentEqual,
    #[token("**=")]
    PowerEqual,
    #[token("&=")]
    AmpersandEqual,
    #[token("|=")]
    PipeEqual,
    #[token("^=")]
    CaretEqual,
    #[token("<<=")]
    LeftShiftEqual,

    #[token("++")]
    Increment,
    #[token("--")]
    Decrement,

    #[token("|>")]
    Pipeline,
    #[token("->")]
    Arrow,
    #[token("=>")]
    FatArrow,
    #[token("..")]
    DotDot,
    #[token("..=")]
    DotDotEqual,
    #[token("...")]
    DotDotDot,
    #[token("??")]
    NullCoalesce,
    #[token("?")]
    Question,
    #[token("?.")]
    SafeNav,

    // Delimiters
    #[token("(")]
    LeftParen,
    #[token(")")]
    RightParen,
    #[token("[")]
    LeftBracket,
    #[token("]")]
    RightBracket,
    #[token("{")]
    LeftBrace,
    #[token("}")]
    RightBrace,

    // Punctuation
    #[token(",")]
    Comma,
    #[token(".")]
    Dot,
    #[token(":")]
    Colon,
    #[token("::")]
    ColonColon,
    #[token(";")]
    Semicolon,
    #[token("_", priority = 2)]
    Underscore,

    // Attribute support
    #[token("#")]
    Hash,
}

impl Token {
    #[must_use]
    pub fn is_binary_op(&self) -> bool {
        matches!(
            self,
            Token::Plus
                | Token::Minus
                | Token::Star
                | Token::Slash
                | Token::Percent
                | Token::Power
                | Token::EqualEqual
                | Token::NotEqual
                | Token::Less
                | Token::LessEqual
                | Token::Greater
                | Token::GreaterEqual
                | Token::AndAnd
                | Token::OrOr
                | Token::Ampersand
                | Token::Pipe
                | Token::Caret
                | Token::LeftShift
        )
    }

    #[must_use]
    pub fn is_unary_op(&self) -> bool {
        matches!(
            self,
            Token::Bang | Token::Minus | Token::Tilde | Token::Ampersand
        )
    }

    #[must_use]
    pub fn is_assignment_op(&self) -> bool {
        matches!(
            self,
            Token::Equal
                | Token::PlusEqual
                | Token::MinusEqual
                | Token::StarEqual
                | Token::SlashEqual
                | Token::PercentEqual
                | Token::PowerEqual
                | Token::AmpersandEqual
                | Token::PipeEqual
                | Token::CaretEqual
                | Token::LeftShiftEqual
        )
    }
}

pub struct TokenStream<'a> {
    lexer: Lexer<'a, Token>,
    peeked: Option<(Token, Span)>,
}

/// Saved position in the token stream for backtracking
#[derive(Clone)]
pub struct TokenStreamPosition<'a> {
    lexer: Lexer<'a, Token>,
    peeked: Option<(Token, Span)>,
}

impl<'a> TokenStream<'a> {
    #[must_use]
    pub fn new(input: &'a str) -> Self {
        Self {
            lexer: Token::lexer(input),
            peeked: None,
        }
    }

    /// Save the current position for later restoration
    #[must_use]
    pub fn position(&self) -> TokenStreamPosition<'a> {
        TokenStreamPosition {
            lexer: self.lexer.clone(),
            peeked: self.peeked.clone(),
        }
    }

    /// Restore a previously saved position
    pub fn set_position(&mut self, pos: TokenStreamPosition<'a>) {
        self.lexer = pos.lexer;
        self.peeked = pos.peeked;
    }

    #[allow(clippy::should_implement_trait)]
    pub fn next(&mut self) -> Option<(Token, Span)> {
        if let Some(peeked) = self.peeked.take() {
            return Some(peeked);
        }

        self.lexer.next().map(|result| {
            let token = result.unwrap_or(Token::Bang); // Error recovery
            let span = Span::new(self.lexer.span().start, self.lexer.span().end);
            (token, span)
        })
    }

    pub fn peek(&mut self) -> Option<&(Token, Span)> {
        if self.peeked.is_none() {
            self.peeked = self.next();
        }
        self.peeked.as_ref()
    }

    pub fn peek_nth(&mut self, n: usize) -> Option<(Token, Span)> {
        // For simplicity, we'll only support peek_nth(1) by cloning the lexer
        if n == 1 {
            let saved_peeked = self.peeked.clone();
            let saved_lexer = self.lexer.clone();

            // Get first token
            let _first = self.peek();
            self.advance();

            // Get second token
            let result = self.peek().cloned();

            // Restore state
            self.lexer = saved_lexer;
            self.peeked = saved_peeked;

            result
        } else {
            None // Not supported for n > 1
        }
    }

    pub fn peek_nth_is_colon(&mut self, n: usize) -> bool {
        if n == 0 {
            self.peek().is_some_and(|(t, _)| matches!(t, Token::Colon))
        } else {
            self.peek_nth(n)
                .is_some_and(|(t, _)| matches!(t, Token::Colon))
        }
    }

    /// Expect a specific token and return its span
    ///
    /// # Errors
    ///
    /// Returns an error if the next token doesn't match the expected token or if we reached EOF
    pub fn expect(&mut self, expected: &Token) -> anyhow::Result<Span> {
        match self.next() {
            Some((token, span)) if token == *expected => Ok(span),
            Some((token, _)) => anyhow::bail!("Expected {:?}, found {:?}", expected, token),
            None => anyhow::bail!("Expected {:?}, found EOF", expected),
        }
    }

    // Alias for next() to avoid clippy warning about Iterator trait
    pub fn advance(&mut self) -> Option<(Token, Span)> {
        self.next()
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used)]
#[allow(clippy::panic)]
mod tests {
    use super::*;
    use proptest::prelude::*;

    #[test]
    #[allow(clippy::approx_constant)] // Intentional literal for test
    fn test_tokenize_basic() {
        let mut stream = TokenStream::new("let x = 42 + 3.14");

        assert_eq!(stream.next().map(|(t, _)| t), Some(Token::Let));
        assert_eq!(
            stream.next().map(|(t, _)| t),
            Some(Token::Identifier("x".to_string()))
        );
        assert_eq!(stream.next().map(|(t, _)| t), Some(Token::Equal));
        assert_eq!(stream.next().map(|(t, _)| t), Some(Token::Integer(42)));
        assert_eq!(stream.next().map(|(t, _)| t), Some(Token::Plus));
        assert_eq!(stream.next().map(|(t, _)| t), Some(Token::Float(3.14))); // Intentional literal for test
        assert_eq!(stream.next().map(|(t, _)| t), None);
    }

    #[test]
    fn test_tokenize_pipeline() {
        let mut stream = TokenStream::new("[1, 2, 3] >> map(|x| x * 2)");

        assert_eq!(stream.next().map(|(t, _)| t), Some(Token::LeftBracket));
        assert_eq!(stream.next().map(|(t, _)| t), Some(Token::Integer(1)));
        // ... rest of tokens
    }

    #[test]
    fn test_tokenize_comments() {
        let mut stream = TokenStream::new("x // comment\n+ /* block */ y");

        assert_eq!(
            stream.next().map(|(t, _)| t),
            Some(Token::Identifier("x".to_string()))
        );
        assert_eq!(stream.next().map(|(t, _)| t), Some(Token::Plus));
        assert_eq!(
            stream.next().map(|(t, _)| t),
            Some(Token::Identifier("y".to_string()))
        );
    }

    proptest! {
        #[test]
        fn test_tokenize_identifiers(s in "[a-zA-Z_][a-zA-Z0-9_]{0,100}") {
            // Skip reserved keywords that should tokenize as their respective tokens
            let reserved_keywords = [
                "true", "false", "fun", "fn", "let", "var", "mod", "if", "else", "match",
                "for", "in", "while", "loop", "async", "await", "throw", "try", "catch",
                "return", "command", "Ok", "Err", "Some", "None", "null", "Result", "Option",
                "break", "continue", "struct", "enum", "impl", "trait", "extend", "actor",
                "state", "receive", "send", "ask", "type", "where", "const", "static",
                "mut", "pub", "import", "use", "as", "module", "export", "df"
            ];
            
            if reserved_keywords.contains(&s.as_str()) {
                return Ok(()); // Skip test for reserved keywords
            }
            
            let mut stream = TokenStream::new(&s);
            match stream.advance() {
                Some((Token::Identifier(id), _)) => prop_assert_eq!(id, s),
                Some((Token::Underscore, _)) if s == "_" => {}, // Special case for underscore
                _ => panic!("Failed to tokenize identifier: {s}"),
            }
        }

        #[test]
        fn test_tokenize_integers(n in 0i64..1_000_000) {
            let s = n.to_string();
            let mut stream = TokenStream::new(&s);
            match stream.advance() {
                Some((Token::Integer(i), _)) => prop_assert_eq!(i, n),
                _ => panic!("Failed to tokenize integer"),
            }
        }
    }
}
