use chumsky::prelude::*;
use chumsky::span::SimpleSpan;

/// Keywords and reserved type names cannot be used as identifiers (variable names, param names, let bindings).
/// Keywords are recognised as distinct tokens; type names are emitted as `DTypeName` so they are not `Ident`.
pub const KEYWORDS: &[&str] = &["if", "then", "else", "cast", "true", "false", "e", "pi"];

/// Reserved type names (dtype literals). Must match `crate::ir::dtype::all_dtype_names()`.
pub const RESERVED_TYPE_NAMES: &[&str] = &[
    "bool", "u8", "u16", "u32", "u64", "i8", "i16", "i32", "i64", "f16", "bf16", "f32", "f64",
    "c64", "c128",
];

#[derive(Debug, Clone, PartialEq)]
pub enum TokenKind {
    If,
    Then,
    Else,
    Cast,
    True,
    False,
    ConstE,
    ConstPi,
    /// User identifier (variable / param / call name). Not a keyword or reserved type name.
    Ident(String),
    /// Reserved type name (e.g. f64, i32). Accepted only in type positions; cannot be used as identifier.
    DTypeName(String),
    Int(i64),
    Float(f64),
    Complex(f64),
    Plus,
    Minus,
    Star,
    Slash,
    SlashSlash,
    Percent,
    Pow,
    AndAnd,
    OrOr,
    Bang,
    EqEq,
    Ne,
    Le,
    Ge,
    Lt,
    Gt,
    LParen,
    RParen,
    Comma,
    At,
}

pub(crate) fn is_ident_start(c: char) -> bool {
    c == '_' || c.is_ascii_alphabetic()
}

pub(crate) fn is_ident_continue(c: char) -> bool {
    c == '_' || c.is_ascii_alphanumeric()
}

pub fn lexer<'src>()
-> impl Parser<'src, &'src str, Vec<(TokenKind, SimpleSpan<usize>)>, extra::Err<Rich<'src, char>>> {
    let int_lit =
        text::int(10)
            .to_slice()
            .validate(|lit: &str, e, emitter| match lit.parse::<i64>() {
                Ok(value) => TokenKind::Int(value),
                Err(_) => {
                    emitter.emit(Rich::custom(e.span(), "invalid integer literal"));
                    TokenKind::Int(0)
                }
            });

    let digits = text::digits(10).to_slice();
    let exp = just('e')
        .or(just('E'))
        .then(one_of("+-").or_not())
        .then(digits)
        .to_slice()
        .map(ToString::to_string);

    let float_with_dot = digits
        .then(just('.').then(digits))
        .then(exp.or_not())
        .to_slice()
        .map(ToString::to_string);

    let float_with_exp = digits.then(exp).to_slice().map(ToString::to_string);

    let float_lit = float_with_dot.or(float_with_exp);

    let complex_lit =
        float_lit
            .then_ignore(just('j'))
            .validate(|lit: String, e, emitter| match lit.parse::<f64>() {
                Ok(value) => TokenKind::Complex(value),
                Err(_) => {
                    emitter.emit(Rich::custom(e.span(), "invalid complex literal"));
                    TokenKind::Complex(0.0)
                }
            });

    let float_token = float_lit.validate(|lit: String, e, emitter| match lit.parse::<f64>() {
        Ok(value) => TokenKind::Float(value),
        Err(_) => {
            emitter.emit(Rich::custom(e.span(), "invalid float literal"));
            TokenKind::Float(0.0)
        }
    });

    let ident = text::ident().map(|ident: &str| {
        if RESERVED_TYPE_NAMES.contains(&ident) {
            TokenKind::DTypeName(ident.to_string())
        } else if KEYWORDS.contains(&ident) {
            match ident {
                "if" => TokenKind::If,
                "then" => TokenKind::Then,
                "else" => TokenKind::Else,
                "cast" => TokenKind::Cast,
                "true" => TokenKind::True,
                "false" => TokenKind::False,
                "e" => TokenKind::ConstE,
                "pi" => TokenKind::ConstPi,
                _ => TokenKind::Ident(ident.to_string()),
            }
        } else {
            TokenKind::Ident(ident.to_string())
        }
    });

    let op = choice((
        just("**").to(TokenKind::Pow),
        just("//").to(TokenKind::SlashSlash),
        just("&&").to(TokenKind::AndAnd),
        just("||").to(TokenKind::OrOr),
        just("==").to(TokenKind::EqEq),
        just("!=").to(TokenKind::Ne),
        just("<=").to(TokenKind::Le),
        just(">=").to(TokenKind::Ge),
        just('+').to(TokenKind::Plus),
        just('-').to(TokenKind::Minus),
        just('*').to(TokenKind::Star),
        just('/').to(TokenKind::Slash),
        just('%').to(TokenKind::Percent),
        just('!').to(TokenKind::Bang),
        just('<').to(TokenKind::Lt),
        just('>').to(TokenKind::Gt),
        just('(').to(TokenKind::LParen),
        just(')').to(TokenKind::RParen),
        just(',').to(TokenKind::Comma),
        just('@').to(TokenKind::At),
    ));

    choice((complex_lit, float_token, int_lit, ident, op))
        .map_with(|tok, e| (tok, e.span()))
        .padded()
        .repeated()
        .collect()
}
