//! Lexer for the pax DSL.
//!
//! Slice 1 handles just the tokens needed for a single `var` declaration with
//! an integer literal: the `var` keyword, identifiers, `:`, `=`, integers,
//! whitespace, and `//` line comments.

use chumsky::prelude::*;
use std::fmt;

pub type Span = SimpleSpan;
pub type Spanned<T> = (T, Span);

#[derive(Clone, Debug, PartialEq)]
pub enum Token<'src> {
    Var,
    Let,
    If,
    Else,
    Foreach,
    In,
    Until,
    Pa,
    Debug,
    Terminate,
    Switch,
    Case,
    Default,
    Scope,
    On,
    Null,
    Ident(&'src str),
    Int(i64),
    Float(f64),
    Str(String),
    Bool(bool),
    Colon,
    Eq,
    PlusEq,
    MinusEq,
    AmpEq,
    Amp,
    Plus,
    Minus,
    Star,
    Slash,
    Lt,
    Gt,
    Le,
    Ge,
    EqEq,
    BangEq,
    AmpAmp,
    PipePipe,
    Bang,
    Comma,
    Dot,
    LBracket,
    RBracket,
    LBrace,
    RBrace,
    LParen,
    RParen,
    Question,
}

impl fmt::Display for Token<'_> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Token::Var => f.write_str("`var`"),
            Token::Let => f.write_str("`let`"),
            Token::If => f.write_str("`if`"),
            Token::Else => f.write_str("`else`"),
            Token::Foreach => f.write_str("`foreach`"),
            Token::In => f.write_str("`in`"),
            Token::Until => f.write_str("`until`"),
            Token::Pa => f.write_str("`pa`"),
            Token::Debug => f.write_str("`debug`"),
            Token::Terminate => f.write_str("`terminate`"),
            Token::Switch => f.write_str("`switch`"),
            Token::Case => f.write_str("`case`"),
            Token::Default => f.write_str("`default`"),
            Token::Scope => f.write_str("`scope`"),
            Token::On => f.write_str("`on`"),
            Token::Null => f.write_str("`null`"),
            Token::Ident(s) => write!(f, "`{s}`"),
            Token::Int(n) => write!(f, "`{n}`"),
            Token::Float(x) => write!(f, "`{x}`"),
            Token::Str(s) => write!(f, "string \"{s}\""),
            Token::Bool(b) => write!(f, "`{b}`"),
            Token::Colon => f.write_str("`:`"),
            Token::Eq => f.write_str("`=`"),
            Token::PlusEq => f.write_str("`+=`"),
            Token::MinusEq => f.write_str("`-=`"),
            Token::AmpEq => f.write_str("`&=`"),
            Token::Amp => f.write_str("`&`"),
            Token::Plus => f.write_str("`+`"),
            Token::Minus => f.write_str("`-`"),
            Token::Star => f.write_str("`*`"),
            Token::Slash => f.write_str("`/`"),
            Token::Lt => f.write_str("`<`"),
            Token::Gt => f.write_str("`>`"),
            Token::Le => f.write_str("`<=`"),
            Token::Ge => f.write_str("`>=`"),
            Token::EqEq => f.write_str("`==`"),
            Token::BangEq => f.write_str("`!=`"),
            Token::AmpAmp => f.write_str("`&&`"),
            Token::PipePipe => f.write_str("`||`"),
            Token::Bang => f.write_str("`!`"),
            Token::Comma => f.write_str("`,`"),
            Token::Dot => f.write_str("`.`"),
            Token::LBracket => f.write_str("`[`"),
            Token::RBracket => f.write_str("`]`"),
            Token::LBrace => f.write_str("`{`"),
            Token::RBrace => f.write_str("`}`"),
            Token::LParen => f.write_str("`(`"),
            Token::RParen => f.write_str("`)`"),
            Token::Question => f.write_str("`?`"),
        }
    }
}

pub fn lexer<'src>()
-> impl Parser<'src, &'src str, Vec<Spanned<Token<'src>>>, extra::Err<Rich<'src, char, Span>>> {
    // Numbers: `<digits>` → Int, `<digits>.<digits>` → Float. The fractional
    // tail is `or_not` so `obj.field` and `3.foo` still tokenize as int + dot
    // + rest (or_not rewinds when `.` is present but no digits follow).
    let fraction = just('.').then(text::digits(10)).to_slice();
    let number =
        text::int(10)
            .then(fraction.or_not())
            .map(|(int_part, frac): (&str, Option<&str>)| match frac {
                Some(frac) => Token::Float(format!("{int_part}{frac}").parse::<f64>().unwrap()),
                None => Token::Int(int_part.parse::<i64>().unwrap()),
            });

    let escape = just('\\').ignore_then(choice((
        just('n').to('\n'),
        just('t').to('\t'),
        just('r').to('\r'),
        just('"').to('"'),
        just('\\').to('\\'),
    )));

    let str_char = escape.or(none_of("\\\""));

    let str_ = just('"')
        .ignore_then(str_char.repeated().collect::<String>())
        .then_ignore(just('"'))
        .map(Token::Str);

    let compound = choice((
        just("+=").to(Token::PlusEq),
        just("-=").to(Token::MinusEq),
        just("&&").to(Token::AmpAmp),
        just("&=").to(Token::AmpEq),
        just("||").to(Token::PipePipe),
        just("<=").to(Token::Le),
        just(">=").to(Token::Ge),
        just("==").to(Token::EqEq),
        just("!=").to(Token::BangEq),
    ));

    let ctrl = choice((
        just(':').to(Token::Colon),
        just('=').to(Token::Eq),
        just('&').to(Token::Amp),
        just('+').to(Token::Plus),
        just('-').to(Token::Minus),
        just('*').to(Token::Star),
        just('/').to(Token::Slash),
        just('<').to(Token::Lt),
        just('>').to(Token::Gt),
        just('!').to(Token::Bang),
        just(',').to(Token::Comma),
        just('.').to(Token::Dot),
        just('[').to(Token::LBracket),
        just(']').to(Token::RBracket),
        just('{').to(Token::LBrace),
        just('}').to(Token::RBrace),
        just('(').to(Token::LParen),
        just(')').to(Token::RParen),
        just('?').to(Token::Question),
    ));

    let ident = text::ascii::ident().map(|s: &str| match s {
        "var" => Token::Var,
        "let" => Token::Let,
        "if" => Token::If,
        "else" => Token::Else,
        "foreach" => Token::Foreach,
        "in" => Token::In,
        "until" => Token::Until,
        "pa" => Token::Pa,
        "debug" => Token::Debug,
        "terminate" => Token::Terminate,
        "switch" => Token::Switch,
        "case" => Token::Case,
        "default" => Token::Default,
        "scope" => Token::Scope,
        "on" => Token::On,
        "null" => Token::Null,
        "true" => Token::Bool(true),
        "false" => Token::Bool(false),
        _ => Token::Ident(s),
    });

    let token = number.or(str_).or(compound).or(ctrl).or(ident);

    let comment = just("//")
        .then(any().and_is(just('\n').not()).repeated())
        .padded();

    token
        .map_with(|tok, e| (tok, e.span()))
        .padded_by(comment.repeated())
        .padded()
        .repeated()
        .collect()
}

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

    fn lex(src: &str) -> Vec<Token<'_>> {
        lexer()
            .parse(src)
            .into_result()
            .expect("lex failed")
            .into_iter()
            .map(|(t, _)| t)
            .collect()
    }

    #[test]
    fn slice1_var_decl() {
        assert_eq!(
            lex("var counter: int = 1"),
            vec![
                Token::Var,
                Token::Ident("counter"),
                Token::Colon,
                Token::Ident("int"),
                Token::Eq,
                Token::Int(1),
            ]
        );
    }

    #[test]
    fn slice2_string_and_bool() {
        assert_eq!(
            lex(r#"var greeting: string = "hello""#),
            vec![
                Token::Var,
                Token::Ident("greeting"),
                Token::Colon,
                Token::Ident("string"),
                Token::Eq,
                Token::Str("hello".to_string()),
            ]
        );
        assert_eq!(
            lex("var ok: bool = true"),
            vec![
                Token::Var,
                Token::Ident("ok"),
                Token::Colon,
                Token::Ident("bool"),
                Token::Eq,
                Token::Bool(true),
            ]
        );
    }

    #[test]
    fn slice6_compound_assign_ops() {
        assert_eq!(
            lex("counter += 1"),
            vec![Token::Ident("counter"), Token::PlusEq, Token::Int(1),]
        );
        assert_eq!(
            lex("counter -= 1"),
            vec![Token::Ident("counter"), Token::MinusEq, Token::Int(1),]
        );
    }

    #[test]
    fn slice31_float_literals() {
        // Standard float literal.
        assert_eq!(
            lex("var rate: float = 1.5"),
            vec![
                Token::Var,
                Token::Ident("rate"),
                Token::Colon,
                Token::Ident("float"),
                Token::Eq,
                Token::Float(1.5),
            ]
        );
        // Leading zero in fractional part is fine.
        match lex("0.05").as_slice() {
            [Token::Float(f)] => assert!((f - 0.05).abs() < 1e-12),
            other => panic!("expected single float token, got {other:?}"),
        }
        // Trailing `.` with no digits stays int + dot (member-access path).
        assert_eq!(
            lex("obj.field"),
            vec![Token::Ident("obj"), Token::Dot, Token::Ident("field")]
        );
        assert_eq!(
            lex("3.foo"),
            vec![Token::Int(3), Token::Dot, Token::Ident("foo")]
        );
    }

    #[test]
    fn slice13_string_escapes() {
        assert_eq!(
            lex(r#""a\nb\tc\"d\\e\re""#),
            vec![Token::Str("a\nb\tc\"d\\e\re".to_string())]
        );
    }

    #[test]
    fn slice45a_question_subscript() {
        assert_eq!(
            lex(r#"obj?["key"]"#),
            vec![
                Token::Ident("obj"),
                Token::Question,
                Token::LBracket,
                Token::Str("key".to_string()),
                Token::RBracket,
            ]
        );
        assert_eq!(
            lex("arr?[0]"),
            vec![
                Token::Ident("arr"),
                Token::Question,
                Token::LBracket,
                Token::Int(0),
                Token::RBracket,
            ]
        );
    }

    #[test]
    fn skips_line_comment() {
        assert_eq!(
            lex("// hello\nvar x: int = 42"),
            vec![
                Token::Var,
                Token::Ident("x"),
                Token::Colon,
                Token::Ident("int"),
                Token::Eq,
                Token::Int(42),
            ]
        );
    }
}
