vamp/src/markdown.rs

use std::collections::HashMap;
use std::path::PathBuf;

use tree_sitter::{Node, Parser, Range, Tree};
use unicode_segmentation::UnicodeSegmentation;

use crate::plugin::TextPosition;

#[derive(Debug, Eq, PartialEq)]
pub(crate) struct LinkTarget {
    pub path: Option<PathBuf>,
    pub fragment: Option<String>,
}

pub(crate) fn link_at(text: &str, position: TextPosition) -> Result<LinkTarget, String> {
    let byte = position_byte(text, position)?;
    let blocks = block_tree(text)?;
    let mut node = blocks
        .root_node()
        .descendant_for_byte_range(byte, byte)
        .ok_or_else(no_link)?;
    loop {
        if node.kind() == "link_reference_definition" {
            return destination(child(node, "link_destination"), text);
        }
        if matches!(node.kind(), "inline" | "pipe_table_cell") {
            break;
        }
        node = node.parent().ok_or_else(no_link)?;
    }
    let (inline, ranges) = inline_tree(text, node)?;
    let mut node = inline
        .root_node()
        .descendant_for_byte_range(byte, byte)
        .ok_or_else(no_link)?;
    let mut link = None;
    loop {
        if !(node.start_byte() <= byte && byte < node.end_byte()) {
            return Err(no_link());
        }
        match node.kind() {
            "image" => return Err("Markdown images are not navigation links".into()),
            "uri_autolink" | "email_autolink" => {
                return Err("External Markdown URLs are not supported".into());
            }
            "inline_link"
            | "full_reference_link"
            | "collapsed_reference_link"
            | "shortcut_link" => {
                link = Some(node);
            }
            _ => {}
        }
        let Some(parent) = node.parent() else { break };
        node = parent;
    }
    let link = link.ok_or_else(no_link)?;
    if link.kind() == "inline_link" {
        return destination(child(link, "link_destination"), text);
    }
    let label = reference_label(link, text, &ranges).ok_or_else(no_link)?;
    let definitions = definitions(blocks.root_node(), text);
    let target = definitions
        .get(&label)
        .ok_or_else(|| format!("Markdown reference [{label}] is not defined"))?;
    destination(Some(*target), text)
}

pub(crate) fn heading_row(text: &str, fragment: &str) -> Result<usize, String> {
    if fragment.is_empty() {
        return Err("Markdown heading fragment is empty".into());
    }
    let blocks = block_tree(text)?;
    let definitions = definitions(blocks.root_node(), text);
    let mut slugs = HashMap::<String, usize>::new();
    let mut result = None;
    let mut error = None;
    walk(blocks.root_node(), |heading| {
        if result.is_some() || !matches!(heading.kind(), "atx_heading" | "setext_heading") {
            return;
        }
        let mut title = String::new();
        let mut heading_error = None;
        walk(heading, |node| {
            if node.kind() != "inline" {
                return;
            }
            match inline_tree(text, node) {
                Ok((tree, ranges)) => {
                    if let Err(err) =
                        visible_text(tree.root_node(), text, &ranges, &definitions, &mut title)
                    {
                        heading_error = Some(err);
                    }
                }
                Err(err) => heading_error = Some(err),
            }
        });
        if let Some(err) = heading_error {
            error = Some(err);
            return;
        }
        // The block grammar leaves optional closing ATX hashes in the inline node.
        if heading.kind() == "atx_heading"
            && let Some(content) = heading.child_by_field_name("heading_content")
        {
            let raw = source(content, text).trim_end();
            let without_hashes = raw.trim_end_matches('#');
            if without_hashes.len() < raw.len()
                && without_hashes
                    .chars()
                    .next_back()
                    .is_some_and(char::is_whitespace)
            {
                let hashes = raw.len() - without_hashes.len();
                title.truncate(title.trim_end().len().saturating_sub(hashes));
            }
        }
        let base = slug(title.trim());
        let mut suffix = *slugs.get(&base).unwrap_or(&0);
        let mut candidate = base.clone();
        while slugs.contains_key(&candidate) {
            suffix += 1;
            candidate = format!("{base}-{suffix}");
        }
        slugs.insert(base, suffix);
        slugs.entry(candidate.clone()).or_insert(0);
        if candidate == fragment {
            result = Some(heading.start_position().row);
        }
    });
    if let Some(row) = result {
        return Ok(row);
    }
    if let Some(error) = error {
        return Err(error);
    }
    Err(format!("Markdown heading #{fragment} was not found"))
}

fn no_link() -> String {
    "No Markdown link under the cursor".into()
}

fn position_byte(text: &str, position: TextPosition) -> Result<usize, String> {
    let mut start = 0;
    for (row, line) in text.split_inclusive('\n').enumerate() {
        if row == position.row {
            let line = line.trim_end_matches(['\r', '\n']);
            return line
                .grapheme_indices(true)
                .nth(position.column)
                .map(|(offset, _)| start + offset)
                .ok_or_else(no_link);
        }
        start += line.len();
    }
    Err(no_link())
}

fn block_tree(text: &str) -> Result<Tree, String> {
    let mut parser = Parser::new();
    parser
        .set_language(&tree_sitter_md::LANGUAGE.into())
        .map_err(|error| format!("Cannot load Markdown grammar: {error}"))?;
    parser
        .parse(text, None)
        .ok_or_else(|| "Cannot parse Markdown document".into())
}

fn inline_tree(text: &str, node: Node<'_>) -> Result<(Tree, Vec<Range>), String> {
    let mut ranges = Vec::new();
    let mut remaining = node.range();
    let mut cursor = node.walk();
    for continuation in node.named_children(&mut cursor) {
        if continuation.kind() != "block_continuation" {
            continue;
        }
        if remaining.start_byte < continuation.start_byte() {
            ranges.push(Range {
                end_byte: continuation.start_byte(),
                end_point: continuation.start_position(),
                ..remaining
            });
        }
        remaining.start_byte = continuation.end_byte();
        remaining.start_point = continuation.end_position();
    }
    if remaining.start_byte < remaining.end_byte {
        ranges.push(remaining);
    }
    let mut parser = Parser::new();
    parser
        .set_language(&tree_sitter_md::INLINE_LANGUAGE.into())
        .map_err(|error| format!("Cannot load Markdown inline grammar: {error}"))?;
    parser
        .set_included_ranges(&ranges)
        .map_err(|error| format!("Cannot parse Markdown inline ranges: {error}"))?;
    let tree = parser
        .parse(text, None)
        .ok_or_else(|| "Cannot parse Markdown inline content".to_string())?;
    Ok((tree, ranges))
}

fn walk<'a>(root: Node<'a>, mut visit: impl FnMut(Node<'a>)) {
    let mut cursor = root.walk();
    loop {
        visit(cursor.node());
        if cursor.goto_first_child() {
            continue;
        }
        while !cursor.goto_next_sibling() {
            if !cursor.goto_parent() {
                return;
            }
        }
    }
}

fn source<'a>(node: Node<'_>, text: &'a str) -> &'a str {
    &text[node.byte_range()]
}

fn child<'a>(node: Node<'a>, kind: &str) -> Option<Node<'a>> {
    let mut cursor = node.walk();
    node.named_children(&mut cursor)
        .find(|child| child.kind() == kind)
}

fn definitions<'a>(root: Node<'a>, text: &str) -> HashMap<String, Node<'a>> {
    let mut definitions = HashMap::new();
    walk(root, |node| {
        if node.kind() == "link_reference_definition"
            && let (Some(label), Some(destination)) =
                (child(node, "link_label"), child(node, "link_destination"))
        {
            let raw = source(label, text);
            let mut label_text = String::with_capacity(raw.len());
            let mut offset = label.start_byte() + 1;
            let mut cursor = label.walk();
            for continuation in label.named_children(&mut cursor) {
                if continuation.kind() == "block_continuation" {
                    label_text.push_str(&text[offset..continuation.start_byte()]);
                    offset = continuation.end_byte();
                }
            }
            label_text.push_str(&text[offset..label.end_byte() - 1]);
            definitions
                .entry(normalize_label(&label_text))
                .or_insert(destination);
        }
    });
    definitions
}

fn normalize_label(label: &str) -> String {
    // Normalize Unicode case, including ß/ẞ/SS and final sigma.
    let mut normalized = String::with_capacity(label.len());
    for word in label.split_whitespace() {
        if !normalized.is_empty() {
            normalized.push(' ');
        }
        normalized.extend(
            word.chars()
                .flat_map(char::to_lowercase)
                .flat_map(char::to_uppercase),
        );
    }
    normalized
}

fn reference_label(node: Node<'_>, text: &str, ranges: &[Range]) -> Option<String> {
    let label = child(
        node,
        if node.kind() == "full_reference_link" {
            "link_label"
        } else {
            "link_text"
        },
    )?;
    let mut raw = String::new();
    let range = if node.kind() == "full_reference_link" {
        label.start_byte() + 1..label.end_byte() - 1
    } else {
        label.byte_range()
    };
    append_source(text, ranges, range, &mut raw);
    Some(normalize_label(&raw))
}

fn append_source(
    text: &str,
    ranges: &[Range],
    source: std::ops::Range<usize>,
    output: &mut String,
) {
    for range in ranges {
        let start = range.start_byte.max(source.start);
        let end = range.end_byte.min(source.end);
        if start < end {
            output.push_str(&text[start..end]);
        }
    }
}

fn destination(node: Option<Node<'_>>, text: &str) -> Result<LinkTarget, String> {
    let raw = node
        .map(|node| decoded_text(node, text))
        .transpose()?
        .unwrap_or_default();
    let angled = node.is_some_and(|node| source(node, text).starts_with('<'));
    let raw = if angled {
        &raw[1..raw.len() - 1]
    } else {
        raw.as_str()
    };
    let (path, fragment) = raw
        .split_once('#')
        .map_or((raw, None), |(path, fragment)| (path, Some(fragment)));
    let path = percent_decode(path)?;
    let fragment = fragment.map(percent_decode).transpose()?;
    if path.starts_with("//") || path.starts_with("\\\\") || has_scheme(&path) {
        return Err("External Markdown URLs are not supported; use a local file path".into());
    }
    if path.contains('\0')
        || fragment
            .as_ref()
            .is_some_and(|fragment| fragment.contains('\0'))
    {
        return Err("Markdown link contains a NUL character".into());
    }
    if fragment.as_deref() == Some("") {
        return Err("Markdown heading fragment is empty".into());
    }
    Ok(LinkTarget {
        path: (!path.is_empty()).then(|| PathBuf::from(path)),
        fragment,
    })
}

fn has_scheme(path: &str) -> bool {
    let Some((scheme, _)) = path.split_once(':') else {
        return false;
    };
    scheme
        .as_bytes()
        .first()
        .is_some_and(u8::is_ascii_alphabetic)
        && scheme
            .bytes()
            .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'+' | b'-' | b'.'))
}

fn percent_decode(value: &str) -> Result<String, String> {
    if !value.contains('%') {
        return Ok(value.to_owned());
    }
    let mut bytes = Vec::with_capacity(value.len());
    let raw = value.as_bytes();
    let mut index = 0;
    while index < raw.len() {
        if raw[index] == b'%' {
            let hex = raw
                .get(index + 1..index + 3)
                .and_then(|hex| std::str::from_utf8(hex).ok())
                .and_then(|hex| u8::from_str_radix(hex, 16).ok())
                .ok_or_else(|| "Invalid percent escape in Markdown link".to_string())?;
            bytes.push(hex);
            index += 3;
        } else {
            bytes.push(raw[index]);
            index += 1;
        }
    }
    String::from_utf8(bytes).map_err(|_| "Markdown link percent escapes are not valid UTF-8".into())
}

fn append_unescaped(value: &str, result: &mut String) {
    let mut chars = value.chars().peekable();
    while let Some(character) = chars.next() {
        if character == '\\' && chars.peek().is_some_and(char::is_ascii_punctuation) {
            result.push(chars.next().unwrap());
        } else {
            result.push(character);
        }
    }
}

fn decoded_text(node: Node<'_>, text: &str) -> Result<String, String> {
    let mut output = String::with_capacity(node.end_byte() - node.start_byte());
    append_decoded(node, text, &mut output)?;
    Ok(output)
}

fn append_decoded(node: Node<'_>, text: &str, output: &mut String) -> Result<(), String> {
    if matches!(
        node.kind(),
        "entity_reference" | "numeric_character_reference"
    ) {
        return decode_entity(source(node, text), output);
    }
    let mut offset = node.start_byte();
    let mut cursor = node.walk();
    for child in node.named_children(&mut cursor) {
        append_unescaped(&text[offset..child.start_byte()], output);
        append_decoded(child, text, output)?;
        offset = child.end_byte();
    }
    append_unescaped(&text[offset..node.end_byte()], output);
    Ok(())
}

fn decode_entity(entity: &str, output: &mut String) -> Result<(), String> {
    let name = &entity[1..entity.len() - 1];
    let value = match name {
        "amp" | "AMP" => "&",
        "lt" | "LT" => "<",
        "gt" | "GT" => ">",
        "quot" | "QUOT" => "\"",
        "apos" => "'",
        "nbsp" | "NonBreakingSpace" => "\u{a0}",
        _ => {
            if let Some(number) = name.strip_prefix('#') {
                let (number, radix) = number
                    .strip_prefix(['x', 'X'])
                    .map_or((number, 10), |number| (number, 16));
                let value = u32::from_str_radix(number, radix)
                    .ok()
                    .filter(|value| *value != 0)
                    .and_then(char::from_u32)
                    .unwrap_or('\u{fffd}');
                output.push(value);
                return Ok(());
            }
            return Err(format!(
                "Unsupported Markdown HTML entity {entity}; use literal Unicode or a numeric entity"
            ));
        }
    };
    output.push_str(value);
    Ok(())
}

fn visible_text(
    node: Node<'_>,
    text: &str,
    ranges: &[Range],
    definitions: &HashMap<String, Node<'_>>,
    output: &mut String,
) -> Result<(), String> {
    match node.kind() {
        "html_tag" | "link_destination" | "link_title" | "block_continuation" => return Ok(()),
        kind if kind.ends_with("_delimiter") => return Ok(()),
        "entity_reference" | "numeric_character_reference" => {
            return decode_entity(source(node, text), output);
        }
        "backslash_escape" => {
            output.push_str(
                source(node, text)
                    .strip_prefix('\\')
                    .unwrap_or(source(node, text)),
            );
            return Ok(());
        }
        "code_span" => {
            let mut raw = String::new();
            append_source(text, ranges, node.byte_range(), &mut raw);
            let delimiter = raw.bytes().take_while(|byte| *byte == b'`').count();
            let mut code = String::with_capacity(raw.len() - delimiter * 2);
            let mut chars = raw[delimiter..raw.len() - delimiter].chars().peekable();
            while let Some(ch) = chars.next() {
                if ch == '\r' {
                    if chars.peek() == Some(&'\n') {
                        chars.next();
                    }
                    code.push(' ');
                } else {
                    code.push(if ch == '\n' { ' ' } else { ch });
                }
            }
            if code.starts_with(' ') && code.ends_with(' ') && code.chars().any(|ch| ch != ' ') {
                output.push_str(&code[1..code.len() - 1]);
            } else {
                output.push_str(&code);
            }
            return Ok(());
        }
        "inline_link" => {
            if let Some(label) = child(node, "link_text") {
                visible_text(label, text, ranges, definitions, output)?;
            }
            return Ok(());
        }
        "image" => {
            let mut cursor = node.walk();
            let inline_image = node.children(&mut cursor).any(|child| child.kind() == "(");
            let label = child(node, "image_description");
            if inline_image {
                if let Some(label) = label {
                    visible_text(label, text, ranges, definitions, output)?;
                }
                return Ok(());
            }
            let reference = if let Some(reference) = child(node, "link_label") {
                let raw = source(reference, text);
                Some(normalize_label(&raw[1..raw.len() - 1]))
            } else {
                label.map(|label| {
                    let mut raw = String::new();
                    append_source(text, ranges, label.byte_range(), &mut raw);
                    normalize_label(&raw)
                })
            };
            if reference.is_some_and(|label| definitions.contains_key(&label)) {
                if let Some(label) = label {
                    visible_text(label, text, ranges, definitions, output)?;
                }
                return Ok(());
            }
        }
        "full_reference_link" | "collapsed_reference_link" | "shortcut_link" => {
            if reference_label(node, text, ranges)
                .is_some_and(|label| definitions.contains_key(&label))
            {
                if let Some(label) = child(node, "link_text") {
                    visible_text(label, text, ranges, definitions, output)?;
                }
                return Ok(());
            }
        }
        _ => {}
    }
    let mut offset = node.start_byte();
    let mut cursor = node.walk();
    for child in node.named_children(&mut cursor) {
        append_source(text, ranges, offset..child.start_byte(), output);
        visible_text(child, text, ranges, definitions, output)?;
        offset = child.end_byte();
    }
    append_source(text, ranges, offset..node.end_byte(), output);
    Ok(())
}

fn slug(title: &str) -> String {
    let lower = title.to_lowercase();
    let mut slug = String::with_capacity(lower.len());
    for grapheme in lower.graphemes(true) {
        if grapheme.chars().next().is_some_and(char::is_whitespace) {
            slug.push('-');
        } else if grapheme
            .chars()
            .next()
            .is_some_and(|ch| ch.is_alphanumeric() || ch == '-' || ch == '_')
        {
            // Keep combining marks attached to Unicode letters, not just ASCII words.
            slug.push_str(grapheme);
        }
    }
    slug
}

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

    fn at(text: &str, needle: &str) -> Result<LinkTarget, String> {
        let byte = text.find(needle).unwrap();
        let prefix = &text[..byte];
        let row = prefix.bytes().filter(|byte| *byte == b'\n').count();
        let column = prefix.rsplit('\n').next().unwrap().graphemes(true).count();
        link_at(text, TextPosition { row, column })
    }

    fn target(path: Option<&str>, fragment: Option<&str>) -> LinkTarget {
        LinkTarget {
            path: path.map(PathBuf::from),
            fragment: fragment.map(str::to_owned),
        }
    }

    #[test]
    fn follows_inline_labels_destinations_and_grapheme_columns() {
        let text = "👩‍💻 [go **there**](<../My%20Notes/%E6%97%A5%E6%9C%AC\\(1\\).md#caf%C3%A9>)";
        let expected = target(Some("../My Notes/日本(1).md"), Some("café"));
        assert_eq!(at(text, "there").unwrap(), expected);
        assert_eq!(at(text, "My%20").unwrap(), expected);
        assert_eq!(
            at("[absolute](/Tmp/Note.md)", "absolute").unwrap(),
            target(Some("/Tmp/Note.md"), None)
        );
        assert_eq!(
            at("[here](#title)", "here").unwrap(),
            target(None, Some("title"))
        );
        assert_eq!(at("[top]()", "top").unwrap(), target(None, None));
        assert_eq!(
            at("[here](#title)", "[here").unwrap(),
            target(None, Some("title"))
        );
        assert_eq!(
            at("[file](literal%23name.md)", "file").unwrap(),
            target(Some("literal#name.md"), None)
        );
        assert_eq!(
            at("[file](My\\[Notes\\].md)", "file").unwrap(),
            target(Some("My[Notes].md"), None)
        );
    }

    #[test]
    fn resolves_all_reference_forms_and_multiline_block_definitions() {
        let text = "[full][ A  B ] [a b][] [A B]\n\n> [a b]:\n>   <My%20Notes.md#title>\n>   \"A title\"\n\n[a b]: ignored.md\n";
        let expected = target(Some("My Notes.md"), Some("title"));
        for needle in ["full", " A  B ", "a b][]", "A B]"] {
            assert_eq!(at(text, needle).unwrap(), expected);
        }
        assert_eq!(
            at("[Straße]\n\n[STRASSE]: file.md", "Straße").unwrap(),
            target(Some("file.md"), None)
        );
        let text = "> [multi\n> line][target]\n\n[target]: other.md\n";
        assert_eq!(at(text, "line]").unwrap(), target(Some("other.md"), None));
    }

    #[test]
    fn rejects_images_code_plaintext_and_external_urls() {
        for (text, needle) in [
            ("![image](file.md)", "image"),
            ("`[code](file.md)`", "code"),
            ("```md\n[code](file.md)\n```", "file"),
            ("    [code](file.md)", "code"),
            ("plain file.md text", "file"),
            ("[web](https://example.com)", "web"),
            ("[web](//example.com/a)", "web"),
            ("[web](mailto:a@example.com)", "web"),
            ("<https://example.com>", "example"),
        ] {
            assert!(at(text, needle).is_err(), "{text}");
        }
        assert!(
            at("[missing]", "missing")
                .unwrap_err()
                .contains("not defined")
        );
        assert!(
            at("[bad](file%ZZ.md)", "bad")
                .unwrap_err()
                .contains("percent")
        );
        assert!(
            at("[bad](file%FF.md)", "bad")
                .unwrap_err()
                .contains("UTF-8")
        );
        assert!(at("[bad](#)", "bad").unwrap_err().contains("empty"));
    }

    #[test]
    fn locates_visible_unicode_headings_setext_and_unique_slugs() {
        let text = "# **Café** [日本](destination.md) `code` ###\n\n> ## Café 日本 code\n\nCafé 日本 code\n====\n\n# Café 日本 code-1\n\n# Cafe\u{301} _name_ ![alt](hidden.md)\n\n```md\n# Hidden\n```\n";
        assert_eq!(heading_row(text, "café-日本-code").unwrap(), 0);
        assert_eq!(heading_row(text, "café-日本-code-1").unwrap(), 2);
        assert_eq!(heading_row(text, "café-日本-code-2").unwrap(), 4);
        assert_eq!(heading_row(text, "café-日本-code-1-1").unwrap(), 7);
        assert_eq!(heading_row(text, "cafe\u{301}-name-alt").unwrap(), 9);
        assert!(heading_row(text, "hidden").is_err());
        assert!(heading_row(text, "destinationmd").is_err());
        assert!(heading_row(text, "").is_err());
        assert_eq!(
            heading_row("# [Heading][ref]\n\n[ref]: ignored.md", "heading").unwrap(),
            0
        );
        assert_eq!(
            heading_row("# ![Alt][ref]\n\n[ref]: ignored.md", "alt").unwrap(),
            0
        );
        assert_eq!(heading_row("# ![Alt][missing]", "altmissing").unwrap(), 0);
    }

    #[test]
    fn renders_entities_or_reports_the_explicit_entity_ceiling() {
        assert_eq!(
            at("[file](a&amp;b&#x20;c.md)", "file").unwrap(),
            target(Some("a&b c.md"), None)
        );
        assert_eq!(
            at("[file](&lt;name&gt;)", "file").unwrap(),
            target(Some("<name>"), None)
        );
        assert!(
            at("[file](a&copy;b.md)", "file")
                .unwrap_err()
                .contains("Unsupported Markdown HTML entity")
        );
        assert_eq!(heading_row("# A &amp; B &#xE9;\n", "a--b-é").unwrap(), 0);
        assert!(
            heading_row("# &copy;\n", "copy")
                .unwrap_err()
                .contains("Unsupported Markdown HTML entity")
        );
        assert_eq!(heading_row("# &copy;\n# Valid\n", "valid").unwrap(), 1);
        assert_eq!(
            heading_row("> First\n> Second\n> ------\n", "first-second").unwrap(),
            0
        );
    }
}