//! # Query helpers
//!
//! Read-only convenience APIs for inspecting a [`TokenizedFile`] without
//! exposing its internal representation.

use crate::lexer::{Line, TokenData, TokenPosition, TokenizedFile, VisibleLineSpan};

impl<'src> TokenizedFile<'src> {
    pub const fn eof(&self) -> TokenPosition {
        TokenPosition(self.buffer.len())
    }

    pub fn is_eof(&self, position: &TokenPosition) -> bool {
        position == &self.eof()
    }

    /// Returns the number of physical lines stored in this file.
    ///
    /// Empty line after the trailing newline sequence isn't counted as a line
    /// by this method.
    #[must_use]
    pub const fn line_count(&self) -> usize {
        self.lines.len()
    }

    /// Returns an iterator over tokens that *start* on physical line
    /// `line_number`.
    ///
    /// The yielded items are `(TokenPosition, TokenData)` pairs, matching the
    /// canonical file-wide token arena.
    ///
    /// If the line ends with a newline token, that newline token is included.
    ///
    /// If the line begins with a carried fragment of a multi-line token that
    /// started on an earlier line, that fragment is **not** yielded here.
    /// Use [`TokenizedFile::line_text`] to reconstruct the visible content of
    /// the full line.
    ///
    /// If `line_number` is out of bounds, the returned iterator is empty.
    #[must_use]
    pub fn line_tokens(
        &self,
        line_number: usize,
    ) -> std::vec::IntoIter<(TokenPosition, TokenData<'src>)> {
        let Some(line) = self.lines.get(line_number) else {
            return Vec::new().into_iter();
        };
        let Some(local_range) = line.local_range() else {
            return Vec::new().into_iter();
        };
        let mut out = Vec::with_capacity(local_range.len());
        for buffer_index in local_range {
            // Invariant:
            // `Line::local_range()` is always constructed from contiguous
            // slices of `self.buffer` during tokenization, so every index in
            // this range must be valid for `self.buffer`.
            let token_data = self.buffer[buffer_index];
            out.push((TokenPosition(buffer_index), token_data));
        }
        out.into_iter()
    }

    /// Returns the token stored at `position`, if that position is valid.
    ///
    /// This is a direct lookup into the file-wide token buffer.
    #[must_use]
    pub fn token_at(&self, position: TokenPosition) -> Option<TokenData<'src>> {
        self.buffer.get(position.0).copied()
    }

    /// Reconstructs the visible text of physical line `line_index`.
    ///
    /// The returned string does **not** include a trailing line terminator.
    ///
    /// Unlike [`TokenizedFile::line_tokens`], this method includes the visible
    /// fragment of a multi-line token carried from an earlier line.
    ///
    /// Returns [`None`] iff `line_index >= self.line_count()`.
    #[must_use]
    pub fn line_text(&self, line_index: usize) -> Option<String> {
        let line = self.lines.get(line_index)?;
        let mut out = String::new();

        if let Some(piece) = self.carried_piece_for_line(line_index) {
            out.push_str(piece);
        }
        let Some(range) = line.local_range() else {
            return Some(out);
        };
        for buffer_index in range.clone() {
            let token_piece = self.buffer[buffer_index];
            if token_piece.token.is_newline() {
                // Must be last token
                debug_assert_eq!(buffer_index + 1, range.end);
                break;
            }
            if token_piece.token.can_span_lines()
                && let Some(first_segment) = self
                    .multi_line_map
                    .get(&buffer_index)
                    .and_then(|segments| segments.first())
            {
                out.push_str(&token_piece.lexeme[first_segment.clone()]);
                // Must be last token
                debug_assert_eq!(buffer_index + 1, range.end);
                break;
            }
            out.push_str(token_piece.lexeme);
        }
        Some(out)
    }

    /// Returns the 0-based physical line on which the token at `position`
    /// starts.
    ///
    /// For multi-line tokens, this is the line where the token begins, not
    /// every physical line it spans.
    ///
    /// Returns `None` if `position` is out of bounds.
    #[must_use]
    pub fn token_line(&self, position: TokenPosition) -> Option<usize> {
        // EOF is a valid virtual position: past the end of the last stored line.
        if position == self.eof() {
            return self.line_count().checked_sub(1);
        }

        // Reject invalid non-EOF positions early.
        self.buffer.get(position.0)?;

        let line_index = self
            .lines
            .partition_point(|line| self.line_search_upper_bound(line) <= position.0);

        (line_index < self.lines.len()).then_some(line_index)
    }

    /// Returns the exclusive upper token index bound for binary-searching
    /// lines by token position.
    ///
    /// In other words: every token that "belongs" to this line in start-line
    /// terms has index `< returned_value`.
    fn line_search_upper_bound(&self, line: &Line) -> usize {
        if let Some(local_range) = line.local_range() {
            local_range.end
        } else {
            // Pure continuation line: it contains only the carried fragment of
            // a multi-line token that started earlier.
            //
            // That token is always the last local token on the origin line, so
            // its token index + 1 acts as the exclusive upper bound.
            let origin_line = line
                .continued_from
                .expect("empty line entry must be a continuation line");
            self.carried_token_index(origin_line)
                .expect("continuation line must point to a valid origin token")
                + 1
        }
    }

    /// If `line_index` begins with a fragment of a multi-line token that
    /// started earlier, returns the visible slice of that token for this line.
    fn carried_piece_for_line(&self, line_index: usize) -> Option<&'src str> {
        // Find carried, multiline token
        let origin_line = self.lines.get(line_index)?.continued_from?;
        let carried_token_index = self.carried_token_index(origin_line)?;
        // Find right part of the multiline token's lexeme
        let segments = self.multi_line_map.get(&carried_token_index)?;
        let segment_index = line_index.checked_sub(origin_line)?;
        let boundary = segments.get(segment_index)?;
        self.buffer
            .get(carried_token_index)?
            .lexeme
            .get(boundary.clone())
    }

    /// Recovers the token index of the multi-line token that started on
    /// `origin_line` and is carried into later lines.
    ///
    /// In the current representation, this is always the last local token that
    /// started on the origin line.
    fn carried_token_index(&self, origin_line: usize) -> Option<usize> {
        let range = self.lines.get(origin_line)?.local_range()?;
        let token_index = range.end.checked_sub(1)?;

        debug_assert!(self.buffer[token_index].token.can_span_lines());
        Some(token_index)
    }

    /// Returns the visible per-line spans occupied by the token at `position`.
    ///
    /// Coordinates are expressed in visible character columns inside
    /// `line_text(line)`, with an exclusive end bound.
    ///
    /// Newline-only tokens have no visible text, so they return an empty vector.
    ///
    /// Returns `None` if `position` is invalid.
    #[must_use]
    pub fn token_visible_spans(&self, position: TokenPosition) -> Option<Vec<VisibleLineSpan>> {
        // EOF is a virtual zero-width span at the end of the last stored line.
        if position == self.eof() {
            let line = self.line_count().checked_sub(1)?;
            let column = self.line_text(line)?.chars().count();

            return Some(vec![VisibleLineSpan {
                line,
                columns: column..column,
            }]);
        }

        let token_piece = self.buffer.get(position.0).copied()?;
        let start_line = self.token_line(position)?;
        let start_column = self.token_start_visible_column(position)?;

        if token_piece.token.is_newline() {
            return Some(Vec::new());
        }

        if let Some(segments) = self.multi_line_map.get(&position.0) {
            let mut out = Vec::with_capacity(segments.len());

            for (segment_index, byte_range) in segments.iter().enumerate() {
                let visible_text = &token_piece.lexeme[byte_range.clone()];
                let width = visible_text.chars().count();

                if width == 0 {
                    continue;
                }

                let line = start_line + segment_index;

                if line >= self.line_count() {
                    break;
                }

                let column_start = if segment_index == 0 { start_column } else { 0 };
                out.push(VisibleLineSpan {
                    line,
                    columns: column_start..(column_start + width),
                });
            }

            return Some(out);
        }

        let width = token_piece.lexeme.chars().count();
        Some(vec![VisibleLineSpan {
            line: start_line,
            columns: start_column..(start_column + width),
        }])
    }

    /// Returns the visible start column of the token at `position` inside
    /// `line_text(token_line(position))`.
    ///
    /// Column is measured in visible characters, excluding line terminators.
    fn token_start_visible_column(&self, position: TokenPosition) -> Option<usize> {
        let line_index = self.token_line(position)?;
        let line = self.lines.get(line_index)?;

        let mut column = self
            .carried_piece_for_line(line_index)
            .map_or(0, |text| text.chars().count());

        let local_range = line.local_range()?;
        for buffer_index in local_range {
            if buffer_index == position.0 {
                return Some(column);
            }

            let token_piece = self.buffer.get(buffer_index)?;

            if token_piece.token.is_newline() {
                break;
            }

            if token_piece.token.can_span_lines() && self.multi_line_map.contains_key(&buffer_index)
            {
                //debug_assert_eq!(buffer_index + 1, local_range.end);
                return None;
            }

            column += token_piece.lexeme.chars().count();
        }

        None
    }

    #[must_use]
    pub fn span_visible_on_line(&self, span: crate::lexer::TokenSpan) -> Option<VisibleLineSpan> {
        let start = self.token_visible_spans(span.start)?.into_iter().next()?;
        let end = self.token_visible_spans(span.end)?.into_iter().last()?;

        if start.line != end.line {
            return None;
        }

        Some(VisibleLineSpan {
            line: start.line,
            columns: start.columns.start..end.columns.end,
        })
    }

    pub fn token_text(&self, pos: TokenPosition) -> Option<&'src str> {
        self.token_at(pos).map(|t| t.lexeme)
    }

    pub fn same_line(&self, a: TokenPosition, b: TokenPosition) -> bool {
        match (self.token_line(a), self.token_line(b)) {
            (Some(x), Some(y)) => x == y,
            _ => false,
        }
    }
}
