use crate::token::{Token, TokenKind};

pub struct Lexer<'a> {
    input: &'a str,
    position: usize,
}

impl<'a> Lexer<'a> {
    pub fn new(input: &'a str) -> Self {
        Self { input, position: 0 }
    }

    pub fn next(&mut self) -> Option<Result<Token, LexerError>> {
        let maybe_chunk = self.input.get(self.position..);
        if maybe_chunk.is_none() {
            return None;
        }

        let mut chunk = maybe_chunk.unwrap();
        if chunk.is_empty() {
            // in case we're done
            return None;
        }
        while chunk.starts_with(&[' ', '\t', '\r', '\n']) {
            // ignore whitespace
            self.position += 1;
            let maybe_chunk = self.input.get(self.position..);
            if maybe_chunk.is_none() {
                return None;
            }
            chunk = maybe_chunk.unwrap();
            if chunk.is_empty() {
                return None;
            }
        }

        let mut chars = chunk.chars();
        let current = chars.next().unwrap(); // safe because we return None if chunk is empty above
        let peek = chars.next();

        let (end, kind) = match current {
            '-' => {
                if let Some(peek) = peek {
                    if peek == '>' {
                        (self.position + 2, TokenKind::RightArrow)
                    } else {
                        (self.position + 1, TokenKind::Minus)
                    }
                } else {
                    (self.position + 1, TokenKind::Minus)
                }
            }
            '(' => (self.position + 1, TokenKind::LeftParentheses),
            ')' => (self.position + 1, TokenKind::RightParentheses),
            '[' => (self.position + 1, TokenKind::LeftSquare),
            ']' => (self.position + 1, TokenKind::RightSquare),
            '<' => (self.position + 1, TokenKind::Lt),
            '>' => (self.position + 1, TokenKind::Gt),
            '*' => (self.position + 1, TokenKind::Star),
            '/' => (self.position + 1, TokenKind::Slash),
            '%' => (self.position + 1, TokenKind::Modulo),
            '+' => (self.position + 1, TokenKind::Plus),
            '&' => (self.position + 1, TokenKind::Ampersand),
            '^' => (self.position + 1, TokenKind::Caret),
            '|' => (self.position + 1, TokenKind::Bar),
            '=' => (self.position + 1, TokenKind::Equals),
            ',' => (self.position + 1, TokenKind::Comma),
            '.' => (self.position + 1, TokenKind::Dot),
            ':' => (self.position + 1, TokenKind::Colon),
            _ => {
                // more than one char token
                if chunk.starts_with(|c: char| c.is_ascii_digit()) {
                    // number literal
                    let mut end = self.position;
                    let mut whole_chars = chunk.chars();
                    while let Some(c) = whole_chars.next() {
                        if c.is_ascii_digit() {
                            end += 1;
                        } else {
                            break;
                        }
                    }

                    let mut fraction_chars = chunk.chars().skip(end - self.position);
                    if fraction_chars.next().map(|c| c == '.').unwrap_or(false) {
                        let mut found_fraction = false;
                        while let Some(c) = fraction_chars.next() {
                            if c.is_ascii_digit() {
                                end += 1;
                                if !found_fraction {
                                    end += 1; // to account for decimal point
                                    found_fraction = true;
                                }
                            } else {
                                break;
                            }
                        }
                        if found_fraction {
                            (end, TokenKind::DoubleLiteral)
                        } else {
                            (end, TokenKind::IntegerLiteral)
                        }
                    } else {
                        (end, TokenKind::IntegerLiteral)
                    }
                } else if chunk.starts_with("\"") {
                    // string literal
                    let mut end = self.position;
                    let mut terminated = false;
                    let mut chars = chunk.chars();
                    chars.next(); // skip opening quote
                    end += 1;
                    for char in chars {
                        end += 1;
                        if char == '"' {
                            terminated = true;
                            break;
                        }
                    }
                    if !terminated {
                        return Some(Err(LexerError::new(
                            self.position,
                            end,
                            LexerErrorKind::UnterminatedString,
                        )));
                    }
                    (end, TokenKind::String)
                } else {
                    // keyword or identifier
                    let mut prefix = String::new();
                    for char in chunk.chars() {
                        if char.is_alphanumeric() || char == '_' {
                            prefix.push(char);
                        } else {
                            break;
                        }
                    }

                    if prefix.len() == 0 {
                        return Some(Err(LexerError::new(
                            self.position,
                            self.position + 1,
                            LexerErrorKind::UnrecognizedCharacter(chunk.chars().next().unwrap()),
                        )));
                    }

                    let token_kind = match prefix.as_str() {
                        "fn" => TokenKind::Fn,
                        "end" => TokenKind::End,
                        "let" => TokenKind::Let,
                        "extern" => TokenKind::Extern,
                        "class" => TokenKind::Class,
                        "self" => TokenKind::SelfKw,
                        "pub" => TokenKind::Public,
                        "mut" => TokenKind::Mut,
                        "ctor" => TokenKind::Ctor,
                        _ => TokenKind::Identifier,
                    };
                    (self.position + prefix.len(), token_kind)
                }
            }
        };
        let token = Token::new(self.position, end, kind);
        self.position += token.end() - token.start();
        Some(Ok(token))
    }
}

#[derive(Debug, Eq, PartialEq)]
pub struct LexerError {
    start: usize,
    end: usize,
    kind: LexerErrorKind,
}

impl LexerError {
    pub fn new(start: usize, end: usize, kind: LexerErrorKind) -> Self {
        Self { start, end, kind }
    }

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

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

    pub fn kind(&self) -> LexerErrorKind {
        self.kind
    }
}

#[derive(Debug, Clone, Copy, Eq, PartialEq)]
pub enum LexerErrorKind {
    UnterminatedString,
    UnrecognizedCharacter(char),
}

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

    fn assert_next(lexer: &mut Lexer, kind: TokenKind, length: usize) {
        let token = lexer.next().unwrap().unwrap();
        println!("{:?}", token);
        assert_eq!(token.kind(), kind);
        assert_eq!(token.end() - token.start(), length);
    }

    #[test]
    fn forty_two() {
        let mut lexer = Lexer::new("fn main() 42 end");
        assert_next(&mut lexer, TokenKind::Fn, 2);
        assert_next(&mut lexer, TokenKind::Identifier, 4);
        assert_next(&mut lexer, TokenKind::LeftParentheses, 1);
        assert_next(&mut lexer, TokenKind::RightParentheses, 1);
        assert_next(&mut lexer, TokenKind::IntegerLiteral, 2);
        assert_next(&mut lexer, TokenKind::End, 3);
        assert_eq!(lexer.next(), None);
    }

    #[test]
    fn hello_world() {
        let mut lexer = Lexer::new("fn main() println(\"Hello, World!\") end");
        assert_next(&mut lexer, TokenKind::Fn, 2);
        assert_next(&mut lexer, TokenKind::Identifier, 4);
        assert_next(&mut lexer, TokenKind::LeftParentheses, 1);
        assert_next(&mut lexer, TokenKind::RightParentheses, 1);
        assert_next(&mut lexer, TokenKind::Identifier, 7);
        assert_next(&mut lexer, TokenKind::LeftParentheses, 1);
        assert_next(&mut lexer, TokenKind::String, 15);
        assert_next(&mut lexer, TokenKind::RightParentheses, 1);
        assert_next(&mut lexer, TokenKind::End, 3);
        assert_eq!(lexer.next(), None);
    }

    #[test]
    fn blank_after_last_token_returns_none() {
        let mut lexer = Lexer::new("fn ");
        assert_next(&mut lexer, TokenKind::Fn, 2);
        assert_eq!(lexer.next(), None);
    }

    #[test]
    fn extern_returned() {
        let mut lexer = Lexer::new("extern");
        assert_next(&mut lexer, TokenKind::Extern, 6);
    }

    #[test]
    fn double_literal() {
        let mut lexer = Lexer::new("123.45");
        assert_next(&mut lexer, TokenKind::DoubleLiteral, 6);
    }

    #[test]
    fn one_digit_whole_part_double() {
        let mut lexer = Lexer::new("1.0");
        assert_next(&mut lexer, TokenKind::DoubleLiteral, 3);
    }

    #[test]
    fn two_digits_fraction_double() {
        let mut lexer = Lexer::new("1.23");
        assert_next(&mut lexer, TokenKind::DoubleLiteral, 4);
    }

    #[test]
    fn int_literal() {
        let mut lexer = Lexer::new("123");
        assert_next(&mut lexer, TokenKind::IntegerLiteral, 3);
    }

    #[test]
    fn doubles() {
        let mut lexer = Lexer::new("println(1.23 + 2.34)");
        assert_next(&mut lexer, TokenKind::Identifier, 7);
        assert_next(&mut lexer, TokenKind::LeftParentheses, 1);
        assert_next(&mut lexer, TokenKind::DoubleLiteral, 4);
        assert_next(&mut lexer, TokenKind::Plus, 1);
        assert_next(&mut lexer, TokenKind::DoubleLiteral, 4);
        assert_next(&mut lexer, TokenKind::RightParentheses, 1);
        assert!(lexer.next().is_none());
    }

    #[test]
    fn class_with_fields_smoke_test() {
        let mut lexer = Lexer::new(
            "
            class Foo
                pub bar: Int
                mut baz: String
                car = 42                

                fn new(bar: Int, baz: String)
                    self.bar = bar
                    self.baz = baz
                end
            end
            ",
        );
        while let Some(result) = lexer.next() {
            match result {
                Ok(_) => {}
                Err(lexer_error) => {
                    panic!("{:?}", lexer_error);
                }
            }
        }
        assert!(lexer.next().is_none());
    }
}