//! # Tokenizer
//!
//! Converts raw source text into a lossless, position-aware stream of lexical
//! [`Token`]s, grouped *per physical line*, and returns it as
//! a [`TokenizedFile`]. A trailing newline terminates the last physical line
//! rather than introducing an additional empty line.
//!
//! Design goals:
//!
//! 1. **Lossless**: preserving complete information for each token, enough to
//!     recreate the original bytes without loss.
//! 2. **LSP readiness**: the LSP heavily relies on UTF-16 support, so we
//!     precompute lengths of each token in that encoding, making interfacing
//!     easier.
//!
//! ## Opt-in debug helpers
//!
//! Extra diagnostics become available in **debug builds** or when the crate is
//! compiled with `debug` feature enabled. They live in the [`debug_tools`]
//! extension trait, implemented for [`TokenizedFile`].

mod queries;
mod raw_lexer;
#[cfg(test)]
mod tests;
mod token;

use std::collections::HashMap;
use std::ops::Range;

use logos::Logos;

use raw_lexer::RawToken;

pub use raw_lexer::BraceKind;
pub use token::Keyword;
pub use token::Token;

/// Empirically chosen starting size for token buffer (used during tokenization)
/// that provides good performance.
const DEFAULT_TOKEN_BUFFER_CAPACITY: usize = 20_000;

// TODO: check this!!!
/// Visible fragment of a token on one physical line.
///
/// `columns` is an end-exclusive range inside the string returned by
/// [`TokenizedFile::line_text`] for that line.
#[derive(Clone, Debug, Hash, PartialEq, Eq)]
pub struct VisibleLineSpan {
    pub line: usize,
    pub columns: std::ops::Range<usize>,
}

/// A token together with its source text and precomputed UTF-16 length.
///
/// It does not store an absolute file position.
#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
pub struct TokenData<'src> {
    /// Kind of token that was lexed.
    pub token: Token,
    /// Underlying text that was lexed as the corresponding token.
    pub lexeme: &'src str,
    /// Length of the token in UTF-16 code units for the needs of easy seeking
    /// using given LSP cursor coordinates (line + UTF-16 offset).
    /// Precomputed for convenience.
    pub utf16_length: usize,
}

/// 0-based index of a token within the file-wide token buffer.
#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, PartialOrd, Ord, Default)]
pub struct TokenPosition(pub usize);

// All inclusive!
#[derive(Debug, Hash, PartialEq, Eq, PartialOrd, Ord, Clone, Copy)]
pub struct TokenSpan {
    pub start: TokenPosition,
    pub end: TokenPosition,
}

impl TokenSpan {
    // -------- existing coord-based API (unchanged externally) --------

    #[must_use]
    pub const fn merge(left_span: &Self, right_span: &Self) -> Self {
        Self {
            // assumes both were constructed in the same style; good enough for the refactor
            start: left_span.start,
            end: right_span.end,
        }
    }

    // -------- NEW: 4 constructors based on TokenIndex --------

    /// Single-token span from an index (coords are dummy for now).
    #[inline]
    #[must_use]
    pub const fn new(single_index: TokenPosition) -> Self {
        Self {
            start: single_index,
            end: single_index,
        }
    }

    /// Span from two indices (coords are dummy for now).
    #[inline]
    #[must_use]
    pub const fn range(from: TokenPosition, to: TokenPosition) -> Self {
        Self {
            start: from,
            end: to,
        }
    }

    /// Immutable extension by index (keeps coords as-is).
    #[inline]
    #[must_use]
    pub fn extended(&self, right_most_index: TokenPosition) -> Self {
        Self {
            start: self.start,
            end: std::cmp::max(self.end, right_most_index),
        }
    }

    /// In-place extension by index (coords unchanged).
    #[inline]
    pub fn extend_to(&mut self, right_most_index: TokenPosition) {
        if right_most_index > self.end {
            self.end = right_most_index;
        }
    }
}

/// A tokenized, lossless representation of an `UnrealScript` source file.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct TokenizedFile<'src> {
    /// Arena of every token span in this file.
    buffer: Vec<TokenData<'src>>,
    /// Mapping from physical line number to the tokens that belong to it.
    lines: Vec<Line>,
    /// Mapping token index to ranges of bytes that correspond to
    /// visible characters (i.e. all non line terminators) in its lines.
    ///
    /// Records only exists for multiline tokens and ranges can be empty for
    /// lines that only contain line break boundary.
    multi_line_map: HashMap<BufferIndex, Vec<VisibleByteRange>>,
    /// Simple flag for marking erroneous state.
    had_errors: bool,
}

/// An immutable iterator over all tokens in a [`TokenizedFile`], preserving
/// their order of appearance in the original source file.
///
/// After exhaustion it keeps returning [`None`].
#[must_use]
#[derive(Clone, Debug)]
pub struct Tokens<'file, 'src> {
    /// Position of the next token to be returned in the canonical file-wide
    /// token arena.
    cursor: TokenPosition,
    /// [`TokenizedFile`] whose tokens we're iterating over.
    source_file: &'file TokenizedFile<'src>,
}

/// Type for referring to line numbers.
type LineNumber = usize;

/// Type for specific tokens inside each [`Line`].
type BufferIndex = usize;

/// Type for describing sub-range of visible characters of a single line for
/// some token.
type VisibleByteRange = Range<usize>;

/// Representation of a single physical line of the source file.
///
/// Uses ranges instead of slices to avoid a self-referential relationship
/// with [`TokenizedFile`], which Rust forbids.
#[derive(Clone, Debug, Hash, PartialEq, Eq)]
struct Line {
    /// Token that began on an earlier line (`None` for standalone lines).
    continued_from: Option<LineNumber>,
    /// Contiguous tokens that started on this line (`start >= end` iff empty).
    local_range: Range<BufferIndex>,
}

/// Mutable state used while tokenizing a source file.
#[derive(Debug)]
struct Tokenizer<'src> {
    /// Arena that owns every [`TokenData`] produced for the file.
    buffer: Vec<TokenData<'src>>,
    /// Mapping from physical line number to the tokens that belong to it.
    lines: Vec<Line>,
    /// Mapping token index to ranges of bytes that correspond to
    /// visible characters in its lines.
    multi_line_map: HashMap<BufferIndex, Vec<VisibleByteRange>>,
    /// The 0-based physical line number that is currently being scanned.
    line_number: usize,
    /// Points to the first token (index in [`Tokenizer::buffer`]) not yet
    /// committed to `lines`, e.g. where the current *line* starts.
    uncommitted_start_index: usize,
    /// When a multi-line token is being scanned, stores the 0-based line
    /// on which it started; [`None`] otherwise.
    ///
    /// `Some(line_number)` iff the current line is within a multi-line token
    /// that started on `line_number`; it is consumed exactly once by
    /// [`Self::commit_current_line`].
    multi_line_start_line: Option<LineNumber>,
    /// Set to `true` if the lexer reported any error tokens.
    had_errors: bool,
}

impl<'src> TokenizedFile<'src> {
    /// Tokenizes `source` and returns a fresh [`TokenizedFile`].
    ///
    /// Its output is lossless and groups resulting tokens by physical lines.
    /// Error spans are preserved as [`Token::Error`].
    #[must_use]
    pub fn tokenize(source: &'src str) -> Self {
        let mut tokenizer = Tokenizer::new();
        let mut lexer = RawToken::lexer(source);

        while let Some(token_result) = lexer.next() {
            // Add `Token::Error` manually, since Logos won't do it for us.
            let token = token_result.unwrap_or_else(|()| {
                tokenizer.had_errors = true;
                RawToken::Error
            });
            let token_piece = make_token_data(Token::from(token), lexer.slice());
            tokenizer.process_token_piece(token_piece);
        }
        tokenizer.into_tokenized_file()
    }

    /// Returns `true` if tokenization produced any error tokens.
    #[must_use]
    pub const fn has_errors(&self) -> bool {
        self.had_errors
    }

    /// Returns an iterator over all contained tokens in the order they appear
    /// in the original source file.
    ///
    /// Returns pairs of position and token data: `(TokenPosition, TokenData)`.
    pub const fn iter(&self) -> Tokens<'_, 'src> {
        Tokens::new(self)
    }
}

impl Line {
    /// Creates a standalone line that owns a contiguous slice in
    /// the [`TokenizedFile::buffer`] arena.
    const fn standalone(locals: Range<BufferIndex>) -> Self {
        Self {
            continued_from: None,
            local_range: locals,
        }
    }

    /// Creates a line that is part of a multi-line token started on
    /// another line, referencing the 0-based index of its origin.
    const fn continued(carried: LineNumber) -> Self {
        Self {
            continued_from: Some(carried),
            local_range: 0..0,
        }
    }

    /// Creates a line that is part of a multi-line token started on
    /// another line and also contains additional tokens local to itself.
    const fn continued_with_tokens(carried: LineNumber, locals: Range<BufferIndex>) -> Self {
        Self {
            continued_from: Some(carried),
            local_range: locals,
        }
    }

    /// Returns a range of tokens inside [`TokenizedFile::buffer`] that start
    /// on this line.
    ///
    /// [`None`] means there are no such tokens. Otherwise range is guaranteed
    /// to not be empty.
    fn local_range(&self) -> Option<Range<BufferIndex>> {
        if self.local_range.is_empty() {
            None
        } else {
            Some(self.local_range.clone())
        }
    }
}

impl<'src> Tokenizer<'src> {
    /// Returns an empty tokenizer state.
    fn new() -> Self {
        Self {
            buffer: Vec::with_capacity(DEFAULT_TOKEN_BUFFER_CAPACITY),
            lines: Vec::new(),
            multi_line_map: HashMap::new(),
            line_number: 0,
            uncommitted_start_index: 0,
            multi_line_start_line: None,
            had_errors: false,
        }
    }

    /// Handles a token span and dispatches to the appropriate handler.
    fn process_token_piece(&mut self, token_piece: TokenData<'src>) {
        if token_piece.token.can_span_lines() {
            self.process_multi_line_token(token_piece);
        } else {
            self.process_single_line_token(token_piece);
        }
    }

    /// Handles simple tokens that *never* span multiple lines, allowing us to
    /// skip a lot of work.
    fn process_single_line_token(&mut self, token_piece: TokenData<'src>) {
        if token_piece.token.is_newline() {
            self.line_number += 1;
            self.buffer.push(token_piece);
            self.commit_current_line();
        } else {
            self.buffer.push(token_piece);
        }
    }

    /// Handles tokens that might contain one or more newline characters.
    fn process_multi_line_token(&mut self, token_piece: TokenData<'src>) {
        let start_line = self.line_number;
        let line_break_map = split_visible_line_segments(token_piece.lexeme);
        let newline_count = line_break_map.len().saturating_sub(1);

        // Needed for unterminated multi-line error tokens that reach EOF right
        // after a line break.
        let ends_with_newline =
            token_piece.lexeme.ends_with('\n') || token_piece.lexeme.ends_with('\r');

        let multi_line_token_index = self.buffer.len();
        self.buffer.push(token_piece);
        if !line_break_map.is_empty() {
            self.multi_line_map
                .insert(multi_line_token_index, line_break_map);
        }
        // A line is committed only once the token stream has actually crossed
        // a physical line boundary.
        if newline_count > 0 {
            // This clears `multi_line_start_line`
            self.commit_current_line();
            // We only need to insert one `Line::continued(start_line)` per
            // *interior* line:
            //
            // standalone           | local int i = /* Now we start long comment
            // continued            | with three line breaks and *exactly* two
            // continued            | inner lines that contain nothing but
            // continued_with_tokens| comment bytes! */ = 0;
            let inner_lines_count = newline_count - 1;
            for _ in 0..inner_lines_count {
                self.lines.push(Line::continued(start_line));
            }
            self.multi_line_start_line = if ends_with_newline {
                None
            } else {
                Some(start_line)
            };
        }

        self.line_number = start_line + newline_count;
    }

    /// Commits the tokens of the current physical line into `self.lines`.
    fn commit_current_line(&mut self) {
        let slice_end = self.buffer.len();
        // A trailing newline terminates the current physical line rather than
        // creating an additional empty line entry.
        if slice_end > self.uncommitted_start_index {
            let slice = self.uncommitted_start_index..slice_end;

            // If we were in the middle of a multi-line token, we
            // *always* consume `multi_line_start` here, ensuring that each call
            // to `commit_current_line()` only applies it once.
            // This guarantees no "bleed" between adjacent multi-line tokens.
            if let Some(from) = self.multi_line_start_line.take() {
                self.lines.push(Line::continued_with_tokens(from, slice));
            } else {
                self.lines.push(Line::standalone(slice));
            }
            self.uncommitted_start_index = slice_end;
        }
    }

    /// Finishes tokenization, converting accumulated data into
    /// [`TokenizedFile`].
    fn into_tokenized_file(mut self) -> TokenizedFile<'src> {
        // Commits the final line when the file does not end with a newline.
        self.commit_current_line();
        // If we still have a `multi_line_start`
        // (i.e. a pure multi-line token with no local tokens on its last line),
        // push a bare `Line::continued` entry.
        if let Some(from) = self.multi_line_start_line.take() {
            self.lines.push(Line::continued(from));
        }

        self.buffer.shrink_to_fit();
        self.lines.shrink_to_fit();

        TokenizedFile {
            buffer: self.buffer,
            lines: self.lines,
            multi_line_map: self.multi_line_map,
            had_errors: self.had_errors,
        }
    }
}

fn make_token_data(token: Token, text: &str) -> TokenData<'_> {
    let length_utf16 = text.encode_utf16().count();
    TokenData {
        lexeme: text,
        token,
        utf16_length: length_utf16,
    }
}

/// Returns byte ranges of visible text characters for each physical line
/// spanned by `text`.
///
/// Returns an empty vector if `text` contains no line breaks.
fn split_visible_line_segments(text: &str) -> Vec<Range<usize>> {
    let bytes = text.as_bytes();
    let mut segments = Vec::new();

    let mut segment_start = 0usize;
    let mut saw_line_break = false;
    let mut bytes_iterator = bytes.iter().enumerate().peekable();
    while let Some((next_byte_index, &next_byte)) = bytes_iterator.next() {
        // Logos' regex rule is "\r\n|\n|\r", so we agree with it on new line
        // character treatment
        match next_byte {
            b'\r' => {
                saw_line_break = true;
                let visible_end = next_byte_index;
                let next_start =
                    if let Some((next_line_break_index, b'\n')) = bytes_iterator.peek().copied() {
                        bytes_iterator.next(); // consume '\n' of `\r\n`
                        next_line_break_index + 1
                    } else {
                        next_byte_index + 1
                    };
                segments.push(segment_start..visible_end);
                segment_start = next_start;
            }
            b'\n' => {
                saw_line_break = true;
                let visible_end = next_byte_index;
                segments.push(segment_start..visible_end);
                segment_start = next_byte_index + 1;
            }
            _ => (),
        }
    }
    // If the token contained at least one line break, include the visible
    // segment of its final physical line as well. This may be empty, e.g.
    // for text ending with '\n' or '\r\n'.
    if saw_line_break {
        segments.push(segment_start..bytes.len());
    }
    segments
}

// Because once `cursor` moves past the end of `buffer`, it can never become
// valid again.
impl std::iter::FusedIterator for Tokens<'_, '_> {}

impl<'file, 'src> Tokens<'file, 'src> {
    /// Advances the iterator cursor by one token.
    const fn advance(&mut self) {
        self.cursor.0 += 1;
    }

    /// Creates a new iterator.
    const fn new(source_file: &'file TokenizedFile<'src>) -> Self {
        Self {
            source_file,
            cursor: TokenPosition(0),
        }
    }
}

impl<'src> Iterator for Tokens<'_, 'src> {
    type Item = (TokenPosition, TokenData<'src>);

    fn next(&mut self) -> Option<Self::Item> {
        if let Some(&token_piece) = self.source_file.buffer.get(self.cursor.0) {
            let position = self.cursor;
            self.advance();
            Some((position, token_piece))
        } else {
            None
        }
    }
}

impl<'file, 'src> IntoIterator for &'file TokenizedFile<'src> {
    type Item = (TokenPosition, TokenData<'src>);
    type IntoIter = Tokens<'file, 'src>;

    fn into_iter(self) -> Self::IntoIter {
        self.iter()
    }
}
