//! Main lexer.
//!
//! Entry point [`lex`] takes a source string and returns the full token
//! stream (terminated with [`TokenKind::Eof`]).  Internal state is threaded
//! purely through recursion as a [`LexState`] struct.

use crate::error::Error;
use crate::keyword;
use crate::number;
use crate::regex;
use crate::string;
use crate::template;
use crate::token::{Token, TokenKind};
use ecma_syntax_cat::span::{Position, Span};

/// Tokenise an ECMAScript source string.
///
/// # Errors
///
/// See [`Error`] for the failure modes.
pub fn lex(source: &str) -> Result<Vec<Token>, Error> {
    let initial = LexState {
        src: source.as_bytes(),
        pos: 0,
        line: 1,
        column: 1,
        tokens: Vec::new(),
        last_kind: None,
        template_stack: Vec::new(),
        brace_depth: 0,
    };
    let final_state = step(initial)?;
    let eof_span = Span::new(final_state.position(), final_state.position());
    let tokens_with_eof: Vec<Token> = final_state
        .tokens
        .into_iter()
        .chain(std::iter::once(Token::new(TokenKind::Eof, eof_span)))
        .collect();
    Ok(tokens_with_eof)
}

/// One peek result.
enum Step {
    End,
    Byte(u8),
}

/// Threaded lexer state.  Every state-modifying operation returns a new
/// `LexState` rather than mutating the receiver, mirroring the persistent
/// data structures the project's spike crates established.
struct LexState<'a> {
    src: &'a [u8],
    pos: usize,
    line: u32,
    column: u32,
    tokens: Vec<Token>,
    last_kind: Option<TokenKind>,
    template_stack: Vec<usize>,
    brace_depth: usize,
}

impl LexState<'_> {
    fn position(&self) -> Position {
        Position::new(
            self.line,
            self.column,
            u32::try_from(self.pos).unwrap_or(u32::MAX),
        )
    }

    fn peek(&self) -> Step {
        self.src
            .get(self.pos)
            .copied()
            .map_or(Step::End, Step::Byte)
    }

    fn at(&self, offset: usize) -> Step {
        self.src
            .get(self.pos + offset)
            .copied()
            .map_or(Step::End, Step::Byte)
    }

    fn advance(self, bytes: usize, columns: u32) -> Self {
        Self {
            pos: self.pos + bytes,
            column: self.column + columns,
            ..self
        }
    }

    fn newline(self, eol_bytes: usize) -> Self {
        Self {
            pos: self.pos + eol_bytes,
            line: self.line + 1,
            column: 1,
            ..self
        }
    }

    fn emit(self, kind: TokenKind, span: Span) -> Self {
        let last_kind = Some(kind.clone());
        let updated_brace = update_brace(self.brace_depth, &kind);
        let new_tokens: Vec<Token> = self
            .tokens
            .into_iter()
            .chain(std::iter::once(Token::new(kind, span)))
            .collect();
        Self {
            tokens: new_tokens,
            last_kind,
            brace_depth: updated_brace,
            ..self
        }
    }

    fn jump_to(self, new_pos: usize, new_line: u32, new_column: u32) -> Self {
        Self {
            pos: new_pos,
            line: new_line,
            column: new_column,
            ..self
        }
    }
}

fn update_brace(depth: usize, kind: &TokenKind) -> usize {
    match kind {
        TokenKind::LBrace => depth + 1,
        TokenKind::RBrace => depth.saturating_sub(1),
        _other => depth,
    }
}

fn step(state: LexState<'_>) -> Result<LexState<'_>, Error> {
    match state.peek() {
        Step::End => Ok(state),
        Step::Byte(b) => dispatch(state, b),
    }
}

fn dispatch(state: LexState<'_>, byte: u8) -> Result<LexState<'_>, Error> {
    match byte {
        b' ' | b'\t' => step(state.advance(1, 1)),
        b'\n' => step(state.newline(1)),
        b'\r' => {
            let lf_follows = matches!(state.at(1), Step::Byte(b'\n'));
            let bytes = if lf_follows { 2 } else { 1 };
            step(state.newline(bytes))
        }
        b'/' => dispatch_slash(state),
        b'\'' | b'"' => emit_string(state, byte),
        b'`' => emit_template_open(state),
        b'0'..=b'9' => emit_number(state),
        b'.' => dispatch_dot(state),
        b'_' | b'$' | b'a'..=b'z' | b'A'..=b'Z' => emit_identifier(state),
        b'#' => emit_private_identifier(state),
        b'(' => emit_single(state, TokenKind::LParen),
        b')' => emit_single(state, TokenKind::RParen),
        b'[' => emit_single(state, TokenKind::LBracket),
        b']' => emit_single(state, TokenKind::RBracket),
        b'{' => emit_single(state, TokenKind::LBrace),
        b'}' => dispatch_close_brace(state),
        b',' => emit_single(state, TokenKind::Comma),
        b';' => emit_single(state, TokenKind::Semicolon),
        b':' => emit_single(state, TokenKind::Colon),
        b'~' => emit_single(state, TokenKind::Tilde),
        b'?' => dispatch_question(state),
        b'=' => dispatch_eq(state),
        b'!' => dispatch_bang(state),
        b'<' => dispatch_lt(state),
        b'>' => dispatch_gt(state),
        b'+' => dispatch_plus(state),
        b'-' => dispatch_minus(state),
        b'*' => dispatch_star(state),
        b'%' => dispatch_percent(state),
        b'&' => dispatch_amp(state),
        b'|' => dispatch_pipe(state),
        b'^' => dispatch_caret(state),
        other => Err(Error::UnexpectedChar {
            at: state.position(),
            ch: char::from(other),
        }),
    }
}

fn emit_single(state: LexState<'_>, kind: TokenKind) -> Result<LexState<'_>, Error> {
    let start = state.position();
    let next = state.advance(1, 1);
    let end = next.position();
    step(next.emit(kind, Span::new(start, end)))
}

fn emit_two_byte(state: LexState<'_>, kind: TokenKind) -> Result<LexState<'_>, Error> {
    let start = state.position();
    let next = state.advance(2, 2);
    let end = next.position();
    step(next.emit(kind, Span::new(start, end)))
}

fn emit_three_byte(state: LexState<'_>, kind: TokenKind) -> Result<LexState<'_>, Error> {
    let start = state.position();
    let next = state.advance(3, 3);
    let end = next.position();
    step(next.emit(kind, Span::new(start, end)))
}

fn emit_four_byte(state: LexState<'_>, kind: TokenKind) -> Result<LexState<'_>, Error> {
    let start = state.position();
    let next = state.advance(4, 4);
    let end = next.position();
    step(next.emit(kind, Span::new(start, end)))
}

fn dispatch_slash(state: LexState<'_>) -> Result<LexState<'_>, Error> {
    match state.at(1) {
        Step::Byte(b'/') => skip_line_comment(state),
        Step::Byte(b'*') => skip_block_comment(state),
        Step::Byte(b'=') => emit_two_byte(state, TokenKind::SlashEq),
        Step::End | Step::Byte(_) => {
            let allow_regex = state
                .last_kind
                .as_ref()
                .is_none_or(TokenKind::allows_regex_after);
            if allow_regex {
                emit_regex(state)
            } else {
                emit_single(state, TokenKind::Slash)
            }
        }
    }
}

fn skip_line_comment(state: LexState<'_>) -> Result<LexState<'_>, Error> {
    let start_pos = state.pos;
    let end = scan_until_eol(state.src, start_pos);
    let delta = end - start_pos;
    let consumed_cols = u32::try_from(delta).unwrap_or(u32::MAX);
    step(state.advance(delta, consumed_cols))
}

fn skip_block_comment(state: LexState<'_>) -> Result<LexState<'_>, Error> {
    scan_block_comment(
        state.src,
        state.pos + 2,
        state.line,
        state.column + 2,
        state.position(),
    )
    .map(|(new_pos, new_line, new_column)| state.jump_to(new_pos, new_line, new_column))
    .and_then(step)
}

fn scan_until_eol(src: &[u8], from: usize) -> usize {
    match src.get(from).copied() {
        None | Some(b'\n' | b'\r') => from,
        Some(_other) => scan_until_eol(src, from + 1),
    }
}

fn scan_block_comment(
    src: &[u8],
    pos: usize,
    line: u32,
    column: u32,
    at: Position,
) -> Result<(usize, u32, u32), Error> {
    match src.get(pos).copied() {
        None => Err(Error::UnterminatedComment { at }),
        Some(b'*') if matches!(src.get(pos + 1).copied(), Some(b'/')) => {
            Ok((pos + 2, line, column + 2))
        }
        Some(b'\n') => scan_block_comment(src, pos + 1, line + 1, 1, at),
        Some(b'\r') => {
            let lf_follows = matches!(src.get(pos + 1).copied(), Some(b'\n'));
            let bytes = if lf_follows { 2 } else { 1 };
            scan_block_comment(src, pos + bytes, line + 1, 1, at)
        }
        Some(_other) => scan_block_comment(src, pos + 1, line, column + 1, at),
    }
}

fn dispatch_close_brace(state: LexState<'_>) -> Result<LexState<'_>, Error> {
    let inside_template = state
        .template_stack
        .last()
        .copied()
        .is_some_and(|depth| depth == state.brace_depth);
    if inside_template {
        emit_template_continue(state)
    } else {
        emit_single(state, TokenKind::RBrace)
    }
}

fn emit_string(state: LexState<'_>, _quote: u8) -> Result<LexState<'_>, Error> {
    let start = state.position();
    let start_pos = state.pos;
    let (content, end_pos) = string::scan(state.src, start_pos, start)?;
    let delta = end_pos - start_pos;
    let columns = u32::try_from(delta).unwrap_or(u32::MAX);
    let advanced = state.advance(delta, columns);
    let end = advanced.position();
    step(advanced.emit(TokenKind::String(content), Span::new(start, end)))
}

fn emit_template_open(state: LexState<'_>) -> Result<LexState<'_>, Error> {
    let start = state.position();
    let start_pos = state.pos;
    let (piece, end_pos) = template::scan_open(state.src, start_pos, start)?;
    let delta = end_pos - start_pos;
    let columns = u32::try_from(delta).unwrap_or(u32::MAX);
    let advanced = state.advance(delta, columns);
    let end = advanced.position();
    match piece {
        template::TemplatePiece::NoSubstitution(content) => step(advanced.emit(
            TokenKind::TemplateNoSubstitution(content),
            Span::new(start, end),
        )),
        template::TemplatePiece::Head(content) => {
            let with_token = advanced.emit(TokenKind::TemplateHead(content), Span::new(start, end));
            let pushed = push_template_marker(with_token);
            step(pushed)
        }
    }
}

fn emit_template_continue(state: LexState<'_>) -> Result<LexState<'_>, Error> {
    let start = state.position();
    let start_pos = state.pos;
    let (piece, end_pos) = template::scan_continue(state.src, start_pos, start)?;
    let delta = end_pos - start_pos;
    let columns = u32::try_from(delta).unwrap_or(u32::MAX);
    let popped = pop_template_marker(state);
    let advanced = popped.advance(delta, columns);
    let end = advanced.position();
    match piece {
        template::TemplateContinuation::Tail(content) => {
            step(advanced.emit(TokenKind::TemplateTail(content), Span::new(start, end)))
        }
        template::TemplateContinuation::Middle(content) => {
            let with_token =
                advanced.emit(TokenKind::TemplateMiddle(content), Span::new(start, end));
            let pushed = push_template_marker(with_token);
            step(pushed)
        }
    }
}

fn push_template_marker(state: LexState<'_>) -> LexState<'_> {
    let depth = state.brace_depth;
    let new_stack: Vec<usize> = state
        .template_stack
        .into_iter()
        .chain(std::iter::once(depth))
        .collect();
    LexState {
        template_stack: new_stack,
        ..state
    }
}

fn pop_template_marker(state: LexState<'_>) -> LexState<'_> {
    let new_stack: Vec<usize> = state
        .template_stack
        .iter()
        .take(state.template_stack.len().saturating_sub(1))
        .copied()
        .collect();
    LexState {
        template_stack: new_stack,
        ..state
    }
}

fn emit_number(state: LexState<'_>) -> Result<LexState<'_>, Error> {
    let start = state.position();
    let start_pos = state.pos;
    let (kind, end_pos) = number::scan(state.src, start_pos, start)?;
    let delta = end_pos - start_pos;
    let columns = u32::try_from(delta).unwrap_or(u32::MAX);
    let advanced = state.advance(delta, columns);
    let end = advanced.position();
    step(advanced.emit(kind, Span::new(start, end)))
}

fn emit_regex(state: LexState<'_>) -> Result<LexState<'_>, Error> {
    let start = state.position();
    let start_pos = state.pos;
    let (kind, end_pos) = regex::scan(state.src, start_pos, start)?;
    let delta = end_pos - start_pos;
    let columns = u32::try_from(delta).unwrap_or(u32::MAX);
    let advanced = state.advance(delta, columns);
    let end = advanced.position();
    step(advanced.emit(kind, Span::new(start, end)))
}

fn dispatch_dot(state: LexState<'_>) -> Result<LexState<'_>, Error> {
    match state.at(1) {
        Step::Byte(b'.') if matches!(state.at(2), Step::Byte(b'.')) => {
            emit_three_byte(state, TokenKind::Spread)
        }
        Step::Byte(b'0'..=b'9') => emit_number(state),
        Step::End | Step::Byte(_) => emit_single(state, TokenKind::Dot),
    }
}

fn emit_identifier(state: LexState<'_>) -> Result<LexState<'_>, Error> {
    let start = state.position();
    let start_pos = state.pos;
    let end = scan_identifier(state.src, start_pos);
    let bytes = state.src.get(start_pos..end).unwrap_or(&[]);
    let kind = keyword::classify(bytes).unwrap_or_else(|| {
        TokenKind::Identifier(std::str::from_utf8(bytes).unwrap_or("").to_owned())
    });
    let delta = end - start_pos;
    let columns = u32::try_from(delta).unwrap_or(u32::MAX);
    let advanced = state.advance(delta, columns);
    let span_end = advanced.position();
    step(advanced.emit(kind, Span::new(start, span_end)))
}

fn emit_private_identifier(state: LexState<'_>) -> Result<LexState<'_>, Error> {
    let start = state.position();
    let start_pos = state.pos;
    match state.at(1) {
        Step::Byte(b) if is_ident_start(b) => {
            let end = scan_identifier(state.src, start_pos + 1);
            let bytes = state.src.get(start_pos + 1..end).unwrap_or(&[]);
            let name = std::str::from_utf8(bytes).unwrap_or("").to_owned();
            let total_bytes = end - start_pos;
            let columns = u32::try_from(total_bytes).unwrap_or(u32::MAX);
            let advanced = state.advance(total_bytes, columns);
            let span_end = advanced.position();
            step(advanced.emit(
                TokenKind::PrivateIdentifier(name),
                Span::new(start, span_end),
            ))
        }
        Step::End | Step::Byte(_) => Err(Error::UnexpectedChar {
            at: state.position(),
            ch: '#',
        }),
    }
}

fn scan_identifier(src: &[u8], from: usize) -> usize {
    src.get(from)
        .copied()
        .filter(|b| is_ident_continue(*b))
        .map_or(from, |_| scan_identifier(src, from + 1))
}

fn is_ident_start(b: u8) -> bool {
    b.is_ascii_alphabetic() || b == b'_' || b == b'$'
}

fn is_ident_continue(b: u8) -> bool {
    b.is_ascii_alphanumeric() || b == b'_' || b == b'$'
}

fn dispatch_question(state: LexState<'_>) -> Result<LexState<'_>, Error> {
    match state.at(1) {
        Step::Byte(b'.') => emit_two_byte(state, TokenKind::OptionalChain),
        Step::Byte(b'?') if matches!(state.at(2), Step::Byte(b'=')) => {
            emit_three_byte(state, TokenKind::QQEq)
        }
        Step::Byte(b'?') => emit_two_byte(state, TokenKind::QQ),
        Step::End | Step::Byte(_) => emit_single(state, TokenKind::Question),
    }
}

fn dispatch_eq(state: LexState<'_>) -> Result<LexState<'_>, Error> {
    match state.at(1) {
        Step::Byte(b'=') if matches!(state.at(2), Step::Byte(b'=')) => {
            emit_three_byte(state, TokenKind::EqEqEq)
        }
        Step::Byte(b'=') => emit_two_byte(state, TokenKind::EqEq),
        Step::Byte(b'>') => emit_two_byte(state, TokenKind::Arrow),
        Step::End | Step::Byte(_) => emit_single(state, TokenKind::Eq),
    }
}

fn dispatch_bang(state: LexState<'_>) -> Result<LexState<'_>, Error> {
    match state.at(1) {
        Step::Byte(b'=') if matches!(state.at(2), Step::Byte(b'=')) => {
            emit_three_byte(state, TokenKind::BangEqEq)
        }
        Step::Byte(b'=') => emit_two_byte(state, TokenKind::BangEq),
        Step::End | Step::Byte(_) => emit_single(state, TokenKind::Bang),
    }
}

fn dispatch_lt(state: LexState<'_>) -> Result<LexState<'_>, Error> {
    match state.at(1) {
        Step::Byte(b'=') => emit_two_byte(state, TokenKind::LtEq),
        Step::Byte(b'<') if matches!(state.at(2), Step::Byte(b'=')) => {
            emit_three_byte(state, TokenKind::LtLtEq)
        }
        Step::Byte(b'<') => emit_two_byte(state, TokenKind::LtLt),
        Step::End | Step::Byte(_) => emit_single(state, TokenKind::Lt),
    }
}

fn dispatch_gt(state: LexState<'_>) -> Result<LexState<'_>, Error> {
    match state.at(1) {
        Step::Byte(b'=') => emit_two_byte(state, TokenKind::GtEq),
        Step::Byte(b'>')
            if matches!(state.at(2), Step::Byte(b'>'))
                && matches!(state.at(3), Step::Byte(b'=')) =>
        {
            emit_four_byte(state, TokenKind::GtGtGtEq)
        }
        Step::Byte(b'>') if matches!(state.at(2), Step::Byte(b'>')) => {
            emit_three_byte(state, TokenKind::GtGtGt)
        }
        Step::Byte(b'>') if matches!(state.at(2), Step::Byte(b'=')) => {
            emit_three_byte(state, TokenKind::GtGtEq)
        }
        Step::Byte(b'>') => emit_two_byte(state, TokenKind::GtGt),
        Step::End | Step::Byte(_) => emit_single(state, TokenKind::Gt),
    }
}

fn dispatch_plus(state: LexState<'_>) -> Result<LexState<'_>, Error> {
    match state.at(1) {
        Step::Byte(b'+') => emit_two_byte(state, TokenKind::PlusPlus),
        Step::Byte(b'=') => emit_two_byte(state, TokenKind::PlusEq),
        Step::End | Step::Byte(_) => emit_single(state, TokenKind::Plus),
    }
}

fn dispatch_minus(state: LexState<'_>) -> Result<LexState<'_>, Error> {
    match state.at(1) {
        Step::Byte(b'-') => emit_two_byte(state, TokenKind::MinusMinus),
        Step::Byte(b'=') => emit_two_byte(state, TokenKind::MinusEq),
        Step::End | Step::Byte(_) => emit_single(state, TokenKind::Minus),
    }
}

fn dispatch_star(state: LexState<'_>) -> Result<LexState<'_>, Error> {
    match state.at(1) {
        Step::Byte(b'*') if matches!(state.at(2), Step::Byte(b'=')) => {
            emit_three_byte(state, TokenKind::StarStarEq)
        }
        Step::Byte(b'*') => emit_two_byte(state, TokenKind::StarStar),
        Step::Byte(b'=') => emit_two_byte(state, TokenKind::StarEq),
        Step::End | Step::Byte(_) => emit_single(state, TokenKind::Star),
    }
}

fn dispatch_percent(state: LexState<'_>) -> Result<LexState<'_>, Error> {
    match state.at(1) {
        Step::Byte(b'=') => emit_two_byte(state, TokenKind::PercentEq),
        Step::End | Step::Byte(_) => emit_single(state, TokenKind::Percent),
    }
}

fn dispatch_amp(state: LexState<'_>) -> Result<LexState<'_>, Error> {
    match state.at(1) {
        Step::Byte(b'&') if matches!(state.at(2), Step::Byte(b'=')) => {
            emit_three_byte(state, TokenKind::AmpAmpEq)
        }
        Step::Byte(b'&') => emit_two_byte(state, TokenKind::AmpAmp),
        Step::Byte(b'=') => emit_two_byte(state, TokenKind::AmpEq),
        Step::End | Step::Byte(_) => emit_single(state, TokenKind::Amp),
    }
}

fn dispatch_pipe(state: LexState<'_>) -> Result<LexState<'_>, Error> {
    match state.at(1) {
        Step::Byte(b'|') if matches!(state.at(2), Step::Byte(b'=')) => {
            emit_three_byte(state, TokenKind::PipePipeEq)
        }
        Step::Byte(b'|') => emit_two_byte(state, TokenKind::PipePipe),
        Step::Byte(b'=') => emit_two_byte(state, TokenKind::PipeEq),
        Step::End | Step::Byte(_) => emit_single(state, TokenKind::Pipe),
    }
}

fn dispatch_caret(state: LexState<'_>) -> Result<LexState<'_>, Error> {
    match state.at(1) {
        Step::Byte(b'=') => emit_two_byte(state, TokenKind::CaretEq),
        Step::End | Step::Byte(_) => emit_single(state, TokenKind::Caret),
    }
}
