#[cfg(test)]
mod test;

use std::marker::PhantomData;

use crate::{
    interner,
    lexer::{
        span::Span,
        token::{Keyword, Literal, Token, TokenKind, TokenTag},
    },
    source::CharSource,
};

pub mod span;
pub mod token;

pub struct Lexer<'source, Kws: Keyword> {
    source: CharSource<'source>,
    marker: PhantomData<Kws>,
    peeked: Option<Token<Kws>>,
    last_span: Span,
}

impl<'source, Kws: Keyword> Lexer<'source, Kws> {
    pub fn from(source: &'source str) -> Self {
        Self {
            last_span: (0..0).into(),
            marker: PhantomData,
            source: CharSource::from(source),
            peeked: None,
        }
    }

    pub fn last_span(&self) -> Span {
        self.last_span
    }

    pub fn char_before_span(&self, span: Span) -> Span {
        let slice = &self.source.source()[..span.start];
        slice
            .char_indices()
            .next_back()
            .map(|(offset, c)| (offset..offset + c.len_utf8()).into())
            .unwrap_or(Span::new(span.start, span.start))
    }

    pub fn last_char_span_of(&self, span: Span) -> Span {
        let slice = &self.source.source()[span.start..span.end];
        slice
            .char_indices()
            .next_back()
            .map(|(offset, c)| {
                let start = span.start + offset;
                (start..start + c.len_utf8()).into()
            })
            .unwrap_or(span)
    }

    pub fn last_char_span(&self) -> Span {
        self.last_char_span_of(self.last_span)
    }

    pub fn span_before(&self, span: Span) -> Span {
        (self.last_span.end..span.start).into()
    }

    pub fn spanning(&self, span: Span) -> &'source str {
        self.source.spanning(span)
    }

    pub fn pos(&self) -> usize {
        self.source.pos()
    }

    pub fn peek(&mut self) -> Token<Kws> {
        self.peeked
            .get_or_insert_with(|| Self::raw_next(&mut self.source))
            .clone()
    }

    pub fn bump(&mut self) -> Token<Kws> {
        let tok = self
            .peeked
            .take()
            .unwrap_or_else(|| Self::raw_next(&mut self.source));
        self.last_span = tok.span();
        tok
    }

    pub fn at(&mut self, tag: TokenTag<Kws>) -> bool {
        self.peek().kind().tag().eq(&tag)
    }

    pub fn at_kw(&mut self, tag: Kws::Tag) -> bool {
        matches!(self.peek().kind(), TokenKind::Keyword(kw) if kw.tag().eq(&tag))
    }

    pub fn expect(&mut self, tag: TokenTag<Kws>) -> Result<Token<Kws>, Token<Kws>> {
        let peek = self.peek();
        if peek.kind().tag().eq(&tag) {
            Ok(self.bump())
        } else {
            Err(peek)
        }
    }

    pub fn expect_kw(&mut self, tag: Kws::Tag) -> Result<Token<Kws>, Token<Kws>> {
        let peek = self.peek();
        if let TokenKind::Keyword(kw) = peek.kind()
            && kw.tag().eq(&tag)
        {
            Ok(self.bump())
        } else {
            Err(peek)
        }
    }

    pub fn eat(&mut self, tag: TokenTag<Kws>) -> bool {
        self.expect(tag).is_ok()
    }

    pub fn eat_kw(&mut self, tag: Kws::Tag) -> bool {
        self.expect_kw(tag).is_ok()
    }

    pub fn eat_while<F>(&mut self, f: F)
    where
        F: Fn(&Token<Kws>) -> bool,
    {
        while f(&self.peek()) {
            self.bump();
        }
    }

    pub fn is_finished(&mut self) -> bool {
        matches!(self.peek().kind(), TokenKind::Eof)
    }

    fn raw_next(source: &mut CharSource<'source>) -> Token<Kws> {
        let start = source.pos();
        Self::skip_whitespace(source);
        Self::lex_eof(source)
            .or_else(|| Self::lex_op(source))
            .or_else(|| Self::lex_punct(source))
            .or_else(|| Self::lex_literal(source))
            .or_else(|| Self::lex_words(source))
            .unwrap_or_else(|| Self::error_token(source, start))
    }

    fn error_token(source: &mut CharSource<'source>, start: usize) -> Token<Kws> {
        Token::new(start..source.pos(), TokenKind::Error)
    }

    fn lex_eof(source: &mut CharSource<'source>) -> Option<Token<Kws>> {
        if source.is_empty() {
            let len = source.source().len().saturating_sub(1);
            Some(Token::new(len..len, TokenKind::Eof))
        } else {
            None
        }
    }

    fn lex_punct(source: &mut CharSource<'source>) -> Option<Token<Kws>> {
        let start = source.pos();
        source
            .try_next(|c, _| match c {
                '{' => Some(TokenKind::LeftBrace),
                '}' => Some(TokenKind::RightBrace),
                '(' => Some(TokenKind::LeftParen),
                ')' => Some(TokenKind::RightParen),
                '[' => Some(TokenKind::LeftBracket),
                ']' => Some(TokenKind::RightBracket),
                ':' => Some(TokenKind::Colon),
                ',' => Some(TokenKind::Comma),
                ';' => Some(TokenKind::Semicolon),
                '.' => Some(TokenKind::Dot),
                '#' => Some(TokenKind::HashTag),
                _ => None,
            })
            .map(|kind| Token::new(start..source.pos(), kind))
    }

    fn lex_op(source: &mut CharSource<'source>) -> Option<Token<Kws>> {
        let start = source.pos();
        source
            .try_next(|c, src| match c {
                '=' => Some(
                    src.try_next(|c, _| match c {
                        '=' => Some(TokenKind::Eq),
                        '>' => Some(TokenKind::StrongArrow),
                        _ => None,
                    })
                    .unwrap_or(TokenKind::Assign),
                ),
                '>' => Some(
                    src.try_next(|c, _| match c {
                        '=' => Some(TokenKind::Ge),
                        _ => None,
                    })
                    .unwrap_or(TokenKind::Gt),
                ),
                '<' => Some(
                    src.try_next(|c, _| match c {
                        '=' => Some(TokenKind::Le),
                        _ => None,
                    })
                    .unwrap_or(TokenKind::Lt),
                ),
                '-' => Some(
                    src.try_next(|c, _| match c {
                        '>' => Some(TokenKind::Arrow),
                        _ => None,
                    })
                    .unwrap_or(TokenKind::Sub),
                ),
                '+' => Some(TokenKind::Add),
                '*' => Some(TokenKind::Mul),
                '/' => Some(TokenKind::Div),
                '%' => Some(TokenKind::Rem),
                '!' => Some(
                    src.try_next(|c, _| match c {
                        '=' => Some(TokenKind::Ne),
                        _ => None,
                    })
                    .unwrap_or(TokenKind::Not),
                ),
                '|' => Some(
                    src.try_next(|c, _| match c {
                        '|' => Some(TokenKind::DoublePipe),
                        _ => None,
                    })
                    .unwrap_or(TokenKind::Pipe),
                ),
                '&' => Some(
                    src.try_next(|c, _| match c {
                        '&' => Some(TokenKind::And),
                        _ => None,
                    })
                    .unwrap_or(TokenKind::Ampersand),
                ),
                _ => None,
            })
            .map(|kind| Token::new(start..source.pos(), kind))
    }

    fn lex_ident(source: &mut CharSource<'source>) -> Option<Span> {
        let start = source.pos();
        source.next_if(unicode_ident::is_xid_start)?;
        while source.next_if(unicode_ident::is_xid_continue).is_some() {}
        Some((start..source.pos()).into())
    }

    fn matches_bool(source: &mut CharSource<'source>, span: Span) -> Option<TokenKind<Kws>> {
        let ident = source.spanning(span);
        match ident {
            "true" => Some(TokenKind::Literal(Literal::Bool(true))),
            "false" => Some(TokenKind::Literal(Literal::Bool(false))),
            _ => None,
        }
    }

    fn lex_words(source: &mut CharSource<'source>) -> Option<Token<Kws>> {
        let span = Self::lex_ident(source)?;
        let ident = source.spanning(span);
        Self::matches_bool(source, span)
            .or_else(|| Kws::matches(ident).map(|kws| TokenKind::Keyword(kws)))
            .or_else(|| {
                let sym = interner::intern_ident(ident);
                Some(TokenKind::Ident(sym))
            })
            .map(|kind| Token::new(span, kind))
    }

    fn lex_literal(source: &mut CharSource<'source>) -> Option<Token<Kws>> {
        Self::lex_string(source)
            .or_else(|| Self::lex_char(source))
            .or_else(|| Self::lex_number(source))
    }

    fn lex_string(source: &mut CharSource<'source>) -> Option<Token<Kws>> {
        let start = source.pos();
        source
            .try_next(|c, src| match c {
                '"' => loop {
                    match src.next() {
                        None => break None,
                        Some('"') => break Some(()),
                        Some('\\') => src.advance(),
                        _ => (),
                    }
                },
                _ => None,
            })
            .map(|_| {
                let span: Span = (start..source.pos()).into();
                let string = source.spanning(span);
                let sym =
                    interner::intern_str(string).expect("string should be valid at this point");
                Token::new(span, TokenKind::Literal(Literal::String(sym)))
            })
    }

    fn lex_char(source: &mut CharSource<'source>) -> Option<Token<Kws>> {
        let start = source.pos();
        source.next_if(|c| c == '\'')?;
        source
            .try_next(|c, src| {
                let value = match c {
                    '\\' => match src.next() {
                        None => return None,
                        Some('n') => '\n',
                        Some('t') => '\t',
                        Some('r') => '\r',
                        Some('\\') => '\\',
                        Some('\'') => '\'',
                        Some('"') => '"',
                        Some(_) => return None,
                    },
                    c => c,
                };
                // consume closing quote for both branches
                src.next_if(|c| c == '\'').map(|_| value)
            })
            .map(|c| {
                let span = start..source.pos();
                Token::new(span, TokenKind::Literal(Literal::Char(c)))
            })
    }

    // first we check if the first char is a digit
    // if it isn't return early
    // then we check if that digit and the next one are any of the
    // radix variants aka a 0 follows by
    // bnf:
    // <number>     ::= <hex> | <oct> | <bin> | <dec>
    //
    // <hex>        ::= "0x" ( <hex_digit> )+
    // <hex_digit>  ::= "0".."9" | "a" .."f" | "A".."F"
    //
    // <oct>        ::= "0o" ( <oct_digit> )+
    // <oct_digit>  ::= "0".."7"
    //
    // <bin>        ::= "0b" ( <bin_digit> )+
    // <bin_digit>  ::= "0" | "1"
    //
    // <dec>        ::= ( <dec_digit> )+ ("." (<dec_digit>)* )?
    // <dec_digit>  ::= "0" .. "9"

    fn lex_radix(
        source: &mut CharSource<'source>,
        radix: u32,
        kind: &str,
    ) -> Result<TokenKind<Kws>, String> {
        let start = source.pos();
        // check if we have at least one digit
        source
            .peek()
            .filter(|c| c.is_digit(radix))
            .ok_or_else(|| format!("expected {kind} digit"))?;
        while source.next_if(|c| c.is_digit(radix)).is_some() {}
        let span: Span = (start..source.pos()).into();
        let int = u64::from_str_radix(source.spanning(span), radix).map_err(|e| e.to_string())?;
        Ok(TokenKind::Literal(Literal::Int(int)))
    }

    fn lex_dec(source: &mut CharSource<'source>) -> Result<TokenKind<Kws>, String> {
        let radix = 10;
        let start = source.pos();
        // check if we have at least one digit
        source
            .peek()
            .filter(|c| c.is_digit(radix))
            .ok_or("expected decimal digit")?;
        while source.next_if(|c| c.is_digit(radix)).is_some() {}

        // now, we need to diverge
        // if we find a '.' we need to keep parsing
        // if we don't, we just return a IntLit

        let token = match source.next_if(|c| c == '.') {
            Some(_) => {
                while source.next_if(|c| c.is_digit(radix)).is_some() {}
                let span: Span = (start..source.pos()).into();
                let float = source
                    .spanning(span)
                    .parse::<f64>()
                    .map_err(|e| e.to_string())?;
                TokenKind::Literal(Literal::Float(float))
            }
            None => {
                let span: Span = (start..source.pos()).into();
                let int =
                    u64::from_str_radix(source.spanning(span), radix).map_err(|e| e.to_string())?;
                TokenKind::Literal(Literal::Int(int))
            }
        };

        Ok(token)
    }

    fn lex_number(source: &mut CharSource<'source>) -> Option<Token<Kws>> {
        let start = source.pos();
        let kind = match source.peek_by::<2>() {
            Some(['0', 'x']) => {
                source.advance_by(2);
                Self::lex_radix(source, 16, "hexadecimal")
            }
            Some(['0', 'o']) => {
                source.advance_by(2);
                Self::lex_radix(source, 8, "octonal")
            }
            Some(['0', 'b']) => {
                source.advance_by(2);
                Self::lex_radix(source, 2, "binary")
            }
            _ => Self::lex_dec(source),
        }
        .ok()?;

        let span: Span = (start..source.pos()).into();
        Some(Token::new(span, kind))
    }

    fn skip_whitespace(source: &mut CharSource<'source>) {
        loop {
            // Skip whitespace
            while source.next_if(char::is_whitespace).is_some() {}
            // Skip line comments // ...
            if source.peek_by::<2>() == Some(['/', '/']) {
                source.advance_by(2);
                while let Some(c) = source.peek() {
                    if c == '\n' {
                        break;
                    }
                    source.advance();
                }
                continue;
            }
            // Skip block comments /* ... */
            if source.peek_by::<2>() == Some(['/', '*']) {
                source.advance_by(2);
                loop {
                    match source.peek_by::<2>() {
                        Some(['*', '/']) => {
                            source.advance_by(2);
                            break;
                        }
                        Some(_) => {
                            source.advance();
                        }
                        None => break,
                    }
                }
                continue;
            }
            break;
        }
    }
}

impl<'source, Kws: Keyword> Iterator for Lexer<'source, Kws> {
    type Item = Token<Kws>;

    fn next(&mut self) -> Option<Self::Item> {
        let next = self.bump();
        match next.kind() {
            TokenKind::Eof => None,
            _ => Some(next),
        }
    }
}
