lang-check 0.6.0

use tree_sitter::Node;

use super::{ProseRange, gap, shared};

/// Built-in environment types that tree-sitter-latex recognises as dedicated
/// node kinds (not `generic_environment`). Skip these entirely.
const SKIP_ENV_KINDS: &[&str] = &[
    "verbatim_environment",
    "minted_environment",
    "listing_environment",
    "comment_environment",
    "math_environment",
    "asy_environment",
    "luacode_environment",
    "pycode_environment",
    "sageblock_environment",
    "sagesilent_environment",
];

/// Generic environment names (the `{name}` in `\begin{name}`) that should be
/// skipped. These are environments tree-sitter-latex parses as
/// `generic_environment` rather than a dedicated node kind.
const SKIP_GENERIC_ENVS: &[&str] = &[
    "algorithm",
    "algorithmic",
    "lstlisting",
    "equation",
    "equation*",
    "align",
    "align*",
    "gather",
    "gather*",
    "multline",
    "multline*",
    "flalign",
    "flalign*",
    "split",
    "mathpar",
    "mathpar*",
    "IEEEeqnarray",
    "IEEEeqnarray*",
    "tikzpicture",
    "pgfpicture",
    "forest",
    "tabular",
    "tabular*",
    "array",
    "matrix",
    "bmatrix",
    "pmatrix",
    "vmatrix",
    "Bmatrix",
    "Vmatrix",
    "cases",
    "bnf",
];

/// Node types that represent math and should be skipped.
const SKIP_NODES: &[&str] = &["inline_formula", "displayed_equation"];

/// Node types that are always structural (never contain prose `word` nodes).
const STRUCTURAL_NODES: &[&str] = &[
    "command_name",
    "graphics_include",
    "label_definition",
    "label_reference",
    "citation",
    "package_include",
    "bibstyle_include",
];

/// Generic command names (the `\name` in `\name{...}`) whose arguments are
/// non-prose and should be skipped entirely. When `collect_words` encounters
/// a `generic_command` whose command name matches, the entire subtree is
/// dropped.
const SKIP_GENERIC_COMMANDS: &[&str] = &[
    "thispagestyle",
    "pagestyle",
    "bibliographystyle",
    "bibliography",
    "setcounter",
    "addtocounter",
    "setlength",
    "addtolength",
    "newcommand",
    "renewcommand",
    "newenvironment",
    "renewenvironment",
    "DeclareMathOperator",
    "definecolor",
    "hypersetup",
    "geometry",
    "input",
    "include",
    "hfill",
    "vfill",
    "hspace",
    "vspace",
    "smallskip",
    "medskip",
    "bigskip",
    "hrule",
    "vrule",
    "newpage",
    "clearpage",
    "maketitle",
    "tableofcontents",
    "listoffigures",
    "listoftables",
    "texttt",
    "verb",
    "lstinline",
    "mintinline",
    "url",
    "href",
    "path",
];

/// User-configurable extras for LaTeX prose extraction.
#[derive(Default)]
pub struct LatexExtras<'a> {
    pub skip_envs: &'a [String],
    pub skip_commands: &'a [String],
}

/// Extract prose ranges from a LaTeX AST.
///
/// Walks the tree collecting `word` leaf nodes, skipping preamble, math,
/// verbatim, and other non-prose environments. Adjacent words are merged
/// into sentence-level prose chunks.
pub(crate) fn extract(text: &str, root: Node, extras: &LatexExtras) -> Vec<ProseRange> {
    let doc_start = find_document_body_start(root, text);

    let mut word_ranges: Vec<(usize, usize)> = Vec::new();
    let mut skip_until = 0usize;
    collect_words(
        root,
        text,
        doc_start,
        false,
        extras,
        &mut skip_until,
        &mut word_ranges,
    );

    shared::merge_ranges(&word_ranges, text, latex_gap)
}

/// Check whether a node kind represents a structural (non-prose) container.
///
/// Includes all `brack_group*` variants, specialised `curly_group_*` variants
/// (but NOT plain `curly_group` which holds prose), and explicitly listed
/// structural node types.
fn is_structural_node(kind: &str) -> bool {
    if kind.starts_with("brack_group") {
        return true;
    }
    if kind.starts_with("curly_group_") {
        return true;
    }
    STRUCTURAL_NODES.contains(&kind)
}

/// Find the byte offset where `\begin{document}` body starts.
/// Returns 0 if no document environment is found (single-file snippets).
fn find_document_body_start(root: Node, text: &str) -> usize {
    let mut cursor = root.walk();
    for child in root.children(&mut cursor) {
        if child.kind() == "generic_environment"
            && let Some(begin_node) = child.child_by_field_name("begin")
        {
            let begin_text = &text[begin_node.start_byte()..begin_node.end_byte()];
            if begin_text.contains("document") {
                return begin_node.end_byte();
            }
        }
    }
    0
}

/// Recursively collect prose leaf nodes (`word`), skipping excluded subtrees.
///
/// The `in_structural` flag propagates through the tree so that `word` nodes
/// nested inside structural parents (at any depth) are skipped.
/// Whether `name` is a command whose argument is delimited by a character
/// chosen at the call site (`\verb|x|`, `\lstinline!x!`) instead of by braces.
///
/// The grammar does not model that argument, so it arrives as an ordinary
/// `word` node next to the command and has to be cut out by byte offset.
const fn is_verbatim_delimited(name: &[u8]) -> bool {
    matches!(name, b"verb" | b"lstinline" | b"mintinline")
}

/// End of the delimiter-delimited argument starting at `i`.
///
/// `i` is the byte just past the command name *and past a starred variant's
/// star* — every caller consumes the star before asking, because a non-verbatim
/// command has to consume it either way.
///
/// The first character is the delimiter and the argument runs to its next
/// occurrence; `\mintinline{rust}|x|` takes a brace group before it. Returns
/// the byte just past the closing delimiter, or `None` when what follows is not
/// a delimited argument after all.
fn verbatim_argument_end(bytes: &[u8], mut i: usize) -> Option<usize> {
    if bytes.get(i) == Some(&b'{') {
        i = shared::skip_balanced_bytes(bytes, i + 1, b'{', b'}', Some(b'\\'));
    }
    let delimiter = *bytes.get(i)?;
    // A letter or a space is the next word, not a delimiter.
    if delimiter.is_ascii_alphanumeric() || delimiter.is_ascii_whitespace() {
        return None;
    }
    let close = bytes[i + 1..].iter().position(|&b| b == delimiter)?;
    Some(i + 1 + close + 1)
}

/// Whether `node` is a `\verb`-style command, and if so where its delimited
/// argument ends.
fn verbatim_command_end(node: Node, text: &str) -> Option<usize> {
    let name_node = shared::child_of_kind(node, "command_name")?;
    let raw = &text[name_node.byte_range()];
    let name = raw.strip_prefix('\\').unwrap_or(raw).trim_end_matches('*');
    if !is_verbatim_delimited(name.as_bytes()) {
        return None;
    }
    verbatim_argument_end(text.as_bytes(), name_node.end_byte())
}

fn collect_words(
    node: Node,
    text: &str,
    doc_start: usize,
    in_structural: bool,
    extras: &LatexExtras,
    skip_until: &mut usize,
    out: &mut Vec<(usize, usize)>,
) {
    if node.end_byte() <= doc_start {
        return;
    }

    let kind = node.kind();

    if SKIP_ENV_KINDS.contains(&kind) || SKIP_NODES.contains(&kind) {
        return;
    }

    if kind == "generic_environment" && should_skip_generic_env(node, text, extras.skip_envs) {
        return;
    }

    if kind == "generic_command" && should_skip_generic_command(node, text, extras.skip_commands) {
        // `\verb|code|` puts its argument outside the command node, so the words
        // it covers have to be dropped by offset as the walk reaches them.
        if let Some(end) = verbatim_command_end(node, text) {
            *skip_until = (*skip_until).max(end);
        }
        return;
    }

    let structural = in_structural || is_structural_node(kind);

    if kind == "word" {
        if !structural {
            let start = node.start_byte();
            let end = node.end_byte();
            if start >= doc_start && start < end && start >= *skip_until {
                out.push((start, end));
            }
        }
        return;
    }

    let mut cursor = node.walk();
    for child in node.children(&mut cursor) {
        collect_words(child, text, doc_start, structural, extras, skip_until, out);
    }
}

/// Check if a `generic_environment` node should be skipped based on its name.
fn should_skip_generic_env(node: Node, text: &str, extra_skip_envs: &[String]) -> bool {
    let mut cursor = node.walk();
    for child in node.children(&mut cursor) {
        if child.kind() != "begin" {
            continue;
        }
        let mut inner = child.walk();
        for bc in child.children(&mut inner) {
            if bc.kind() != "curly_group_text" {
                continue;
            }
            let mut name_cursor = bc.walk();
            for name_child in bc.children(&mut name_cursor) {
                if name_child.kind() != "text" {
                    continue;
                }
                let env_name = &text[name_child.start_byte()..name_child.end_byte()];
                let env_name = env_name.trim();
                if SKIP_GENERIC_ENVS.contains(&env_name) {
                    return true;
                }
                return extra_skip_envs.iter().any(|e| e == env_name);
            }
        }
        break;
    }
    false
}

/// Check if a `generic_command` node should be skipped based on its command name.
fn should_skip_generic_command(node: Node, text: &str, extra_skip_commands: &[String]) -> bool {
    let Some(name_node) = shared::child_of_kind(node, "command_name") else {
        return false;
    };
    // command_name text is e.g. `\thispagestyle` — strip leading `\`. A starred
    // variant (`\verb*`) is the same command for this purpose.
    let raw = &text[name_node.byte_range()];
    let name = raw.strip_prefix('\\').unwrap_or(raw).trim_end_matches('*');
    SKIP_GENERIC_COMMANDS.contains(&name) || extra_skip_commands.iter().any(|c| c == name)
}

// ---------------------------------------------------------------------------
// Word-range merging with LaTeX-aware gap analysis
// ---------------------------------------------------------------------------

/// LaTeX gap syntax: math, commands, escapes and bare braces.
///
/// Arm order is load-bearing — `\[` and `\(` before `\name`, and `\name`
/// before the bare escape `\X`, which needs two bytes and so lets a lone
/// trailing backslash fall through as ordinary text.
fn latex_gap(b: &[u8], i: usize) -> Option<gap::Match> {
    use gap::Token::{Barrier, Elided, Separator};
    Some(match b[i..] {
        // Inline math: $...$
        [b'$', ..] => gap::Match::at(Separator, i, shared::close_at(b, i + 1, b"$", Some(b'\\'))),
        // Display math: \[...\]. It takes the whitespace on either side into
        // the exclusion as well, so blanking the formula does not leave a
        // double space in the middle of the sentence.
        [b'\\', b'[', ..] => {
            let end = shared::close_at(b, i + 2, b"\\]", Some(b'\\'));
            let padded = pad_whitespace(b, i, end);
            gap::Match::new(Separator, padded.0, padded.1)
        }
        // Inline math: \(...\)
        [b'\\', b'(', ..] => gap::Match::at(
            Separator,
            i,
            shared::close_at(b, i + 2, b"\\)", Some(b'\\')),
        ),
        [b'\\', first, ..] if first.is_ascii_alphabetic() => {
            let name_end = shared::run_end(b, i + 1, |c| c.is_ascii_alphabetic());
            // A block command means the words on either side are in different
            // paragraphs; its arguments are left for the scan to read as text,
            // exactly as before.
            if is_block_command(&b[i + 1..name_end]) {
                gap::Match::at(Barrier, i, name_end)
            } else {
                gap::Match::at(
                    Elided,
                    i,
                    command_args_end(b, &b[i + 1..name_end], name_end),
                )
            }
        }
        // Escape sequence: \\, \, , \; …
        [b'\\', _, ..] => gap::Match::at(Elided, i, i + 2),
        // A bare brace.
        [b'{' | b'}', ..] => gap::Match::at(Elided, i, i + 1),
        _ => return None,
    })
}

/// End of what follows a command name: a `\verb`-style delimited argument, or
/// the braced and bracketed arguments.
fn command_args_end(b: &[u8], name: &[u8], name_end: usize) -> usize {
    // The star of a starred variant belongs to neither the name nor the
    // arguments, so it goes before either is read.
    let mut i = name_end;
    if b.get(i) == Some(&b'*') {
        i += 1;
    }
    if is_verbatim_delimited(name)
        && let Some(end) = verbatim_argument_end(b, i)
    {
        return end;
    }
    shared::skip_command_args_bytes(b, i, &[(b'{', b'}'), (b'[', b']')])
}

/// Widen `start..end` over the whitespace on either side of it.
fn pad_whitespace(b: &[u8], start: usize, end: usize) -> (usize, usize) {
    let mut lo = start;
    while lo > 0 && b[lo - 1].is_ascii_whitespace() {
        lo -= 1;
    }
    let hi = shared::run_end(b, end, |c| c.is_ascii_whitespace());
    (lo, hi)
}

/// Whether a command is block-level layout rather than inline markup.
///
/// These must NOT be bridged: leaving the `\` in the stripped gap fails the
/// bridge check, which is what keeps the ranges on either side separate.
const fn is_block_command(name: &[u8]) -> bool {
    matches!(
        name,
        b"begin"
            | b"end"
            | b"item"
            | b"par"
            | b"section"
            | b"subsection"
            | b"subsubsection"
            | b"paragraph"
            | b"chapter"
            | b"part"
            | b"hfill"
            | b"vfill"
            | b"newline"
            | b"linebreak"
            | b"noindent"
    )
}

/// End of a `\[…\]` / `\(…\)` run: just past the closing `\]` or `\)`.
///
#[cfg(test)]
mod tests {
    use super::LatexExtras;
    use crate::prose::ProseExtractor;
    use anyhow::Result;

    #[test]
    fn test_latex_basic_extraction() -> Result<()> {
        let language: tree_sitter::Language = codebook_tree_sitter_latex::LANGUAGE.into();
        let mut extractor = ProseExtractor::new(language)?;

        let text = r"\documentclass{article}
\usepackage{amsmath}

\begin{document}

\section{Introduction}

This is a simple paragraph with some text.

\textbf{Bold text} and \textit{italic text} here.

\begin{verbatim}
This should be ignored completely.
\end{verbatim}

Another paragraph after verbatim.

\end{document}
";
        let ranges = extractor.extract(text, "latex", &LatexExtras::default())?;
        let extracted: Vec<&str> = ranges
            .iter()
            .map(|r| &text[r.start_byte..r.end_byte])
            .collect();

        assert!(
            extracted.iter().any(|t| t.contains("simple paragraph")),
            "Should extract prose text, got: {extracted:?}"
        );
        assert!(
            extracted.iter().any(|t| t.contains("Bold text")),
            "Should extract text inside \\textbf, got: {extracted:?}"
        );
        assert!(
            !extracted.iter().any(|t| t.contains("ignored completely")),
            "Should NOT extract verbatim content, got: {extracted:?}"
        );
        assert!(
            !extracted.iter().any(|t| t.contains("\\textbf")),
            "Should NOT contain latex commands, got: {extracted:?}"
        );
        assert!(
            !extracted.iter().any(|t| t.contains("\\documentclass")),
            "Should NOT contain preamble, got: {extracted:?}"
        );

        Ok(())
    }

    #[test]
    fn test_latex_math_excluded() -> Result<()> {
        let language: tree_sitter::Language = codebook_tree_sitter_latex::LANGUAGE.into();
        let mut extractor = ProseExtractor::new(language)?;

        let text = r"\documentclass{article}
\begin{document}

Some text before math.

$x^2 + y^2 = z^2$

Text after inline math.

\[
  \int_0^1 f(x) \, dx
\]

Text after display math.

\end{document}
";
        let ranges = extractor.extract(text, "latex", &LatexExtras::default())?;
        let extracted: Vec<&str> = ranges
            .iter()
            .map(|r| &text[r.start_byte..r.end_byte])
            .collect();

        assert!(
            extracted.iter().any(|t| t.contains("before math")),
            "Should extract text before math, got: {extracted:?}"
        );
        assert!(
            extracted.iter().any(|t| t.contains("after inline math")),
            "Should extract text after math, got: {extracted:?}"
        );
        assert!(
            !extracted.iter().any(|t| t.contains("x^2")),
            "Should NOT extract inline math, got: {extracted:?}"
        );
        assert!(
            !extracted.iter().any(|t| t.contains("\\int")),
            "Should NOT extract display math, got: {extracted:?}"
        );

        Ok(())
    }

    #[test]
    fn test_latex_verbatim_delimited_arguments_excluded() -> Result<()> {
        let language: tree_sitter::Language = codebook_tree_sitter_latex::LANGUAGE.into();
        let mut extractor = ProseExtractor::new(language)?;

        let text = "\
\\begin{document}
A paragraph with \\verb|verb_token| inline verbatim.

And \\lstinline!lst_token! plus \\verb*|star_token| and
\\mintinline{rust}|mint_token| too.
\\end{document}
";
        let ranges = extractor.extract(text, "latex", &LatexExtras::default())?;
        let prose: String = ranges.iter().map(|r| r.extract_text(text)).collect();

        assert!(prose.contains("A paragraph with"), "{prose:?}");
        assert!(prose.contains("inline verbatim"), "{prose:?}");
        for token in [
            "verb_token",
            "lst_token",
            "star_token",
            "mint_token",
            "rust",
        ] {
            assert!(!prose.contains(token), "{token} leaked into {prose:?}");
        }
        Ok(())
    }

    #[test]
    fn test_latex_preamble_excluded() -> Result<()> {
        let language: tree_sitter::Language = codebook_tree_sitter_latex::LANGUAGE.into();
        let mut extractor = ProseExtractor::new(language)?;

        let text = r"\documentclass{article}
\usepackage{amsmath}
\title{My Document}
\author{John Doe}

\begin{document}

Hello world.

\end{document}
";
        let ranges = extractor.extract(text, "latex", &LatexExtras::default())?;
        let extracted: Vec<&str> = ranges
            .iter()
            .map(|r| &text[r.start_byte..r.end_byte])
            .collect();

        assert!(
            extracted.iter().any(|t| t.contains("Hello world")),
            "Should extract body text, got: {extracted:?}"
        );
        assert!(
            !extracted.iter().any(|t| t.contains("My Document")),
            "Should NOT extract title from preamble, got: {extracted:?}"
        );
        assert!(
            !extracted.iter().any(|t| t.contains("John Doe")),
            "Should NOT extract author from preamble, got: {extracted:?}"
        );

        Ok(())
    }

    #[test]
    fn test_latex_no_document_env() -> Result<()> {
        let language: tree_sitter::Language = codebook_tree_sitter_latex::LANGUAGE.into();
        let mut extractor = ProseExtractor::new(language)?;

        let text = r"\section{Test}
Some text here.
";
        let ranges = extractor.extract(text, "latex", &LatexExtras::default())?;
        let extracted: Vec<&str> = ranges
            .iter()
            .map(|r| &text[r.start_byte..r.end_byte])
            .collect();

        assert!(
            extracted.iter().any(|t| t.contains("text here")),
            "Should extract text from snippet without document env, got: {extracted:?}"
        );

        Ok(())
    }

    #[test]
    fn test_latex_real_content() -> Result<()> {
        let language: tree_sitter::Language = codebook_tree_sitter_latex::LANGUAGE.into();
        let mut extractor = ProseExtractor::new(language)?;

        let text = r"\documentclass[10pt]{article}
\usepackage{styles/pagestyle}
\usepackage{styles/codestyle}

\begin{document}

{\scshape Notes } \hfill {\scshape \large } \hfill {\scshape \today}

\smallskip
\hrule
\bigskip

\section{Insertion sort}

There are two popular variants of insertsion sort you typically see

\begin{algorithm}[H]
  \caption{InsertionSort A}
  \begin{algorithmic}[1]
    \State $i \gets 1$
    \While{$i < \text{length}(A)$}
    \State $j \gets i$
    \EndWhile
  \end{algorithmic}
\end{algorithm}

\subsection{InsertionSort A}
The invariants for this version are relatively straightforward. The first invariant we specify is that the outer loop variable $i$ is always between $1$ and the length of the array (inclusive). So
\[
  1 \leq i \leq \text{length}(A) \tag{Index Constraint}
\]
Secondly, for the outer loop, we weaken the postcondition with the index variable $i$ to get the invariant that the subarray $A[0..i)$ is sorted.

\begin{grayblock}
  One sidenote we can actually weaken the `elements greater than key' invariant as follows
  \[
    \forall k.\ j < k \leq i \to A[k] \geq \text{key}
  \]
\end{grayblock}

\begin{minted}{dafny}
method InsertionSortA(a : array<int>)
  modifies a
  requires a.Length >= 1
\end{minted}

\end{document}
";
        let ranges = extractor.extract(text, "latex", &LatexExtras::default())?;
        let extracted: Vec<&str> = ranges
            .iter()
            .map(|r| &text[r.start_byte..r.end_byte])
            .collect();

        assert!(
            extracted.iter().any(|t| t.contains("popular variants")),
            "Should extract prose about variants, got: {extracted:?}"
        );
        assert!(
            extracted.iter().any(|t| t.contains("Insertion sort")),
            "Should extract section heading, got: {extracted:?}"
        );
        assert!(
            !extracted.iter().any(|t| t.contains("\\section")),
            "Should NOT contain \\section command, got: {extracted:?}"
        );
        assert!(
            !extracted.iter().any(|t| t.contains("pagestyle")),
            "Should NOT contain preamble, got: {extracted:?}"
        );
        assert!(
            !extracted.iter().any(|t| t.contains("InsertionSortA")),
            "Should NOT contain minted code, got: {extracted:?}"
        );
        assert!(
            !extracted.iter().any(|t| t.contains("\\caption")),
            "Should NOT contain algorithm content, got: {extracted:?}"
        );

        Ok(())
    }

    #[test]
    fn test_latex_algorithm_env() -> Result<()> {
        let language: tree_sitter::Language = codebook_tree_sitter_latex::LANGUAGE.into();
        let mut extractor = ProseExtractor::new(language)?;

        let text = r"\documentclass{article}
\begin{document}

Text before algorithm.

\begin{algorithm}[H]
  \caption{InsertionSort}
  \begin{algorithmic}[1]
    \State $i \gets 1$
  \end{algorithmic}
\end{algorithm}

Text after algorithm.

\end{document}
";
        let ranges = extractor.extract(text, "latex", &LatexExtras::default())?;
        let extracted: Vec<&str> = ranges
            .iter()
            .map(|r| &text[r.start_byte..r.end_byte])
            .collect();

        assert!(
            extracted.iter().any(|t| t.contains("before algorithm")),
            "Should extract text before algorithm, got: {extracted:?}"
        );
        assert!(
            extracted.iter().any(|t| t.contains("after algorithm")),
            "Should extract text after algorithm, got: {extracted:?}"
        );

        Ok(())
    }

    #[test]
    fn test_latex_inline_math_bridges() -> Result<()> {
        let language: tree_sitter::Language = codebook_tree_sitter_latex::LANGUAGE.into();
        let mut extractor = ProseExtractor::new(language)?;

        let text = r"\documentclass{article}
\begin{document}

The variable $i$ is always between $1$ and the length of the array.

Some text, with a comma and more text after it.

\end{document}
";
        let ranges = extractor.extract(text, "latex", &LatexExtras::default())?;
        let extracted: Vec<&str> = ranges
            .iter()
            .map(|r| &text[r.start_byte..r.end_byte])
            .collect();

        assert!(
            extracted.iter().any(|t| t.contains("variable")
                && t.contains("always between")
                && t.contains("length")),
            "Sentence with inline math should be a single chunk bridging across $i$ and $1$, got: {extracted:?}"
        );
        assert!(
            extracted
                .iter()
                .any(|t| t.contains("text,") || (t.contains("text") && t.contains("comma"))),
            "Sentence with comma should stay together, got: {extracted:?}"
        );

        Ok(())
    }

    #[test]
    fn test_latex_includegraphics_not_extracted() -> Result<()> {
        let language: tree_sitter::Language = codebook_tree_sitter_latex::LANGUAGE.into();
        let mut extractor = ProseExtractor::new(language)?;

        let text = r"\documentclass{article}
\begin{document}

Some text before.

\includegraphics[width=0.5\textwidth]{array.pdf}

Some text after.

\end{document}
";
        let ranges = extractor.extract(text, "latex", &LatexExtras::default())?;
        let extracted: Vec<&str> = ranges
            .iter()
            .map(|r| &text[r.start_byte..r.end_byte])
            .collect();

        assert!(
            !extracted.iter().any(|t| *t == "width" || *t == "0.5"),
            "Should NOT extract includegraphics optional args, got: {extracted:?}"
        );
        assert!(
            extracted.iter().any(|t| t.contains("text before")),
            "Should extract prose before includegraphics, got: {extracted:?}"
        );
        assert!(
            extracted.iter().any(|t| t.contains("text after")),
            "Should extract prose after includegraphics, got: {extracted:?}"
        );

        Ok(())
    }

    #[test]
    fn test_latex_display_math_excluded_from_text() -> Result<()> {
        let language: tree_sitter::Language = codebook_tree_sitter_latex::LANGUAGE.into();
        let mut extractor = ProseExtractor::new(language)?;

        let text = r"\documentclass{article}
\begin{document}

We know that
\[
  x^2 + y^2 = z^2
\]
which proves our claim.

\end{document}
";
        let ranges = extractor.extract(text, "latex", &LatexExtras::default())?;

        // The sentence should bridge across the display math
        let bridged = ranges.iter().any(|r| {
            let raw = &text[r.start_byte..r.end_byte];
            raw.contains("know that") && raw.contains("proves our claim")
        });
        assert!(
            bridged,
            "Sentence should bridge across display math, got: {:?}",
            ranges
                .iter()
                .map(|r| &text[r.start_byte..r.end_byte])
                .collect::<Vec<_>>()
        );

        // But extract_text should replace the math with spaces
        let bridged_range = ranges
            .iter()
            .find(|r| {
                let raw = &text[r.start_byte..r.end_byte];
                raw.contains("know that") && raw.contains("proves our claim")
            })
            .expect("Should have a bridged range");

        assert!(
            !bridged_range.exclusions.is_empty(),
            "Should have exclusions for display math"
        );

        let clean_text = bridged_range.extract_text(text);
        assert!(
            !clean_text.contains("x^2"),
            "extract_text should not contain math content, got: {:?}",
            clean_text
        );
        assert!(
            clean_text.contains("know that"),
            "extract_text should still contain prose, got: {:?}",
            clean_text
        );
        // Newlines around display math should also be blanked so that
        // "which" after \] doesn't look like a new sentence start
        assert!(
            !clean_text.contains('\n'),
            "extract_text should blank newlines around display math, got: {:?}",
            clean_text
        );

        Ok(())
    }

    #[test]
    fn test_latex_display_math_no_false_capitalization() -> Result<()> {
        let language: tree_sitter::Language = codebook_tree_sitter_latex::LANGUAGE.into();
        let mut extractor = ProseExtractor::new(language)?;

        // Reproduces the user's exact pattern: display math followed by
        // lowercase continuation that should NOT be flagged for capitalization.
        let text = r"\documentclass{article}
\begin{document}

Thus, our invariant for the inner loop is:
\[
  \forall p, q.\ 0 \leq p < q
\]
the intuition here being that all elements are in sorted order.

\end{document}
";
        let ranges = extractor.extract(text, "latex", &LatexExtras::default())?;

        // Should bridge across \[...\] into one chunk
        let bridged = ranges.iter().find(|r| {
            let raw = &text[r.start_byte..r.end_byte];
            raw.contains("invariant") && raw.contains("intuition")
        });
        assert!(
            bridged.is_some(),
            "Should bridge across display math, got: {:?}",
            ranges
                .iter()
                .map(|r| &text[r.start_byte..r.end_byte])
                .collect::<Vec<_>>()
        );

        let range = bridged.unwrap();
        let clean = range.extract_text(text);

        // The clean text should flow as "is:  the intuition" (with spaces
        // replacing the math), not "is:\n...\nthe" which would trigger
        // capitalization warnings.
        assert!(
            clean.contains("is:") && clean.contains("the intuition"),
            "Prose should flow continuously, got: {:?}",
            clean
        );
        assert!(
            !clean.contains("\\forall"),
            "Math commands should be blanked, got: {:?}",
            clean
        );

        Ok(())
    }

    #[test]
    fn test_latex_mathpar_skipped() -> Result<()> {
        let language: tree_sitter::Language = codebook_tree_sitter_latex::LANGUAGE.into();
        let mut extractor = ProseExtractor::new(language)?;

        let text = r"\documentclass{article}
\begin{document}

We define the rules as follows
\begin{mathpar}
    \inferrule
    { }
    {\Gamma \vdash n : \text{num}} \quad \text{T-Num}

    \inferrule
    {\Gamma (x) = \tau}
    {\Gamma \vdash x : \tau} \quad \text{T-Var}
\end{mathpar}

The proof is complete.

\end{document}
";
        let ranges = extractor.extract(text, "latex", &LatexExtras::default())?;
        let extracted: Vec<&str> = ranges
            .iter()
            .map(|r| &text[r.start_byte..r.end_byte])
            .collect();

        assert!(
            extracted.iter().any(|t| t.contains("define the rules")),
            "Should extract prose before mathpar, got: {extracted:?}"
        );
        assert!(
            extracted.iter().any(|t| t.contains("proof is complete")),
            "Should extract prose after mathpar, got: {extracted:?}"
        );
        assert!(
            !extracted.iter().any(|t| t.contains("T-Num")),
            "Should NOT extract inference rule labels, got: {extracted:?}"
        );
        assert!(
            !extracted.iter().any(|t| t.contains("T-Var")),
            "Should NOT extract inference rule labels, got: {extracted:?}"
        );
        // Single-letter fragments from \text{} inside math should not appear
        assert!(
            !extracted.iter().any(|t| *t == "x" || *t == "n"),
            "Should NOT extract single variable names from mathpar, got: {extracted:?}"
        );

        Ok(())
    }

    #[test]
    fn test_latex_thispagestyle_skipped() -> Result<()> {
        let language: tree_sitter::Language = codebook_tree_sitter_latex::LANGUAGE.into();
        let mut extractor = ProseExtractor::new(language)?;

        let text = r"\documentclass{article}
\begin{document}
\thispagestyle{empty}

Hello world.

\end{document}
";
        let ranges = extractor.extract(text, "latex", &LatexExtras::default())?;
        let extracted: Vec<&str> = ranges
            .iter()
            .map(|r| &text[r.start_byte..r.end_byte])
            .collect();

        assert!(
            !extracted.iter().any(|t| t.contains("empty")),
            "Should NOT extract thispagestyle argument, got: {extracted:?}"
        );
        assert!(
            extracted.iter().any(|t| t.contains("Hello world")),
            "Should extract body prose, got: {extracted:?}"
        );

        Ok(())
    }

    #[test]
    fn test_latex_hfill_breaks_ranges() -> Result<()> {
        let language: tree_sitter::Language = codebook_tree_sitter_latex::LANGUAGE.into();
        let mut extractor = ProseExtractor::new(language)?;

        let text = r"\documentclass{article}
\begin{document}

{\scshape LV } \hfill {\scshape \large Assignment 1} \hfill {\scshape \today}

Some real prose here.

\end{document}
";
        let ranges = extractor.extract(text, "latex", &LatexExtras::default())?;
        let extracted: Vec<&str> = ranges
            .iter()
            .map(|r| &text[r.start_byte..r.end_byte])
            .collect();

        // \hfill should break merging — LV and Assignment 1 should not be in the same range
        assert!(
            !extracted
                .iter()
                .any(|t| t.contains("LV") && t.contains("Assignment")),
            "\\hfill should break ranges, got: {extracted:?}"
        );
        assert!(
            extracted.iter().any(|t| t.contains("real prose")),
            "Should extract body prose, got: {extracted:?}"
        );

        Ok(())
    }

    #[test]
    fn test_latex_bnf_env_skipped() -> Result<()> {
        let language: tree_sitter::Language = codebook_tree_sitter_latex::LANGUAGE.into();
        let mut extractor = ProseExtractor::new(language)?;

        let text = r"\documentclass{article}
\begin{document}

The syntax is defined as follows.
\begin{bnf}(
        prod-delim={--},
        comment={//},
      )[
        colspec = {llcll},
      ]
      e // Expr ::=
      | n // number
    \end{bnf}

That concludes the grammar.

\end{document}
";
        let ranges = extractor.extract(text, "latex", &LatexExtras::default())?;
        let extracted: Vec<&str> = ranges
            .iter()
            .map(|r| &text[r.start_byte..r.end_byte])
            .collect();

        assert!(
            extracted.iter().any(|t| t.contains("syntax is defined")),
            "Should extract prose before bnf, got: {extracted:?}"
        );
        assert!(
            extracted
                .iter()
                .any(|t| t.contains("concludes the grammar")),
            "Should extract prose after bnf, got: {extracted:?}"
        );
        assert!(
            !extracted
                .iter()
                .any(|t| t.contains("prod-delim") || t.contains("colspec")),
            "Should NOT extract bnf parameters, got: {extracted:?}"
        );

        Ok(())
    }

    #[test]
    fn test_latex_inline_math_excluded_from_text() -> Result<()> {
        let language: tree_sitter::Language = codebook_tree_sitter_latex::LANGUAGE.into();
        let mut extractor = ProseExtractor::new(language)?;

        let text = r"\documentclass{article}
\begin{document}

The value $x + 1$ is positive and $y - 2$ is negative.

\end{document}
";
        let ranges = extractor.extract(text, "latex", &LatexExtras::default())?;

        let range = ranges
            .iter()
            .find(|r| {
                let raw = &text[r.start_byte..r.end_byte];
                raw.contains("value") && raw.contains("positive")
            })
            .expect("Should have a range containing the sentence");

        let clean = range.extract_text(text);
        assert!(
            !clean.contains("x + 1"),
            "extract_text should not contain inline math, got: {clean:?}"
        );
        assert!(
            !clean.contains("y - 2"),
            "extract_text should not contain second inline math, got: {clean:?}"
        );
        assert!(
            clean.contains("value"),
            "extract_text should preserve prose, got: {clean:?}"
        );
        assert!(
            clean.contains("positive"),
            "extract_text should preserve prose after math, got: {clean:?}"
        );

        Ok(())
    }

    #[test]
    fn test_latex_paren_math_excluded_from_text() -> Result<()> {
        let language: tree_sitter::Language = codebook_tree_sitter_latex::LANGUAGE.into();
        let mut extractor = ProseExtractor::new(language)?;

        let text = r"\documentclass{article}
\begin{document}

We define \(f(x) = x^2\) for all reals.

\end{document}
";
        let ranges = extractor.extract(text, "latex", &LatexExtras::default())?;

        let range = ranges
            .iter()
            .find(|r| {
                let raw = &text[r.start_byte..r.end_byte];
                raw.contains("define") && raw.contains("reals")
            })
            .expect("Should have a range containing the sentence");

        let clean = range.extract_text(text);
        assert!(
            !clean.contains("f(x)"),
            "extract_text should not contain \\(...\\) math, got: {clean:?}"
        );
        assert!(
            clean.contains("define"),
            "extract_text should preserve prose, got: {clean:?}"
        );

        Ok(())
    }

    #[test]
    fn test_latex_command_excluded_from_text() -> Result<()> {
        let language: tree_sitter::Language = codebook_tree_sitter_latex::LANGUAGE.into();
        let mut extractor = ProseExtractor::new(language)?;

        let text = r"\documentclass{article}
\begin{document}

The \textsc{Foo} method solves \textbf{bar} problems.

\end{document}
";
        let ranges = extractor.extract(text, "latex", &LatexExtras::default())?;

        let range = ranges
            .iter()
            .find(|r| {
                let raw = &text[r.start_byte..r.end_byte];
                raw.contains("method") && raw.contains("solves")
            })
            .expect("Should have a range containing the sentence");

        let clean = range.extract_text(text);
        assert!(
            !clean.contains("\\textsc"),
            "extract_text should not contain \\textsc command, got: {clean:?}"
        );
        assert!(
            !clean.contains("\\textbf"),
            "extract_text should not contain \\textbf command, got: {clean:?}"
        );
        // The word content inside the braces is also excluded (command + args),
        // so Foo and bar should be blanked too
        assert!(
            clean.contains("method"),
            "extract_text should preserve surrounding prose, got: {clean:?}"
        );
        assert!(
            clean.contains("solves"),
            "extract_text should preserve surrounding prose, got: {clean:?}"
        );

        Ok(())
    }

    #[test]
    fn test_latex_custom_skip_env() -> Result<()> {
        let language: tree_sitter::Language = codebook_tree_sitter_latex::LANGUAGE.into();
        let mut extractor = ProseExtractor::new(language)?;

        let text = r"\documentclass{article}
\begin{document}

Text before custom env.

\begin{prooftree}
  Some proof tree content here.
\end{prooftree}

Text after custom env.

\end{document}
";
        // Without extra skip envs, prooftree content is extracted
        let ranges = extractor.extract(text, "latex", &LatexExtras::default())?;
        let extracted: Vec<&str> = ranges
            .iter()
            .map(|r| &text[r.start_byte..r.end_byte])
            .collect();
        assert!(
            extracted.iter().any(|t| t.contains("proof tree content")),
            "Without config, prooftree content should be extracted, got: {extracted:?}"
        );

        // With extra skip envs, prooftree content is skipped
        let extra = vec!["prooftree".to_string()];
        let extras = LatexExtras {
            skip_envs: &extra,
            ..LatexExtras::default()
        };
        let ranges = extractor.extract(text, "latex", &extras)?;
        let extracted: Vec<&str> = ranges
            .iter()
            .map(|r| &text[r.start_byte..r.end_byte])
            .collect();
        assert!(
            !extracted.iter().any(|t| t.contains("proof tree content")),
            "With config, prooftree content should be skipped, got: {extracted:?}"
        );
        assert!(
            extracted.iter().any(|t| t.contains("Text before")),
            "Prose before should still be extracted, got: {extracted:?}"
        );
        assert!(
            extracted.iter().any(|t| t.contains("Text after")),
            "Prose after should still be extracted, got: {extracted:?}"
        );

        Ok(())
    }

    #[test]
    fn test_texttt_content_not_extracted() -> Result<()> {
        let language: tree_sitter::Language = codebook_tree_sitter_latex::LANGUAGE.into();
        let mut extractor = ProseExtractor::new(language)?;

        let text = r"\documentclass{article}
\begin{document}

Use the \texttt{myvar} variable in your code.

\end{document}
";
        let ranges = extractor.extract(text, "latex", &LatexExtras::default())?;

        // The raw range bridges across \texttt{myvar}, but extract_text
        // should strip the command and its argument via gap exclusions.
        let range = ranges
            .iter()
            .find(|r| {
                let raw = &text[r.start_byte..r.end_byte];
                raw.contains("variable")
            })
            .expect("Should have a range containing surrounding prose");

        let clean = range.extract_text(text);
        assert!(
            !clean.contains("myvar"),
            "extract_text should not contain \\texttt argument, got: {clean:?}"
        );
        assert!(
            clean.contains("variable"),
            "extract_text should preserve surrounding prose, got: {clean:?}"
        );

        Ok(())
    }

    #[test]
    fn test_custom_skip_command() -> Result<()> {
        let language: tree_sitter::Language = codebook_tree_sitter_latex::LANGUAGE.into();
        let mut extractor = ProseExtractor::new(language)?;

        let text = r"\documentclass{article}
\begin{document}

The \codefont{badspeling} function works.

\end{document}
";
        // Without skip_commands, codefont content IS extracted as prose words
        let ranges = extractor.extract(text, "latex", &LatexExtras::default())?;
        let clean_texts: Vec<_> = ranges.iter().map(|r| r.extract_text(text)).collect();
        assert!(
            clean_texts.iter().any(|t| t.contains("badspeling")),
            "Without config, codefont content should appear in prose, got: {clean_texts:?}"
        );

        // With skip_commands, the word node inside \codefont is not collected,
        // and extract_text strips the gap command — so "badspeling" disappears.
        let skip_cmds = vec!["codefont".to_string()];
        let extras = LatexExtras {
            skip_commands: &skip_cmds,
            ..LatexExtras::default()
        };
        let ranges = extractor.extract(text, "latex", &extras)?;
        let range = ranges
            .iter()
            .find(|r| {
                let raw = &text[r.start_byte..r.end_byte];
                raw.contains("function works")
            })
            .expect("Should have a range with surrounding prose");

        let clean = range.extract_text(text);
        assert!(
            !clean.contains("badspeling"),
            "With config, codefont argument should not appear in extract_text, got: {clean:?}"
        );
        assert!(
            clean.contains("function works"),
            "Surrounding prose should be preserved, got: {clean:?}"
        );

        Ok(())
    }
}